Fix missing MH proposal-SD adaptation under NUTS for GGM and MixedMRF

src/mcmc/execution/warmup_schedule.h

@@ -2,6 +2,8 @@
 
 #include <vector>
 #include <algorithm>
+#include <cmath>
+#include <optional>
 
 
 /**
@@ -167,6 +169,24 @@ struct WarmupSchedule {
     return learn_proposal_sd && !stage3b_skipped && in_stage3b(i);
   }
 
+  /// Robbins-Monro decay rate for proposal-SD adaptation. Single source of
+  /// truth; every model's tune_proposal_sd consults this.
+  static constexpr double proposal_sd_rm_decay = 0.75;
+
+  /// Robbins-Monro weight for proposal-SD adaptation at the given iteration.
+  ///
+  /// Returns the RM weight (1-indexed since stage 3b began, decay
+  /// `proposal_sd_rm_decay`) iff `adapt_proposal_sd(iter)` is true. Returns
+  /// nullopt otherwise. Use this in every `tune_proposal_sd` rather than
+  /// computing the weight inline — keeps the policy (which iterations adapt,
+  /// what the weight schedule looks like) in one place and makes stage-3b
+  /// adaptation sampler-agnostic by construction.
+  std::optional<double> rm_weight_for_proposal_sd(int iter) const {
+    if (!adapt_proposal_sd(iter)) return std::nullopt;
+    const double t = static_cast<double>(iter - stage3b_start + 1);
+    return std::pow(t, -proposal_sd_rm_decay);
+  }
+
   /// Current Stage-2 window index (-1 outside Stage-2)
   int current_window(int i) const {
     for (size_t k = 0; k < window_ends.size(); ++k)

src/models/ggm/ggm_model.cpp

@@ -673,10 +673,10 @@ double GGMModel::log_density_impl_diag(size_t j) const {
 
 }
 
-void GGMModel::update_edge_parameter(size_t i, size_t j, int iteration) {
+double GGMModel::update_edge_parameter(size_t i, size_t j) {
 
     if (edge_indicators_(i, j) == 0) {
-        return; // Edge is not included; skip update
+        return 0.0; // Edge is not included; skip update (AR irrelevant, masked out)
     }
 
     get_constants(i, j);
@@ -721,12 +721,7 @@ void GGMModel::update_edge_parameter(size_t i, size_t j, int iteration) {
         cholesky_update_after_edge(omega_ij_old, omega_jj_old, i, j);
     }
 
-    // Robbins-Monro proposal-SD adaptation (warmup only)
-    if (iteration >= 1 && iteration < total_warmup_) {
-        double rm_weight = std::pow(iteration, -0.75);
-        proposal_sds_(e) = update_proposal_sd_with_robbins_monro(
-            proposal_sds_(e), ln_alpha, rm_weight, target_accept_);
-    }
+    return std::min(1.0, std::exp(ln_alpha));
 }
 
 void GGMModel::cholesky_update_after_edge(double omega_ij_old, double omega_jj_old, size_t i, size_t j)
@@ -768,7 +763,7 @@ void GGMModel::cholesky_update_after_edge(double omega_ij_old, double omega_jj_o
 
 }
 
-void GGMModel::update_diagonal_parameter(size_t i, int iteration) {
+double GGMModel::update_diagonal_parameter(size_t i) {
     double logdet_omega = cholesky_helpers::get_log_det(cholesky_of_precision_);
     double logdet_omega_sub_ii = logdet_omega + MY_LOG(covariance_matrix_(i, i));
 
@@ -793,12 +788,7 @@ void GGMModel::update_diagonal_parameter(size_t i, int iteration) {
         cholesky_update_after_diag(omega_ii, i);
     }
 
-    // Robbins-Monro proposal-SD adaptation (warmup only)
-    if (iteration >= 1 && iteration < total_warmup_) {
-        double rm_weight = std::pow(iteration, -0.75);
-        proposal_sds_(e) = update_proposal_sd_with_robbins_monro(
-            proposal_sds_(e), ln_alpha, rm_weight, target_accept_);
-    }
+    return std::min(1.0, std::exp(ln_alpha));
 }
 
 void GGMModel::cholesky_update_after_diag(double omega_ii_old, size_t i)
@@ -955,22 +945,41 @@ void GGMModel::update_edge_indicator_parameter_pair(size_t i, size_t j) {
 }
 
 void GGMModel::do_one_metropolis_step(int iteration) {
+    // Collect per-slot accept probabilities for the Robbins-Monro adapter.
+    // proposal_sds_ is stored as a flat dim_-length vec indexed by the
+    // upper-triangle scheme `e = j * (j + 1) / 2 + i`; we mirror that here
+    // as a dim_ x 1 matrix.
+    arma::mat accept_prob(dim_, 1, arma::fill::zeros);
+    arma::umat index_mask(dim_, 1, arma::fill::zeros);
 
     // Update off-diagonals (upper triangle)
     for (size_t i = 0; i < p_ - 1; ++i) {
         for (size_t j = i + 1; j < p_; ++j) {
-            update_edge_parameter(i, j, iteration);
+            double ap = update_edge_parameter(i, j);
+            if (edge_indicators_(i, j) == 1) {
+                size_t e = j * (j + 1) / 2 + i;
+                accept_prob(e, 0) = ap;
+                index_mask(e, 0) = 1;
+            }
         }
     }
 
     // Update diagonals
     for (size_t i = 0; i < p_; ++i) {
-        update_diagonal_parameter(i, iteration);
+        double ap = update_diagonal_parameter(i);
+        size_t e = i * (i + 3) / 2;
+        accept_prob(e, 0) = ap;
+        index_mask(e, 0) = 1;
+    }
+
+    if (metropolis_adapter_) {
+        metropolis_adapter_->update(index_mask, accept_prob, iteration);
     }
 }
 
 void GGMModel::init_metropolis_adaptation(const WarmupSchedule& schedule) {
-    total_warmup_ = schedule.total_warmup;
+    metropolis_adapter_ = std::make_unique<MetropolisAdaptationController>(
+        proposal_sds_, schedule, target_accept_);
 }
 
 void GGMModel::prepare_iteration() {
@@ -1000,12 +1009,10 @@ void GGMModel::update_edge_indicators() {
 }
 
 void GGMModel::tune_proposal_sd(int iteration, const WarmupSchedule& schedule) {
-    if (!schedule.adapt_proposal_sd(iteration)) return;
-
+    auto rm_weight_opt = schedule.rm_weight_for_proposal_sd(iteration);
+    if (!rm_weight_opt) return;
+    const double rm_weight = *rm_weight_opt;
     const double target_accept = target_accept_;
-    const double rm_decay = 0.75;
-    double t = iteration - schedule.stage3b_start + 1;
-    double rm_weight = std::pow(t, -rm_decay);
 
     // Off-diagonal sweeps
     for (size_t i = 0; i < p_ - 1; ++i) {

src/models/ggm/ggm_model.h

@@ -8,6 +8,7 @@
 #include "models/ggm/graph_constraint_structure.h"
 #include "models/ggm/ggm_gradient.h"
 #include "priors/parameter_prior.h"
+#include "mcmc/samplers/metropolis_adaptation.h"
 
 
 /**
@@ -60,7 +61,7 @@ class GGMModel : public BaseModel {
         edge_indicators_(initial_edge_indicators),
         vectorized_parameters_(dim_),
         vectorized_indicator_parameters_(edge_selection_ ? dim_ : 0),
-        proposal_sds_(arma::vec(dim_, arma::fill::ones) * 0.25),
+        proposal_sds_(arma::mat(dim_, 1, arma::fill::ones) * 0.25),
         num_pairwise_(p_ * (p_ - 1) / 2),
         observations_(na_impute ? observations : arma::mat()),
         precision_proposal_(arma::mat(p_, p_, arma::fill::none))
@@ -108,7 +109,7 @@ class GGMModel : public BaseModel {
         edge_indicators_(initial_edge_indicators),
         vectorized_parameters_(dim_),
         vectorized_indicator_parameters_(edge_selection_ ? dim_ : 0),
-        proposal_sds_(arma::vec(dim_, arma::fill::ones) * 0.25),
+        proposal_sds_(arma::mat(dim_, 1, arma::fill::ones) * 0.25),
         num_pairwise_(p_ * (p_ - 1) / 2),
         precision_proposal_(arma::mat(p_, p_, arma::fill::none))
     {
@@ -139,7 +140,6 @@ class GGMModel : public BaseModel {
           vectorized_parameters_(other.vectorized_parameters_),
           vectorized_indicator_parameters_(other.vectorized_indicator_parameters_),
           proposal_sds_(other.proposal_sds_),
-          total_warmup_(other.total_warmup_),
           shuffled_edge_order_(other.shuffled_edge_order_),
           num_pairwise_(other.num_pairwise_),
           rng_(other.rng_),
@@ -191,9 +191,6 @@ class GGMModel : public BaseModel {
         edge_selection_active_ = active;
     }
 
-    /** Store warmup length for Robbins-Monro proposal-SD adaptation. */
-    void init_metropolis_adaptation(const WarmupSchedule& schedule) override;
-
     /**
      * Set the Robbins-Monro target acceptance rate used by the
      * adaptive-Metropolis updates of this GGM. Honoured by all
@@ -203,6 +200,13 @@ class GGMModel : public BaseModel {
         target_accept_ = target;
     }
 
+    /**
+     * Construct Robbins-Monro adaptation controller for the per-iteration
+     * MH proposal SDs. Called once by MetropolisSampler before warmup;
+     * under NUTS this is never called and the controller stays null.
+     */
+    void init_metropolis_adaptation(const WarmupSchedule& schedule) override;
+
     /** Shuffle edge visit order (random scan). */
     void prepare_iteration() override;
 
@@ -456,6 +460,10 @@ class GGMModel : public BaseModel {
     // to 0.44 (componentwise random-walk Metropolis optimum).
     double target_accept_ = 0.44;
 
+    /// Per-iteration adaptation controller (MH mode only — under NUTS this
+    /// stays null and the stage-3b path in tune_proposal_sd is used instead).
+    std::unique_ptr<MetropolisAdaptationController> metropolis_adapter_;
+
     /** Extract upper triangle of the precision matrix into a vector. */
     arma::vec extract_upper_triangle() const {
         arma::vec result(dim_);
@@ -500,10 +508,10 @@ class GGMModel : public BaseModel {
     /// Pre-allocated storage returned by get_vectorized_indicator_parameters().
     arma::ivec vectorized_indicator_parameters_;
 
-    /// Proposal standard deviations for Metropolis updates (one per element).
-    arma::vec proposal_sds_;
-    /// Total number of warmup iterations (for Robbins-Monro adaptation).
-    int total_warmup_ = 0;
+    /// Proposal standard deviations for Metropolis updates (one per element,
+    /// stored as a (dim_, 1) matrix so it can be wrapped by
+    /// MetropolisAdaptationController).
+    arma::mat proposal_sds_;
 
     /// Shuffled edge visit order for random-scan edge selection.
     arma::uvec shuffled_edge_order_;
@@ -562,21 +570,22 @@ class GGMModel : public BaseModel {
      * on an unconstrained reparameterization. Accepts or rejects with a
      * Metropolis ratio using the Gaussian likelihood and Cauchy prior.
      *
-     * @param i          Row index (i < j)
-     * @param j          Column index
-     * @param iteration  Current iteration (for Robbins-Monro adaptation)
+     * @param i  Row index (i < j)
+     * @param j  Column index
+     * @return   Metropolis acceptance probability min(1, exp(ln_alpha)),
+     *           or 0.0 if the edge is inactive (caller masks it out).
      */
-    void update_edge_parameter(size_t i, size_t j, int iteration);
+    double update_edge_parameter(size_t i, size_t j);
 
     /**
      * Propose a new diagonal precision entry on the log scale.
      * Accepts or rejects with a Metropolis ratio using the Gaussian
      * likelihood, a Gamma(1,1) prior, and a Jacobian correction.
      *
-     * @param i          Diagonal index
-     * @param iteration  Current iteration (for Robbins-Monro adaptation)
+     * @param i  Diagonal index
+     * @return   Metropolis acceptance probability min(1, exp(ln_alpha)).
      */
-    void update_diagonal_parameter(size_t i, int iteration);
+    double update_diagonal_parameter(size_t i);
 
     /**
      * Metropolis-Hastings add-delete move for an edge indicator.