Drop Transducers/Folds in favor of OhMyThreads and drop Optimization umbrella

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 <seth.axen@gmail.com> 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"

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"

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(;

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.

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)

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

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
 

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

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
 

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)
 

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