Adding setting/getting seed for MOIRE

NEWS.md

@@ -1,3 +1,13 @@
+# moire (development version)
+
+## New Features
+
+* `run_mcmc(seed =)` sets the sampler seed. `NULL` draws from the current R RNG
+  a value that still fits after per-chain offsets. Independent MCMCs are spaced
+  by the number of parallel-tempering replicas. The seed used is returned as
+  `$seed` (printed if the run errors), with offsets as `$chain_seeds`. Replay
+  needs the same sampler settings; `num_cores` does not affect the draws.
+
 # moire 3.6.0
 
 ## New Features

R/mcmc.R

@@ -38,9 +38,9 @@
 #' the latent genotypes at each step of the MCMC. WARNING: This will increase
 #' the size of the output object significantly.
 #' @param num_chains Total number of chains to run, possibly simultaneously
-#' @param num_cores Total OMP parallel threads to use to run chains.
-#'  num_cores * pt_num_threads should not exceed the number of cores available
-#'  on your system.
+#' @param num_cores How many independent chains to run at once, passed to
+#'  [parallel::mclapply()] as `mc.cores`. `num_cores * pt_num_threads` should
+#'  not exceed the number of cores available on your system.
 #' @param pt_chains Total number of chains to run with parallel tempering or a
 #' vector containing the temperatures that should be used for parallel tempering.
 #' @param pt_grad Power to raise parallel tempering chains to. A value of 1
@@ -62,6 +62,12 @@
 #' adaptation steps. Only used if `adapt_temp` is TRUE.
 #' @param max_initialization_tries Number of times to try to initialize the
 #' chain before giving up
+#' @param seed Integer seed. `NULL` draws from the current R RNG a value that
+#'  still fits after per-chain offsets. Independent MCMCs are spaced by the
+#'  number of parallel-tempering replicas. The value used is returned as `$seed`
+#'  and printed if the run errors; per-chain offsets are `$chain_seeds`.
+#'  Replay needs the same sampler settings; `num_cores` does not affect the
+#'  draws.
 #' @param max_runtime Maximum runtime in minutes. If the MCMC is running for
 #' more than this amount of time, the function will stop and return the current
 #' state of the MCMC.
@@ -103,10 +109,29 @@ run_mcmc <-
            pre_adapt_steps = 25,
            temp_adapt_steps = 25,
            max_initialization_tries = 10000,
-           max_runtime = Inf) {
+           max_runtime = Inf,
+           seed = NULL) {
     start_time <- Sys.time()
-    args <- as.list(environment())
-    mcmc_args <- as.list(environment())
+    n_pt <- if (length(pt_chains) == 1L) {
+      as.integer(pt_chains)
+    } else {
+      length(pt_chains)
+    }
+    seeds <- resolve_mcmc_seeds(seed, num_chains, n_pt)
+    seed <- seeds$seed
+    chain_seeds <- seeds$chain_seeds
+
+    run_completed <- FALSE
+    on.exit({
+      if (!run_completed) {
+        message(sprintf(
+          "moire::run_mcmc() did not complete. Re-run with seed = %d to reproduce.",
+          seed
+        ))
+      }
+    }, add = TRUE)
+    args <- mget(names(formals()), envir = environment(), inherits = FALSE)
+    mcmc_args <- args
     mcmc_args$data <- data$data
     mcmc_args$sample_ids <- data$sample_ids
     mcmc_args$loci <- data$loci
@@ -159,6 +184,7 @@ run_mcmc <-
         1:num_chains,
         function(i) {
           mcmc_args$chain_number <- i
+          mcmc_args$seed <- chain_seeds[[i]]
           mcmc_args$simple_verbose <- (mcmc_args$num_chains > 1)
           mcmc_args$samples <- round(mcmc_args$samples_per_chain)
           handle_chain_result(run_mcmc_rcpp(mcmc_args))
@@ -177,6 +203,44 @@ run_mcmc <-
     res$runtime <- end_time - start_time
     res$chains <- chains
     res$args <- args
+    res$seed <- seed
+    res$chain_seeds <- chain_seeds
 
+    run_completed <- TRUE
     res
   }
+
+# PT replicas use parent+0, parent+1, ... so independent MCMCs are spaced
+# by the replica count. See Seed::pt_replica in src/seed.h.
+resolve_mcmc_seeds <- function(seed, num_chains, n_pt) {
+  n_pt <- as.integer(n_pt)
+  if (length(n_pt) != 1L || is.na(n_pt) || n_pt < 1L) {
+    stop("`pt_chains` must specify at least one temperature")
+  }
+  seed_stride <- n_pt
+  offset_span <- as.double(seed_stride) * max(num_chains - 1L, 0L)
+  if (offset_span > .Machine$integer.max - 1) {
+    stop("`num_chains` is too large to form integer chain seeds")
+  }
+  max_seed <- as.integer(.Machine$integer.max - offset_span)
+  if (max_seed < 1L) {
+    stop("`num_chains` is too large to form integer chain seeds")
+  }
+  if (is.null(seed)) {
+    seed <- sample.int(max_seed, 1L)
+  }
+  seed <- as.integer(seed)
+  if (length(seed) != 1L || is.na(seed)) {
+    stop("`seed` must be a single integer")
+  }
+  if (seed > max_seed) {
+    stop(
+      "`seed` is too large for num_chains = ", num_chains,
+      " with ", n_pt, " parallel-tempering replicas; maximum is ", max_seed
+    )
+  }
+  list(
+    seed = seed,
+    chain_seeds = seed + (seq_len(num_chains) - 1L) * seed_stride
+  )
+}

README.Rmd

@@ -40,7 +40,7 @@ df <- read.csv("your_data.csv")
 data <- load_long_form_data(df)
 
 # With data in appropriate format, run MCMC as follows
-mcmc_results <- moire::run_mcmc(data, is_missing = data$is_missing)
+mcmc_results <- moire::run_mcmc(data, is_missing = data$is_missing, seed = 1)
 ```
 
 ## Manuscript

README.md

@@ -39,7 +39,7 @@ df <- read.csv("your_data.csv")
 data <- load_long_form_data(df)
 
 # With data in appropriate format, run MCMC as follows
-mcmc_results <- moire::run_mcmc(data, is_missing = data$is_missing)
+mcmc_results <- moire::run_mcmc(data, is_missing = data$is_missing, seed = 1)
 ```
 
 ## Manuscript

data-raw/mcmc_results.R

@@ -33,7 +33,7 @@ mcmc_results <- moire::run_mcmc(
   simulated_data,
   verbose = TRUE, burnin = burnin, samples_per_chain = num_samples,
   pt_chains = pt_chains, pt_num_threads = pt_num_threads,
-  adapt_temp = TRUE
+  adapt_temp = TRUE, seed = 17325
 )
 
 ## ----save_results

src/chain.cpp

@@ -1291,8 +1291,9 @@ void Chain::set_temp(float temp) { this->temp = temp; }
 
 float Chain::get_temp() { return this->temp; }
 
-Chain::Chain(GenotypingData genotyping_data, Parameters params, float temp)
-    : genotyping_data(genotyping_data), params(params), sampler()
+Chain::Chain(GenotypingData genotyping_data, Parameters params, float temp,
+             std::uint32_t seed)
+    : genotyping_data(genotyping_data), params(params), sampler(seed)
 
 {
     p_prop_var = std::vector<std::vector<float>>(genotyping_data.num_loci);

src/chain.h

@@ -3,6 +3,7 @@
 #ifndef CHAIN_H_
 #define CHAIN_H_
 
+#include <cstdint>
 #include "combination_indices_generator.h"
 #include "genotyping_data.h"
 #include "parameters.h"
@@ -141,8 +142,8 @@ class Chain
 
     std::vector<int> sample_accept{};
 
-    Chain(GenotypingData genotyping_data, Parameters params, float temp = 1.0);
-    Chain() {};
+    Chain(GenotypingData genotyping_data, Parameters params, float temp,
+          std::uint32_t seed);
     void update_m(int iteration);
     void update_r(int iteration);
     void update_m_r(int iteration);

src/mcmc.cpp

@@ -2,10 +2,12 @@
 
 #include "include/spline/src/spline.h"
 #include "mcmc_utils.h"
+#include "seed.h"
 
 #include <Rcpp.h>
 
 #include <cmath>
+#include <memory>
 
 #ifdef _OPENMP
 #include <omp.h>
@@ -40,7 +42,9 @@ void record_ill_conditioned_start(Chain &chain,
 }  // namespace
 
 MCMC::MCMC(GenotypingData genotyping_data, Parameters params)
-    : genotyping_data(genotyping_data), params(params)
+    : genotyping_data(genotyping_data),
+      params(params),
+      swap_sampler(Seed::swap_sampler(params.seed))
 {
     p_store.resize(genotyping_data.num_loci);
     latent_genotypes_store.resize(genotyping_data.num_samples);
@@ -59,14 +63,15 @@ MCMC::MCMC(GenotypingData genotyping_data, Parameters params)
 
     omp_set_num_threads(params.pt_num_threads);
 
-    chains.resize(params.pt_chains.size());
-    std::vector<bool> any_ill_conditioned(params.pt_chains.size(), false);
-    std::vector<InitializationDiagnostics> per_chain_diag(params.pt_chains.size());
-    std::vector<int> chains_attempted_flags(params.pt_chains.size(), 0);
+    const size_t n_pt = params.pt_chains.size();
+    std::vector<std::unique_ptr<Chain>> constructed(n_pt);
+    std::vector<bool> any_ill_conditioned(n_pt, false);
+    std::vector<InitializationDiagnostics> per_chain_diag(n_pt);
+    std::vector<int> chains_attempted_flags(n_pt, 0);
     temp_gradient = params.pt_chains;
 
     #pragma omp parallel for
-    for (size_t i = 0; i < params.pt_chains.size(); i++)
+    for (size_t i = 0; i < n_pt; i++)
     {
         if (std::any_of(any_ill_conditioned.begin(), any_ill_conditioned.end(),
                         [](bool v) { return v; }))
@@ -76,25 +81,26 @@ MCMC::MCMC(GenotypingData genotyping_data, Parameters params)
 
         chains_attempted_flags[i] = 1;
         float temp = params.pt_chains[i];
-        chains[i] = Chain(genotyping_data, params, temp);
+        constructed[i] = std::make_unique<Chain>(
+            genotyping_data, params, temp, Seed::pt_replica(params.seed, i));
 
-        bool ill_conditioned = !std::isfinite(chains[i].get_llik());
+        bool ill_conditioned = !std::isfinite(constructed[i]->get_llik());
         int ill_count = 0;
         if (ill_conditioned)
         {
-            record_ill_conditioned_start(chains[i], per_chain_diag[i]);
+            record_ill_conditioned_start(*constructed[i], per_chain_diag[i]);
             ++ill_count;
         }
 
         int max_tries = params.max_initialization_tries;
         while (ill_conditioned and max_tries > 0)
         {
-            chains[i].initialize_parameters();
+            constructed[i]->initialize_parameters();
             max_tries--;
-            ill_conditioned = !std::isfinite(chains[i].get_llik());
+            ill_conditioned = !std::isfinite(constructed[i]->get_llik());
             if (ill_conditioned)
             {
-                record_ill_conditioned_start(chains[i], per_chain_diag[i]);
+                record_ill_conditioned_start(*constructed[i], per_chain_diag[i]);
                 ++ill_count;
             }
         }
@@ -151,6 +157,12 @@ MCMC::MCMC(GenotypingData genotyping_data, Parameters params)
         return;
     }
 
+    chains.reserve(n_pt);
+    for (auto &p : constructed)
+    {
+        chains.push_back(std::move(*p));
+    }
+
     std::iota(swap_indices.begin(), swap_indices.end(), 0);
 };
 
@@ -209,7 +221,7 @@ void MCMC::swap_chains(int step, bool burnin)
             swap_barriers[i] += 1.0 - acceptanceRate;
         }
 
-        float u = log(R::runif(0, 1));
+        float u = swap_sampler.sample_log_mh_acceptance();
 
         if ((acceptanceRatio > 0 || u < acceptanceRatio) and
             !std::isnan(acceptanceRatio))

src/mcmc.h

@@ -18,6 +18,7 @@ class MCMC
    public:
     GenotypingData genotyping_data;
     Parameters params;
+    Sampler swap_sampler;
     std::vector<Chain> chains{};
     bool initialization_failed = false;
     InitializationDiagnostics initialization_diagnostics{};

src/parameters.cpp

@@ -16,6 +16,7 @@ Parameters::Parameters(const Rcpp::List &args)
     samples = UtilFunctions::r_to_int(args["samples"]);
     pt_chains = UtilFunctions::r_to_vector_float(args["pt_chains"]);
     pt_num_threads = UtilFunctions::r_to_int(args["pt_num_threads"]);
+    seed = static_cast<std::uint32_t>(UtilFunctions::r_to_int(args["seed"]));
     adapt_temp = UtilFunctions::r_to_bool(args["adapt_temp"]);
     pre_adapt_steps = UtilFunctions::r_to_int(args["pre_adapt_steps"]);
     temp_adapt_steps = UtilFunctions::r_to_int(args["temp_adapt_steps"]);

src/parameters.h

@@ -1,6 +1,7 @@
 #ifndef PARAMETERS_H_
 #define PARAMETERS_H_
 
+#include <cstdint>
 #include <Rcpp.h>
 #include <vector>
 
@@ -18,6 +19,7 @@ class Parameters
     int samples;
     std::vector<float> pt_chains;
     int pt_num_threads;
+    std::uint32_t seed;
     bool adapt_temp;
     int pre_adapt_steps;
     int temp_adapt_steps;