Implement staged warm-up adaptation schedule

[jennajiali]

This PR implements the staged warm-up adaptation interface proposed in issue #79, allowing users to specify different adapter configurations at different points during warm-up. It builds on the partial implementation started on the mmg/adapter-staging branch and ensures adapter states correctly carry over across stage boundaries.

Background

Previously, sample_chain() ran all adapters from iteration 1 of warm-up with no way to vary which adapters were active over time. This is suboptimal in practice: the shape adapter (which estimates the proposal covariance) should not start until the scale adapter has had enough iterations to stabilise, because early covariance estimates based on poorly-mixed samples can make the sampler worse rather than better.

Furthermore, previously all adapter initialize() functions ignored the current proposal parameters on re-initialisation, causing the scale and shape estimates accumulated during one staged warm-up stage to be discarded at the start of the next. This contradicted the requirement in issue #79 that adaptation parameters be initialised at the values at the end of the previous stage.

Changes

R/chains.R

• check_and_process_adapters() — new internal function that normalises the adapters argument into a canonical list of stage specifications, each with $adapters (a list of adapter objects) and $n_iteration (an integer). Accepts three input forms:
  1. A flat list of adapter objects (existing behaviour, now wrapped into a single stage)
  2. A staged list where each element is a list of adapters with an optional trailing integer iteration count; the last stage may omit the count and receives the remaining warm-up iterations
  3. A function accepting n_warm_up_iteration and returning a staged list (used by progressive_adaptation_schedule())
• sample_chain() — warm-up loop now iterates over stages produced by check_and_process_adapters(), passing each stage's adapter list and iteration count to chain_loop() separately. Several bugs in the initial partial implementation are fixed:
  • Bootstrap warm_up_results initialised with key $final_state (not $state) to be consistent with chain_loop()'s return value
  • Each stage now reads stage$n_iteration rather than always using the full n_warm_up_iteration
  • Each stage now passes stage$adapters (clean list of adapter objects) rather than the whole named stage list
  • Progress bar label now shows current stage number (e.g. Warm-up (stage 1/3))
• combine_warm_up_results() — rewritten to use rbind instead of mapply(c, ...). This correctly handles the bootstrap case (rbind(NULL, matrix) returns the matrix unchanged) and uses the correct $final_state key throughout
• combine_stage_results() — fixed $state → $final_state to match chain_loop()'s return value
• @param adapters documentation — updated to describe all three input forms

R/adaptation.R

• progressive_adaptation_schedule() — new exported convenience function that returns a schedule constructor. When called with n_warm_up_iteration, it produces a three-stage adaptation schedule following common practice:
  • Stage 1: scale adapter only (shape held fixed at identity), for n_fixed_shape_iteration iterations (default 50)
  • Stage 2: scale adapter + diagonal shape adapter (shape_adapter("variance")), for n_diagonal_shape_iteration iterations (default 50)
  • Stage 3: scale adapter + dense shape adapter (shape_adapter("covariance")), for remaining iterations
  • If the sum of stage 1 and stage 2 counts exceeds n_warm_up_iteration, both are reduced proportionally and stage 3 is skipped. All adapter objects and iteration counts are customisable via arguments. Includes full roxygen documentation and a working example.
• stochastic_approximation_scale_adapter and dual_averaging_scale_adapter: when initial_scale is not user-specified, initialize() now reads proposal$parameters()$scale and uses it as the starting log-scale if available, falling back to proposal$default_initial_scale() only when no current scale exists. Note that for dual_averaging_scale_adapter, mu, smoothed_log_scale and accept_prob_error already persisted across stages via closure scope and required no change.
• variance_shape_adapter and covariance_shape_adapter: initialize() now follows a three-priority order:
  1. Explicit initial_shape constructor argument if supplied.
  2. Current proposal shape, guarded by type compatibility checks. variance_shape_adapter requires a vector of matching length to prevent accidentally carrying over a matrix-valued shape. covariance_shape_adapter requires a matrix of matching dimensions (Cholesky factor) to prevent carrying over a vector-valued shape.
  3. Default fallback (unit variances or identity matrix, respectively).
• Added initial_shape parameter (numeric vector of per-dimension scales for variance; lower-triangular Cholesky factor matrix for covariance) to the respective function signatures and roxygen documentation, addressing the request to allow users to supply domain-knowledge-based starting values.

tests/testthat/test-chains.R

New tests appended after existing tests:

• Unit tests for check_and_process_adapters() covering all three input forms, the last-stage remainder logic, and all error cases (non-last stage missing count, counts not summing to n_warm_up_iteration, invalid input types)
• Integration tests for sample_chain() with staged adapters: correct warm-up row counts across two and three stages, correct behaviour with differing adapter sets across stages when trace_warm_up = FALSE, function-form adapters argument, correct final_state type, and invalid adapters error propagation

tests/testthat/test-adaptation.R

New tests appended after existing tests covering progressive_adaptation_schedule():

• Returns a function
• Correct stage count and iteration counts for the normal case (n_warm_up > n_fixed + n_diagonal), the exact-boundary case (no dense stage), and the fallback case (n_warm_up < n_fixed + n_diagonal)
• Edge case of n_warm_up_iteration = 1
• Custom n_fixed_shape_iteration and n_diagonal_shape_iteration are respected
• Correct adapter count in each stage
• Schedule output always parses cleanly through check_and_process_adapters() across a range of n_warm_up_iteration values
• End-to-end sample_chain() integration test

New tests covering state carry-over across stages:

• Four tests for scale adapter carry-over: both stochastic_approximation and dual_averaging variants are checked for (a) reading the current proposal scale on re-initialisation and (b) explicit initial_scale still taking priority over the proposal's current value.
• Four tests for variance_shape_adapter: carry-over from a vector proposal shape; explicit initial_shape override; NULL proposal falls back to unit variances; matrix-valued proposal shape (incompatible type) falls back to unit variances.
• Four tests for covariance_shape_adapter: carry-over from a matrix proposal shape; explicit initial_shape override; NULL proposal falls back to identity; vector-valued proposal shape (incompatible type) falls back to identity.
• One end-to-end integration test via sample_chain with trace_warm_up=TRUE: verifies that the log_scale in warm_up_statistics does not jump back to the default value at the stage boundary between two consecutive scale-adapter-only stages.

Usage examples

Fully custom staged schedule:

results <- sample_chain(
  target_distribution,
  initial_state = rnorm(2),
  n_warm_up_iteration = 1000,
  n_main_iteration = 1000,
  adapters = list(
    list(scale_adapter(), 50),                              # scale only
    list(scale_adapter(), shape_adapter("variance"), 50),   # scale + diagonal shape
    list(scale_adapter(), shape_adapter("covariance"))      # scale + dense shape (remainder)
  )
)

Using the convenience constructor with defaults:

results <- sample_chain(
  target_distribution,
  initial_state = rnorm(2),
  n_warm_up_iteration = 1000,
  n_main_iteration = 1000,
  adapters = progressive_adaptation_schedule()
)

Existing flat-list usage is unchanged:

results <- sample_chain(
  target_distribution,
  initial_state = rnorm(2),
  n_warm_up_iteration = 1000,
  n_main_iteration = 1000,
  adapters = list(scale_adapter(), shape_adapter())  # still works as before
)

Known limitation

When trace_warm_up = TRUE and stages use different adapter sets (e.g. scale-only in stage 1, scale + shape in stage 2), the warm-up statistics matrices from each stage have different column counts and cannot be rbind-ed. This causes an error. A future PR could address this by filling NA for missing columns. For now, trace_warm_up = TRUE with staged adapters is only safe when all stages use the same adapter set. This limitation is documented in the new tests.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (65b86ab) to head (a700117).

Additional details and impacted files

@@            Coverage Diff             @@
##              main      #127    +/-   ##
==========================================
  Coverage   100.00%   100.00%            
==========================================
  Files           11        11            
  Lines          620       722   +102     
==========================================
+ Hits           620       722   +102     

☔ 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.

[matt-graham]

This is looke great @jennajiali - I've added some suggested changes but mainly minor things. Will also need to resolve the linter warnings

[matt-graham]

Thanks @jennajiali for addressing previous comments. I've added a few more small suggestions to resolve the linter and test coverage fails in the checks. Let me know if anything doesn't make sense.

[matt-graham]

Thanks @jennajiali for making the additional changes. Factoring out the process_stage_specifications function I think has helped improve the readability a lot - well done for also catching and fixing the unintentional bug in my suggested changes in the handling of the form 1 case 😅. The additional tests for edge cases are also great. I merged in the changes from #130 and made one small fix to a use of = rather than <- that was causing styler check in pre-commit to fail. This all looks good to go to me now so will merge in.

[matt-graham]

Great idea to have this as an additional test case!

[matt-graham]

This is also a great test to have - it's a testament to the generality of your implementation that this works as expected!