Support multi-chain sampling via future.apply mapping

[jennajiali]

Overview

This PR resolves the immediate goals of issue #18. It targets the in-scope subset agreed upon in the meeting: N independent chains sampling the same target distribution, running the full warm-up and main procedure on N processors, so wall-clock time buys N× samples. Parallel-tempering-style inter-chain communication is explicitly out of scope for this change.

Design decisions & implementation

Following the guidance in the issue, sample_chain() is kept as a single-chain function and made pleasant to future_lapply over, rather than wrapping the parallel loop internally.

• chain_index argument: a new leading argument chain_index accepts an optional integer identifier for the chain. Placing chain_index first makes future.apply::future_lapply(seq_len(n_chain), sample_chain, ...) work out of the box, since future_lapply passes each iteration variable as the first positional argument to the function. When non-NULL, it is prefixed onto progress-bar labels and fallback progress messages, so output from several chains running in parallel remains distinguishable.
• initial_state as a generator: initial_state now additionally accepts a generator function such as stats::rnorm or stats::runif. The function is called with target_distribution$dimension to produce a different random starting position for each chain.
• Target helper updates: target_distribution_from_log_density_formula() and target_distribution_from_stan_model() now populate a dimension field on the returned list, unlocking the generator-function form of initial_state without further user input.
• Parallel-safe progress routing: progress bars now route through the optional progressr package when it is installed, keeping progress reporting well-behaved under parallel futures. The existing progress package and message()-based fallbacks continue to work when progressr is unavailable.
• Result collation: added the exported combine_chain_results() function, which inverts the list-of-lists returned by future_lapply into a single list keyed by result-field name, with each entry a per-chain list. This makes downstream diagnostic workflows (e.g. with the posterior package) simpler to express.

Usage examples

1. Sampling four chains in parallel with future.apply:

library(future.apply)
plan(multisession)

target_distribution <- list(
  log_density = function(x) -sum(x^2) / 2,
  gradient_log_density = function(x) -x,
  dimension = 2
)

results <- future_lapply(
  1:4,
  sample_chain,
  target_distribution = target_distribution,
  initial_state = stats::rnorm,      # generator: different start per chain
  n_warm_up_iteration = 1000,
  n_main_iteration = 1000,
  future.seed = 876287L
) |> combine_chain_results()

# results$traces is a length-4 list, one traces matrix per chain
length(results$traces)

2. Sequential multi-chain run without a future backend (useful for debugging before scaling out to plan(multisession)):
plain base-R lapply instead of future_lapply — no parallelism, everything runs one chain after another in your current R session. This is the "debug it here first" version: if something's wrong with your target distribution or adapter setup, you get a normal R traceback instead of a confusing error from inside a background worker. Once this works, swapping lapply → future_lapply (plus plan(multisession)) is the only change needed to parallelize it.

per_chain <- lapply(
  1:4,
  sample_chain,
  target_distribution = target_distribution,
  initial_state = stats::rnorm,
  n_warm_up_iteration = 1000,
  n_main_iteration = 1000,
  show_progress_bar = FALSE
)
combined <- combine_chain_results(per_chain)

3. Distinguishing chains in progress output, e.g. when running a single chain manually or outside of future_lapply:
This isn't really about parallelism — it shows that chain_index is just an optional label you can set by hand, even for a single ad-hoc call. The only visible effect is the progress output.

sample_chain(
  chain_index = 3L,
  target_distribution = target_distribution,
  initial_state = stats::rnorm,
  n_warm_up_iteration = 500,
  n_main_iteration = 500
)
# Fallback message output is now prefixed, e.g.:
# "Chain 3 | Warm-up (stage 1/1): 20% done (100/500 iterations) | elapsed: 0.4s"

vs. without chain_index, you'd just see Warm-up (stage 1/1): 20% done ... with no "Chain 3 |" prefix. Useful if you're eyeballing interleaved log output from several chains and need to tell which line came from which.

4. Parallel-safe progress reporting with progressr (recommended when running many chains under future.apply, since a local progress:: bar isn't meaningful across worker processes):

library(progressr)
handlers(global = TRUE)

with_progress({
  results <- future_lapply(
    1:8,
    sample_chain,
    target_distribution = target_distribution,
    initial_state = stats::rnorm,
    n_warm_up_iteration = 2000,
    n_main_iteration = 2000,
    future.seed = 876287L
  )
})

We have each worker report progress events back to the main process, which then renders one unified bar/message locally. Two pieces:

• handlers(global = TRUE) — registers how progress should be displayed (default handler prints text updates to console).
• with_progress({ ... }) — wraps the code whose progress you want tracked. Inside chain_loop(), rmcmc checks whether it's running inside a with_progress() block and, if so, calls progressr::progressor() instead of progress::progress_bar$new().

You don't have to change anything about the sample_chain() calls themselves — just wrap them in with_progress() and install the progressr package. If you don't do this and just run Example 1 as-is, it still works fine, you just won't see live progress from the workers.

5. Combining chains that also trace warm-up (combine_chain_results() preserves the extra warm_up_traces / warm_up_statistics fields when trace_warm_up = TRUE):

per_chain <- lapply(1:2, function(i) {
  sample_chain(
    chain_index = i,
    target_distribution = target_distribution,
    initial_state = stats::rnorm,
    n_warm_up_iteration = 200,
    n_main_iteration = 200,
    trace_warm_up = TRUE,
    show_progress_bar = FALSE
  )
})
combined <- combine_chain_results(per_chain)
names(combined)
#> [1] "final_state"        "traces"             "statistics"
#> [4] "warm_up_traces"     "warm_up_statistics"

Note the function(i) { sample_chain(chain_index = i, ...) } wrapper here instead of passing sample_chain directly to lapply. That's because we want to pass i specifically into the named chain_index argument (to also demonstrate setting it), rather than relying on positional matching like Example 1 did. Both approaches work — this is just the more explicit.

[codecov]

Codecov Report

❌ Patch coverage is 93.39623% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 98.81%. Comparing base (eb90251) to head (d1bb047).

Files with missing lines	Patch %	Lines
R/chains.R	93.13%	7 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##              main     #139      +/-   ##
===========================================
- Coverage   100.00%   98.81%   -1.19%     
===========================================
  Files           11       11              
  Lines          762      842      +80     
===========================================
+ Hits           762      832      +70     
- Misses           0       10      +10     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.