diff --git a/Project.toml b/Project.toml index 56d130307..bfe7e7981 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Pathfinder" uuid = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454" -version = "0.9.31" +version = "0.10.0" authors = ["Seth Axen and contributors"] [workspace] @@ -9,13 +9,12 @@ projects = ["test"] [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" -Folds = "41a02a25-b8f0-4f67-bc48-60067656b558" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" +OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5" Optim = "429524aa-4258-5aef-a3af-852621145aeb" -Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" @@ -26,7 +25,6 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" -Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999" [weakdeps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" @@ -47,16 +45,15 @@ Accessors = "0.1.12" Distributions = "0.25.87" DynamicHMC = "3.4.0" DynamicPPL = "0.40.6" -Folds = "0.2.9" ForwardDiff = "0.10.26, 1" IrrationalConstants = "0.1.1, 0.2" LinearAlgebra = "1" LogDensityProblems = "2.1.0" MCMCChains = "6.0.2, 7" -Optim = "1.7.2, 2.0" -Optimization = "3.22, 4, 5" -OptimizationBase = "0.0.5, 1, 2, 3, 4, 5.0" -OptimizationOptimJL = "0.2.1, 0.3, 0.4" +OhMyThreads = "0.8.5" +Optim = "2" +OptimizationBase = "5" +OptimizationOptimJL = "0.4.10" PDMats = "0.11.35" PSIS = "0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9" ProgressLogging = "0.1.4" @@ -65,6 +62,5 @@ Requires = "1.1.3" SciMLBase = "2.30" Statistics = "1" StatsBase = "0.33.17, 0.34" -Transducers = "0.4.81" Turing = "0.43" julia = "1.10" diff --git a/docs/Project.toml b/docs/Project.toml index f174f0d0a..dd56d3e3e 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -39,7 +39,7 @@ ForwardDiff = "0.10.26, 1" LinearAlgebra = "1" LogDensityProblems = "2.1.0" LogDensityProblemsAD = "1.7" -Pathfinder = "0.9.27" +Pathfinder = "0.10" Printf = "1" ReverseDiff = "1.15" StableRNGs = "1" diff --git a/docs/inventories/Transducers.toml b/docs/inventories/Transducers.toml deleted file mode 100644 index 0411804b6..000000000 --- a/docs/inventories/Transducers.toml +++ /dev/null @@ -1,11 +0,0 @@ -# DocInventory version 1 -project = "Transducers.jl" -version = "0.4.84" - -# Filtered to just the types we link to -[[jl.type]] -name = "Transducers.PreferParallel" -uri = "reference/manual/#$" -[[jl.type]] -name = "Transducers.SequentialEx" -uri = "reference/manual/#$" diff --git a/docs/make.jl b/docs/make.jl index c00865816..9993e48de 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -41,12 +41,8 @@ links = InterLinks( ), "Optim" => "https://julianlsolvers.github.io/Optim.jl/stable/", "Optimization" => "https://docs.sciml.ai/Optimization/stable/", + "OhMyThreads" => "https://juliafolds2.github.io/OhMyThreads.jl/stable/", "PSIS" => "https://julia.arviz.org/PSIS/stable/", - "Transducers" => ( - "https://juliafolds2.github.io/Transducers.jl/stable/", # not built for a while - "https://juliafolds2.github.io/Transducers.jl/dev/objects.inv", - joinpath(@__DIR__, "inventories", "Transducers.toml"), - ), ) makedocs(; diff --git a/docs/src/index.md b/docs/src/index.md index 2f51db662..855cb89de 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -41,7 +41,7 @@ Pathfinder uses several packages for extended functionality: - [Optimization.jl](https://optimization.sciml.ai/stable/): This allows the L-BFGS optimizer to be replaced with any of the many Optimization-compatible optimizers and supports use of callbacks. Note that any changes made to Pathfinder using these features would be experimental. - [ADTypes.jl](https://sciml.github.io/ADTypes.jl/stable/): Supports specifying the automatic differentiation engine to be used for computing gradient and Hessian, if needed. -- [Transducers.jl](https://juliafolds.github.io/Transducers.jl/stable/): parallelization support +- [OhMyThreads.jl](https://juliafolds2.github.io/OhMyThreads.jl/stable/): parallelization support - [Distributions.jl](https://juliastats.org/Distributions.jl/stable/)/[PDMats.jl](https://github.com/JuliaStats/PDMats.jl): fits can be used anywhere a `Distribution` can be used - [LogDensityProblems.jl](https://www.tamaspapp.eu/LogDensityProblems.jl/stable/): defining the log-density function, gradient, and Hessian - [ProgressLogging.jl](https://julialogging.github.io/ProgressLogging.jl/stable/): In Pluto, Juno, and VSCode, nested progress bars are shown. In the REPL, use TerminalLoggers.jl to get progress bars. diff --git a/ext/PathfinderTuringExt.jl b/ext/PathfinderTuringExt.jl index cf4db8931..3ed77fa3e 100644 --- a/ext/PathfinderTuringExt.jl +++ b/ext/PathfinderTuringExt.jl @@ -214,7 +214,7 @@ Run multi-path Pathfinder on a Turing `model`. ```jldoctest julia> using Pathfinder, Turing, StableRNGs -julia> rng = StableRNG(42); +julia> rng = StableRNG(3); julia> @model function demo_model() α ~ Normal(0, 1) diff --git a/src/Pathfinder.jl b/src/Pathfinder.jl index e1ab7f8a3..1f66811d2 100644 --- a/src/Pathfinder.jl +++ b/src/Pathfinder.jl @@ -2,12 +2,12 @@ module Pathfinder using ADTypes: ADTypes using Distributions: Distributions -using Folds: Folds using IrrationalConstants: log2π using LinearAlgebra using LogDensityProblems: LogDensityProblems +using OhMyThreads: OhMyThreads using Optim: Optim, LineSearches -using Optimization: Optimization +using OptimizationBase: OptimizationBase, OptimizationState using OptimizationOptimJL: OptimizationOptimJL using PDMats: PDMats using ProgressLogging: ProgressLogging @@ -17,13 +17,6 @@ using Requires: Requires using SciMLBase: SciMLBase using Statistics: Statistics using StatsBase: StatsBase -using Transducers: Transducers - -if isdefined(Optimization, :OptimizationState) - using Optimization: OptimizationState -else - using OptimizationBase: OptimizationState -end # Declare and export the public API export PathfinderResult, MultiPathfinderResult @@ -45,7 +38,7 @@ end # We depend on Optim, and Optim depends on ForwardDiff, so we can offer it as a default. default_ad() = ADTypes.AutoForwardDiff() -include("transducers.jl") +include("utils.jl") include("woodbury.jl") include("optimize.jl") include("inverse_hessian.jl") diff --git a/src/elbo.jl b/src/elbo.jl index e7190e57b..e5f08726a 100644 --- a/src/elbo.jl +++ b/src/elbo.jl @@ -1,13 +1,11 @@ -function maximize_elbo(rng, logp, dists, ndraws, executor) - EE = Core.Compiler.return_type( - elbo_and_samples, Tuple{typeof(rng),typeof(logp),eltype(dists),Int} - ) - estimates = similar(dists, EE) - isempty(estimates) && return 0, estimates - Folds.map!(estimates, dists, executor) do dist - return elbo_and_samples(rng, logp, dist, ndraws) +function maximize_elbo(rng, logp, dists, ndraws, ntasks::Int) + seeds = rand!(rng, similar(dists, UInt64)) + estimates = _chunk_tmap(dists, seeds; ntasks, setup=() -> copy(rng)) do chunk_rng, dist, seed + Random.seed!(chunk_rng, seed) + return elbo_and_samples(chunk_rng, logp, dist, ndraws) end - _, iteration_opt = _findmax(estimates |> Transducers.Map(est -> est.value)) + isempty(estimates) && return 0, estimates + _, iteration_opt = _findmax_skipnan(est -> est.value, estimates) return iteration_opt, estimates end @@ -17,7 +15,7 @@ function elbo_and_samples(rng, logp, dist, ndraws) logpϕ .= logp.(eachcol(ϕ)) logr = logpϕ - logqϕ elbo = Statistics.mean(logr) - elbo_se = sqrt(Statistics.var(logr) / length(logr)) + elbo_se = sqrt(Statistics.var(logr; mean = elbo) / length(logr)) return ELBOEstimate(elbo, elbo_se, ϕ, logpϕ, logqϕ, logr) end diff --git a/src/multipath.jl b/src/multipath.jl index ed5dd1f87..f7d9a12e0 100644 --- a/src/multipath.jl +++ b/src/multipath.jl @@ -101,15 +101,13 @@ $(_ARGUMENT_DOCSTRING) - `importance::Bool=true`: Perform Pareto smoothed importance resampling of draws. - `rng::AbstractRNG=Random.default_rng()`: Pseudorandom number generator. It is recommended to use a parallelization-friendly PRNG like the default PRNG on Julia 1.7 and up. -- `executor::Transducers.Executor`: Transducers.jl executor that determines if and how to - run the single-path runs in parallel, defaulting to - [`Transducers.SequentialEx()`](@extref `Transducers.SequentialEx`). If a transducer for - multi-threaded computation is selected, you must first verify that `rng` and the log - density function are thread-safe. -- `executor_per_run::Transducers.Executor`: Transducers.jl executor used within each run to - parallelize PRNG calls, defaulting to - [`Transducers.SequentialEx()`](@extref `Transducers.SequentialEx`). See - [`pathfinder`](@ref) for further description. +- `ntasks::Int=1`: maximum number of parallel tasks used to run the single-path + Pathfinder runs and to evaluate the target log density across draws. The default + `ntasks = 1` runs sequentially; larger values parallelize across runs and across + importance-resampling log-density evaluations, in which case the log-density function + must be thread-safe. Results are reproducible regardless of `ntasks`. +- `ntasks_per_run::Int=1`: same as `ntasks`, but applied within each single-path run for + parallelizing the ELBO evaluation. See [`pathfinder`](@ref) for details. - `kwargs...` : Remaining keywords are forwarded to [`pathfinder`](@ref). # Returns @@ -141,64 +139,92 @@ function multipathfinder( rng::Random.AbstractRNG=Random.default_rng(), history_length::Int=DEFAULT_HISTORY_LENGTH, optimizer=default_optimizer(history_length), - executor::Transducers.Executor=Transducers.SequentialEx(), - executor_per_run=Transducers.SequentialEx(), + ntasks::Int=1, + ntasks_per_run::Int=1, importance::Bool=true, kwargs..., ) - if init === nothing + _init = if init === nothing nruns > 0 || throw( ArgumentError("A positive `nruns` must be set or `init` must be provided.") ) - _init = fill(init, nruns) + fill(init, nruns) else - _init = init + init end + nruns = length(_init) if ndraws > ndraws_per_run * nruns @warn "More draws requested than total number of draws across replicas. Draws will not be unique." end logp(x) = -optim_fun.f(x, nothing) # run pathfinder independently from each starting point - trans = Transducers.Map() do (init_i, optimizer_i) - return pathfinder( - optim_fun; - rng, - history_length, - optimizer=optimizer_i, - ndraws=ndraws_per_run, - init=init_i, - executor=executor_per_run, - ndraws_elbo, - kwargs..., - ) + run_seeds = rand!(rng, similar(_init, UInt64)) + + pathfinder_results = ProgressLogging.@withprogress name = "Multi-path Pathfinder" begin + progress_taskref = Ref{Task}() + progress_channel = Channel{Bool}( + min(nruns, 1_000); spawn=true, taskref=progress_taskref + ) do ch + # Throttle progress logs as generally the logging system is not super performant: + # - At most 1 progress log every 0.1 seconds + # - At most 1 progress log every 0.5% progress + progress_step = max(1, cld(nruns, 200)) + next_logged = progress_step + next_time = time() + 0.1 + + completed = 0 + while take!(ch) + completed += 1 + now = time() + if completed >= next_logged && now >= next_time + ProgressLogging.@logprogress completed / nruns + next_logged = completed + progress_step + next_time = now + 0.1 + end + end + end + + try + # `Optim` optimizers may carry mutable state, so each chunk gets its own copy. + _chunk_tmap( + _init, run_seeds; ntasks, setup=() -> (copy(rng), deepcopy(optimizer)) + ) do (chunk_rng, chunk_optimizer), init, seed + Random.seed!(chunk_rng, seed) + result = pathfinder( + optim_fun; + rng=chunk_rng, + history_length, + optimizer=chunk_optimizer, + ndraws=ndraws_per_run, + init, + ntasks=ntasks_per_run, + ndraws_elbo, + kwargs..., + ) + put!(progress_channel, true) + yield() + return result + end + finally + put!(progress_channel, false) + close(progress_channel) + wait(progress_taskref[]) + end end - iter_optimizers = fill(optimizer, nruns) - iter_sp = - if executor isa Folds.ThreadedEx - # temporary workaround due to - # https://github.com/JuliaFolds2/Transducers.jl/issues/10 - # also support optimizers that store state - zip(_init, Iterators.map(deepcopy, iter_optimizers)) - else - Transducers.withprogress(zip(_init, iter_optimizers); interval=1e-3) - end |> trans - pathfinder_results = Folds.collect(iter_sp, executor) - fit_distributions = - pathfinder_results |> Transducers.Map(x -> x.fit_distribution) |> collect - draws_all = reduce(hcat, pathfinder_results |> Transducers.Map(x -> x.draws)) + fit_distributions = map(x -> x.fit_distribution, pathfinder_results) + draws_all = mapreduce(x -> x.draws, hcat, pathfinder_results) # draw samples from augmented mixture model inds = axes(draws_all, 2) sample_inds, psis_result = if importance - log_densities_fit = - pathfinder_results |> - Transducers.MapCat() do x - return Distributions.logpdf(x.fit_distribution, x.draws) - end |> - collect - iter_logp = eachcol(draws_all) |> Transducers.Map(logp) - log_densities_target = Folds.collect(iter_logp, executor) + log_densities_fit = _maybe_tmapreduce( + x -> Distributions.logpdf(x.fit_distribution, x.draws), + vcat, + pathfinder_results, + ntasks, + ) + log_densities_target = _maybe_tmap(logp, eachcol(draws_all), ntasks) log_densities_ratios = log_densities_target - log_densities_fit resample(rng, inds, log_densities_ratios, ndraws) else diff --git a/src/mvnormal.jl b/src/mvnormal.jl index c89f978c9..8914b2b83 100644 --- a/src/mvnormal.jl +++ b/src/mvnormal.jl @@ -13,12 +13,10 @@ were rejected due to keeping the inverse Hessian positive definite. """ function fit_mvnormals(θs, ∇logpθs; kwargs...) Σs, num_bfgs_updates_rejected = lbfgs_inverse_hessians(θs, ∇logpθs; kwargs...) - trans = Transducers.MapSplat() do Σ, ∇logpθ, θ + dists = map(Σs, ∇logpθs, θs) do Σ, ∇logpθ, θ μ = muladd(Σ, ∇logpθ, θ) return Distributions.MvNormal(μ, Σ) end - l = length(Σs) - dists = @views(zip(Σs, ∇logpθs[1:l], θs[1:l])) |> trans |> collect return dists, num_bfgs_updates_rejected end diff --git a/src/optimize.jl b/src/optimize.jl index f1d658822..9d7515c8f 100644 --- a/src/optimize.jl +++ b/src/optimize.jl @@ -51,7 +51,7 @@ function optimize_with_trace( _callback = OptimizationCallback( xs, fxs, ∇fxs, progress_name, progress_id, maxiters, callback, fail_on_nonfinite ) - sol = Optimization.solve(prob, optimizer; callback=_callback, maxiters, kwargs...) + sol = SciMLBase.solve(prob, optimizer; callback=_callback, maxiters, kwargs...) _∇fxs = _fill_missing_gradient_values!(∇fxs, xs, sol.cache.f) diff --git a/src/singlepath.jl b/src/singlepath.jl index a2fc8bfe3..73ee98f6b 100644 --- a/src/singlepath.jl +++ b/src/singlepath.jl @@ -111,14 +111,10 @@ $(_ARGUMENT_DOCSTRING) - `ndraws_elbo::Int=$DEFAULT_NDRAWS_ELBO`: Number of draws used to estimate the ELBO - `ndraws::Int=ndraws_elbo`: number of approximate draws to return - `rng::Random.AbstractRNG`: The random number generator to be used for drawing samples -- `executor::Transducers.Executor`: Transducers.jl executor that - determines if and how to perform ELBO computation in parallel. The default - ([`Transducers.SequentialEx()`](@extref `Transducers.SequentialEx`)) performs no - parallelization. If `rng` is known to be thread-safe, and the log-density function is - known to have no internal state, then - [`Transducers.PreferParallel()`](@extref `Transducers.PreferParallel`) may be used to - parallelize log-density evaluation. This is generally only faster for expensive log - density functions. +- `ntasks::Int=1`: maximum number of parallel tasks used to compute ELBO estimates across + the fitted distributions. The default `ntasks = 1` runs sequentially; larger values + parallelize the ELBO evaluation, in which case the log-density function must be + thread-safe. Results are reproducible regardless of `ntasks`. - `history_length::Int=$DEFAULT_HISTORY_LENGTH`: Size of the history used to approximate the inverse Hessian. - `optimizer`: Optimizer to be used for constructing trajectory. Can be any optimizer @@ -293,7 +289,7 @@ function _pathfinder( history_length::Int=DEFAULT_HISTORY_LENGTH, optimizer=default_optimizer(history_length), ndraws_elbo=DEFAULT_NDRAWS_ELBO, - executor::Transducers.Executor=Transducers.SequentialEx(), + ntasks::Int=1, kwargs..., ) # compute trajectory @@ -308,7 +304,7 @@ function _pathfinder( # find ELBO-maximizing distribution fit_iteration, elbo_estimates = @views maximize_elbo( - rng, logp, fit_distributions[(begin + 1):end], ndraws_elbo, executor + rng, logp, fit_distributions[(begin + 1):end], ndraws_elbo, ntasks ) if isempty(elbo_estimates) success = false diff --git a/src/transducers.jl b/src/transducers.jl deleted file mode 100644 index 953e98845..000000000 --- a/src/transducers.jl +++ /dev/null @@ -1,18 +0,0 @@ -# transducer-friendly findmax, ignoring NaNs -function _findmax(x) - return Transducers.foldxl(x |> Transducers.Enumerate(); init=missing) do xmax_imax, i_xi - xmax_imax === missing && return reverse(i_xi) - i, xi = i_xi - isnan(xi) && return xmax_imax - xmax = first(xmax_imax) - isnan(xmax) && return (xi, i) - return xi > xmax ? (xi, i) : xmax_imax - end -end - -if !hasmethod(size, Tuple{Transducers.ProgressLoggingFoldable}) - # WARNING: Type piracy! - # https://github.com/JuliaFolds/Transducers.jl/issues/521 - # this method is necessary for Transducers version earlier than v0.4.82 - Base.size(x::Transducers.ProgressLoggingFoldable) = size(x.foldable) -end diff --git a/src/utils.jl b/src/utils.jl new file mode 100644 index 000000000..b6aa0ecde --- /dev/null +++ b/src/utils.jl @@ -0,0 +1,72 @@ +# Internal helpers + +# Number of parallel chunks for `len` items given `ntasks`; `1` runs sequentially. +function _nchunks(len::Int, ntasks::Int) + nchunks = min(len, ntasks) + if nchunks > 1 && Threads.nthreads() > 1 + return nchunks + else + return 1 + end +end + +# Map `f(s, xs[i]...)` over the elements of the equal-axes arrays `xss`, using at most +# `ntasks` parallel tasks and per-chunk state `s = setup()` built once per chunk. +function _chunk_tmap(f, xss::AbstractArray...; ntasks::Int, setup) + nchunks = _nchunks(length(first(xss)), ntasks) + if nchunks == 1 + let s = setup() + return map((xs...) -> f(s, xs...), xss...) + end + else + return OhMyThreads.tmapreduce( + vcat, OhMyThreads.chunks(eachindex(xss...); n=nchunks); chunking=false + ) do chunk + let s = setup() + return map(i -> f(s, map(xs -> xs[i], xss)...), chunk) + end + end + end +end + +# `map`/`mapreduce` using at most `ntasks` parallel tasks. +function _maybe_tmap(f, xs::AbstractArray, ntasks::Int) + nchunks = _nchunks(length(xs), ntasks) + if nchunks == 1 + return map(f, xs) + else + return OhMyThreads.tmap(f, xs; nchunks) + end +end + +function _maybe_tmapreduce(f, op, xs::AbstractArray, ntasks::Int) + nchunks = _nchunks(length(xs), ntasks) + if nchunks == 1 + return mapreduce(f, op, xs) + else + return OhMyThreads.tmapreduce(f, op, xs; nchunks) + end +end + +# Return `(maxvalue, argmaxindex)` of `key.(xs)`, skipping `NaN`s. If the +# first observed value is `NaN`, it is retained until a non-`NaN` entry +# replaces it; if every entry is `NaN`, the first one is returned (matching +# the semantics the success check in `singlepath.jl` relies on). +_findmax_skipnan(xs) = _findmax_skipnan(identity, xs) + +function _findmax_skipnan(key, xs) + state = missing + for (i, x) in enumerate(xs) + xi = key(x) + if state === missing + state = (xi, i) + continue + end + isnan(xi) && continue + xmax = first(state) + if isnan(xmax) || xi > xmax + state = (xi, i) + end + end + return state +end diff --git a/test/Project.toml b/test/Project.toml index 50697402b..b9782685e 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -5,7 +5,6 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Optim = "429524aa-4258-5aef-a3af-852621145aeb" -Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" OptimizationNLopt = "4e6fcdb7-1186-4e1f-a706-475e75c168bb" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" @@ -16,7 +15,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999" [compat] ADTypes = "0.2.5, 1" @@ -25,18 +23,16 @@ ForwardDiff = "0.10.26, 1" LinearAlgebra = "1" LogDensityProblems = "2.1.0" Optim = "1.7.2, 2.0" -Optimization = "3.22, 4, 5" -OptimizationBase = "0.0.5, 1, 2, 3, 4, 5.0" -OptimizationNLopt = "0.2, 0.3" +OptimizationBase = "5" +OptimizationNLopt = "0.3" PDMats = "0.11.35" PSIS = "0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9" -Pathfinder = "0.9.27" +Pathfinder = "0.10" ProgressLogging = "0.1.4" Random = "1" ReverseDiff = "1.15" SciMLBase = "2.30" Test = "1" -Transducers = "0.4.81" [sources] Pathfinder = {path = ".."} diff --git a/test/elbo.jl b/test/elbo.jl index 4faf2f7f3..d48bbf8cf 100644 --- a/test/elbo.jl +++ b/test/elbo.jl @@ -2,7 +2,6 @@ using Distributions using Pathfinder using Random using Test -using Transducers @testset "ELBO estimation" begin @testset "elbo_and_samples" begin @@ -34,21 +33,20 @@ using Transducers logp(x) = logpdf(target_dist, x[1]) σs = [1e-3, 0.05, σ_target, 1.0, 1.1, 1.2, 5.0, 10.0] dists = Normal.(0, σs) - executors = [SequentialEx(), ThreadedEx()] - @testset "$executor" for executor in executors + @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) rng = Random.seed!(Random.default_rng(), 42) lopt, estimates = @inferred Pathfinder.maximize_elbo( - rng, logp, dists, 100, executor + rng, logp, dists, 100, ntasks ) @test lopt == 3 @test estimates[lopt].value ≈ 0 rng = Random.seed!(Random.default_rng(), 42) - lopt2, estimates2 = Pathfinder.maximize_elbo(rng, logp, dists, 100, executor) + lopt2, estimates2 = Pathfinder.maximize_elbo(rng, logp, dists, 100, ntasks) @test lopt2 == lopt @test getproperty.(estimates2, :value) == getproperty.(estimates, :value) @test getproperty.(estimates2, :std_err) == getproperty.(estimates, :std_err) lopt3, estimates3 = @inferred Pathfinder.maximize_elbo( - rng, logp, dists[2:1], 100, executor + rng, logp, dists[2:1], 100, ntasks ) @test lopt3 == 0 @test isempty(estimates3) diff --git a/test/integration/AdvancedHMC/Project.toml b/test/integration/AdvancedHMC/Project.toml index b0b1dca74..0322c5b7d 100644 --- a/test/integration/AdvancedHMC/Project.toml +++ b/test/integration/AdvancedHMC/Project.toml @@ -26,8 +26,8 @@ LinearAlgebra = "1" LogDensityProblems = "2.1.0" LogDensityProblemsAD = "1.7" MCMCDiagnosticTools = "0.3" -Optim = "1.7.2, 2.0" -Pathfinder = "0.9.21" +Optim = "2" +Pathfinder = "0.10" Random = "1" Statistics = "1" StatsFuns = "1" diff --git a/test/integration/DynamicHMC/Project.toml b/test/integration/DynamicHMC/Project.toml index b29bca51f..cae2bea78 100644 --- a/test/integration/DynamicHMC/Project.toml +++ b/test/integration/DynamicHMC/Project.toml @@ -22,8 +22,8 @@ LinearAlgebra = "1" LogDensityProblems = "2.1.0" LogDensityProblemsAD = "1.7" MCMCDiagnosticTools = "0.3" -Optim = "1.7.2, 2.0" -Pathfinder = "0.9.18" +Optim = "2" +Pathfinder = "0.10" Random = "1" Statistics = "1" StatsFuns = "1" diff --git a/test/integration/Turing/Project.toml b/test/integration/Turing/Project.toml index 6d6713a85..ad32d939a 100644 --- a/test/integration/Turing/Project.toml +++ b/test/integration/Turing/Project.toml @@ -21,7 +21,7 @@ DynamicPPL = "0.40.6" FlexiChains = "0.3.4" LinearAlgebra = "1" LogDensityProblems = "2.1.0" -Pathfinder = "0.9.27" +Pathfinder = "0.10" Random = "1" ReverseDiff = "1.15" Test = "1" diff --git a/test/multipath.jl b/test/multipath.jl index 7088be772..b4124045f 100644 --- a/test/multipath.jl +++ b/test/multipath.jl @@ -2,13 +2,11 @@ using ADTypes using Distributions using ForwardDiff using LinearAlgebra -using Optimization using Pathfinder using PSIS using ReverseDiff using SciMLBase using Test -using Transducers @testset "multi path pathfinder" begin @testset "MvNormal" begin @@ -24,12 +22,10 @@ using Transducers ℓ = build_logdensityproblem(logp, dim, 2) rngs = [MersenneTwister(), Random.default_rng()] seed = 76 - @testset for rng in rngs - executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx() - + @testset for rng in rngs, ntasks in unique((1, Threads.nthreads())) Random.seed!(rng, seed) result = multipathfinder( - ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, executor + ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, ntasks ) @test result isa MultiPathfinderResult @test result.input === ℓ @@ -66,7 +62,7 @@ using Transducers Random.seed!(rng, seed) result2 = multipathfinder( - ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, executor + ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, ntasks ) @test result2.fit_distribution == result.fit_distribution @test result2.draws == result.draws @@ -74,7 +70,7 @@ using Transducers Random.seed!(rng, seed) result3 = multipathfinder( - ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, executor + ℓ, ndraws; nruns, ndraws_elbo, ndraws_per_run, rng, ntasks ) for (c1, c2) in zip(result.fit_distribution.components, result3.fit_distribution.components) @@ -107,4 +103,43 @@ using Transducers @test_throws ArgumentError multipathfinder(ℓ, 10; nruns=0) multipathfinder(ℓ, 10; nruns=2) end + + @testset "reproducibility across ntasks" begin + logp(x) = -sum(abs2, x) / 2 + ℓ = build_logdensityproblem(logp, 5, 2) + seed = 19 + nruns = 6 + ndraws = 200 + nthreads = Threads.nthreads() + + @testset "explicit rng" begin + rng = MersenneTwister(seed) + serial = multipathfinder(ℓ, ndraws; rng, nruns, ntasks=1, ntasks_per_run=1) + Random.seed!(rng, seed) + threaded = multipathfinder( + ℓ, ndraws; rng, nruns, ntasks=nthreads, ntasks_per_run=nthreads + ) + @test serial.draws == threaded.draws + @test serial.draw_component_ids == threaded.draw_component_ids + @test [c.μ for c in serial.fit_distribution.components] == + [c.μ for c in threaded.fit_distribution.components] + @test [c.Σ for c in serial.fit_distribution.components] == + [c.Σ for c in threaded.fit_distribution.components] + end + + @testset "default rng" begin + Random.seed!(seed) + serial = multipathfinder(ℓ, ndraws; nruns, ntasks=1, ntasks_per_run=1) + Random.seed!(seed) + threaded = multipathfinder( + ℓ, ndraws; nruns, ntasks=nthreads, ntasks_per_run=nthreads + ) + @test serial.draws == threaded.draws + @test serial.draw_component_ids == threaded.draw_component_ids + @test [c.μ for c in serial.fit_distribution.components] == + [c.μ for c in threaded.fit_distribution.components] + @test [c.Σ for c in serial.fit_distribution.components] == + [c.Σ for c in threaded.fit_distribution.components] + end + end end diff --git a/test/optimize.jl b/test/optimize.jl index 6f7a6c5bf..493d5e26c 100644 --- a/test/optimize.jl +++ b/test/optimize.jl @@ -1,18 +1,14 @@ using ADTypes using LinearAlgebra using Optim -using Optimization using OptimizationNLopt using Pathfinder using ProgressLogging using ReverseDiff using SciMLBase using Test -if isdefined(Optimization, :OptimizationState) - using Optimization: OptimizationState -else - using OptimizationBase: OptimizationState -end + +using OptimizationBase: OptimizationState @testset "build_optim_function" begin n = 20 diff --git a/test/runtests.jl b/test/runtests.jl index b03bccbf2..f01472419 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,7 +6,7 @@ Random.seed!(0) @testset "Pathfinder.jl" begin include("test_utils.jl") - include("transducers.jl") + include("utils.jl") include("woodbury.jl") include("optimize.jl") include("inverse_hessian.jl") diff --git a/test/singlepath.jl b/test/singlepath.jl index e8f7d69c7..1dde6dd4e 100644 --- a/test/singlepath.jl +++ b/test/singlepath.jl @@ -3,13 +3,11 @@ using Distributions using ForwardDiff using LinearAlgebra using Optim -using Optimization using Pathfinder using Random using ReverseDiff using SciMLBase using Test -using Transducers @testset "single path pathfinder" begin @testset "IsoNormal" begin @@ -18,14 +16,16 @@ using Transducers ndraws = 100 rngs = [MersenneTwister(), Random.default_rng()] seed = 42 - @testset for dim in [1, 5, 10, 100], rng in rngs - executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx() + @testset for dim in [1, 5, 10, 100], + rng in rngs, + ntasks in unique((1, Threads.nthreads())) + ℓ = build_logdensityproblem(logp, 5, 2) init = randn(dim) Random.seed!(rng, seed) # less restrictive type check to work around https://github.com/mlcolab/Pathfinder.jl/issues/142 # TODO: remove this workaround once the issue is fixed - result = @inferred PathfinderResult pathfinder(ℓ; init, ndraws, rng, executor) + result = @inferred PathfinderResult pathfinder(ℓ; init, ndraws, rng, ntasks) @test result isa PathfinderResult @test result.input === ℓ @test result.optim_prob isa SciMLBase.OptimizationProblem @@ -53,14 +53,14 @@ using Transducers argmax(getproperty.(result.elbo_estimates, :value)) Random.seed!(rng, seed) - result2 = pathfinder(ℓ; init, ndraws, rng, executor) + result2 = pathfinder(ℓ; init, ndraws, rng, ntasks) @test result2.fit_iteration == result.fit_iteration @test result2.draws == result.draws @test getproperty.(result2.elbo_estimates, :value) == getproperty.(result.elbo_estimates, :value) ndraws = 2 - result3 = pathfinder(ℓ; init, ndraws, executor) + result3 = pathfinder(ℓ; init, ndraws, ntasks) @test size(result3.draws) == (dim, ndraws) end end @@ -78,25 +78,23 @@ using Transducers logp(x) = -dot(x, P, x) / 2 dim = 5 ℓ = build_logdensityproblem(logp, dim, 2) - ndraws_elbo = 100 + ndraws_elbo = 500 rngs = [MersenneTwister(), Random.default_rng()] x = randn(dim) seed = 38 optimizer = Optim.LBFGS(; m=6) - @testset for rng in rngs - executor = rng isa MersenneTwister ? SequentialEx() : ThreadedEx() - + @testset for rng in rngs, ntasks in unique((1, Threads.nthreads())) Random.seed!(rng, seed) # less restrictive type check to work around https://github.com/mlcolab/Pathfinder.jl/issues/142 # TODO: remove this workaround once the issue is fixed result = @inferred PathfinderResult pathfinder( - ℓ; rng, optimizer, ndraws_elbo, executor + ℓ; rng, optimizer, ndraws_elbo, ntasks ) @test result.input === ℓ @test result.fit_distribution.Σ ≈ Σ rtol = 1e-1 @test result.optimizer == optimizer Random.seed!(rng, seed) - result2 = pathfinder(ℓ; rng, optimizer, ndraws_elbo, executor) + result2 = pathfinder(ℓ; rng, optimizer, ndraws_elbo, ntasks) @test result2.fit_distribution == result.fit_distribution @test result2.draws == result.draws @test getproperty.(result2.elbo_estimates, :value) == @@ -171,4 +169,36 @@ using Transducers @test_throws ArgumentError pathfinder(build_logdensityproblem(logp, 0, 2)) pathfinder(build_logdensityproblem(logp, 3, 2)) end + + @testset "reproducibility across ntasks" begin + logp(x) = -sum(abs2, x) / 2 + ℓ = build_logdensityproblem(logp, 5, 2) + seed = 17 + + @testset "explicit rng" begin + rng = MersenneTwister(seed) + serial = pathfinder(ℓ; rng, ndraws=100, ntasks=1) + Random.seed!(rng, seed) + threaded = pathfinder(ℓ; rng, ndraws=100, ntasks=Threads.nthreads()) + @test serial.draws == threaded.draws + @test serial.fit_iteration == threaded.fit_iteration + @test getproperty.(serial.elbo_estimates, :value) == + getproperty.(threaded.elbo_estimates, :value) + @test serial.fit_distribution.μ == threaded.fit_distribution.μ + @test serial.fit_distribution.Σ == threaded.fit_distribution.Σ + end + + @testset "default rng" begin + Random.seed!(seed) + serial = pathfinder(ℓ; ndraws=100, ntasks=1) + Random.seed!(seed) + threaded = pathfinder(ℓ; ndraws=100, ntasks=Threads.nthreads()) + @test serial.draws == threaded.draws + @test serial.fit_iteration == threaded.fit_iteration + @test getproperty.(serial.elbo_estimates, :value) == + getproperty.(threaded.elbo_estimates, :value) + @test serial.fit_distribution.μ == threaded.fit_distribution.μ + @test serial.fit_distribution.Σ == threaded.fit_distribution.Σ + end + end end diff --git a/test/transducers.jl b/test/transducers.jl deleted file mode 100644 index cd78c8f8c..000000000 --- a/test/transducers.jl +++ /dev/null @@ -1,15 +0,0 @@ -using Pathfinder -using Random -using Test -using Transducers - -@testset "transducers integration" begin - @testset "_findmax" begin - x = randn(100) - @test Pathfinder._findmax(x) == findmax(x) - @test Pathfinder._findmax(x |> Map(sin)) == findmax(sin.(x)) - @test Pathfinder._findmax([NaN, 3.0, 1.0]) === (3.0, 2) - @test Pathfinder._findmax([NaN, NaN, NaN]) === (NaN, 1) - @test Pathfinder._findmax([2.0, NaN, 4.0]) === (4.0, 3) - end -end diff --git a/test/utils.jl b/test/utils.jl new file mode 100644 index 000000000..9c0d93040 --- /dev/null +++ b/test/utils.jl @@ -0,0 +1,73 @@ +using Pathfinder +using Random +using Test + +@testset "parallel helpers" begin + @testset "_findmax_skipnan" begin + x = randn(100) + @test Pathfinder._findmax_skipnan(x) == findmax(x) + @test Pathfinder._findmax_skipnan(sin, x) == findmax(sin.(x)) + @test Pathfinder._findmax_skipnan([NaN, 3.0, 1.0]) === (3.0, 2) + @test Pathfinder._findmax_skipnan([NaN, NaN, NaN]) === (NaN, 1) + @test Pathfinder._findmax_skipnan([2.0, NaN, 4.0]) === (4.0, 3) + end + + @testset "_nchunks" begin + @test Pathfinder._nchunks(0, 4) == 1 + @test Pathfinder._nchunks(1, 4) == 1 + @test Pathfinder._nchunks(10, 1) == 1 + if Threads.nthreads() > 1 + @test Pathfinder._nchunks(10, 4) == 4 + @test Pathfinder._nchunks(2, 4) == 2 + end + end + + @testset "_maybe_tmap" begin + xs = randn(50) + @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) + @test Pathfinder._maybe_tmap(sin, xs, ntasks) == map(sin, xs) + @test Pathfinder._maybe_tmap(sin, Float64[], ntasks) == Float64[] + end + end + + @testset "_maybe_tmapreduce" begin + xs = [randn(3) for _ in 1:20] + @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) + @test Pathfinder._maybe_tmapreduce(identity, vcat, xs, ntasks) == + mapreduce(identity, vcat, xs) + # `+` on floats is non-associative, so parallel reduction order may differ + @test Pathfinder._maybe_tmapreduce(sum, +, xs, ntasks) ≈ mapreduce(sum, +, xs) + end + end + + @testset "_chunk_tmap" begin + xs = collect(1:50) + ys = randn(50) + @testset "ntasks=$ntasks" for ntasks in (1, Threads.nthreads()) + # matches sequential map, with correct multi-array pairing + @test Pathfinder._chunk_tmap(xs, ys; ntasks, setup=() -> 0) do _, x, y + return x + y + end == map(+, xs, ys) + # per-chunk state is passed through + @test Pathfinder._chunk_tmap(xs; ntasks, setup=() -> 10) do s, x + return s * x + end == map(x -> 10x, xs) + # empty arrays return empty + @test isempty(Pathfinder._chunk_tmap(Int[]; ntasks, setup=() -> 0) do _, x + return x + end) + end + + # the reproducible-seeding idiom yields identical output regardless of ntasks + function seeded(seed, n, ntasks) + rng = Random.seed!(Random.default_rng(), seed) + xs = collect(1:n) + seeds = rand!(rng, similar(xs, UInt64)) + return Pathfinder._chunk_tmap(xs, seeds; ntasks, setup=() -> copy(rng)) do r, _, s + Random.seed!(r, s) + return rand(r) + end + end + @test seeded(7, 50, 1) == seeded(7, 50, Threads.nthreads()) + end +end