From b84588a685e667d7c3c5e8274f91763c2f5b373c Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 14:18:14 +0200 Subject: [PATCH 01/10] perf(ggm): defer chol(K) everywhere; log|K| by determinant-lemma increment The Metropolis, tuning, and edge-indicator sweeps read only Sigma and log|K| between accepts, so the per-accept Givens update/downdate pair (2 x O(p^2)) and the p-log get_log_det refresh are dropped. Accepted edge/diagonal moves advance log_det_precision_ by the same determinant-lemma ratio the MH ratio uses (O(1) from cached Sigma), and every sweep ends with one refresh_cholesky(), which rebuilds the factor and resets Sigma and the log-det exactly. This replaces the per-sweep drift-tolerance check: the SMW floating-point drift is now cleared unconditionally at sweep end. Verification: test-ggm-gibbs slow gates 20/20, test-hier-zratio-identity 14/14 (zero chol symmetry warnings), test-sample-ggm-prior 53/53, test-ggm-nuts fast 2/2. --- src/models/ggm/ggm_model.cpp | 103 ++++++++++++----------------------- src/models/ggm/ggm_model.h | 73 ++++++++++--------------- 2 files changed, 64 insertions(+), 112 deletions(-) diff --git a/src/models/ggm/ggm_model.cpp b/src/models/ggm/ggm_model.cpp index 385b66c2..8aae51e5 100644 --- a/src/models/ggm/ggm_model.cpp +++ b/src/models/ggm/ggm_model.cpp @@ -333,6 +333,15 @@ void GGMModel::cholesky_update_after_edge(double omega_ij_old, double omega_jj_o vf2_[i] = v2_[0]; vf2_[j] = v2_[1]; + // log|K| advances by the rank-2 matrix-determinant-lemma ratio -- the + // same quantity the MH ratio uses -- against the pre-update covariance. + // chol(K) itself is not maintained per accept; every sweep ends with a + // refresh_cholesky() that recomputes the factor and an exact log-det. + log_det_precision_ += cholesky_helpers::log_det_ratio_edge_kernel( + covariance_matrix_, i, j, + omega_ij_old, precision_proposal_(i, j), + omega_jj_old, precision_proposal_(j, j)); + const arma::uvec support = {static_cast(i), static_cast(j)}; apply_rank2_chol_smw_update_(support); @@ -345,36 +354,22 @@ void GGMModel::cholesky_update_after_edge(double omega_ij_old, double omega_jj_o } -void GGMModel::apply_rank2_chol_smw_update_(const arma::uvec& support, - bool update_L) +void GGMModel::apply_rank2_chol_smw_update_(const arma::uvec& support) { // K_new = K_old + vf1 vf2^T + vf2 vf1^T = K_old + u1 u1^T - u2 u2^T, - // where u1 = (vf1 + vf2) / sqrt(2), u2 = (vf1 - vf2) / sqrt(2). The - // change of basis diagonalises the symmetric rank-2 update so the chol - // factor advances via one rank-1 update + one rank-1 downdate. + // where u1 = (vf1 + vf2) / sqrt(2), u2 = (vf1 - vf2) / sqrt(2). // // `support` lists the nonzero indices of vf1/vf2 (hence of u1/u2): {i,j} // for the edge accept, {i} + N_i for a row-Gibbs row. Sigma * u then // touches only those columns -- O(p |support|) instead of the dense // O(p^2) gemv. + // + // chol(K) is not maintained here: nothing reads the factor between + // accepts within a sweep (the moves read Sigma and the incrementally + // maintained log-det), and every sweep ends with refresh_cholesky(). u1_ = (vf1_ + vf2_) / sqrt(2); u2_ = (vf1_ - vf2_) / sqrt(2); - // chol(K) update (2 x O(p^2)). Needed when a caller reads chol(K) or the - // log-det before the next full refresh (the edge accept does). The - // row-block Gibbs sweep passes update_L = false: the row draw reads only - // Sigma, and chol(K) is rebuilt once via refresh_cholesky() at sweep end. - // A failed downdate means K_new is not numerically PD along this route; - // rebuild all factors from K. - if (update_L) { - cholesky_update(cholesky_of_precision_, u1_); - if (!cholesky_downdate(cholesky_of_precision_, u2_)) { - refresh_cholesky(); - return; - } - log_det_precision_ = cholesky_helpers::get_log_det(cholesky_of_precision_); - } - // Sherman-Morrison-Woodbury rank-2 update of covariance_matrix_ = inv(K), // O(p^2) total: // K_new = K_old + M D M^T, M = [u1, u2], D = diag(+1, -1) @@ -633,10 +628,9 @@ void GGMModel::update_row_block_gibbs(size_t i) { gibbs_support_[0] = static_cast(i); for (size_t k = 0; k < q; ++k) gibbs_support_[k + 1] = Ni[k]; - // Defer the chol(K) Givens passes: nothing reads chol(K) between rows of - // the sweep. Sigma is still refreshed so the next row's Schur extraction - // is exact; chol(K) is rebuilt once in do_one_gibbs_step after the sweep. - apply_rank2_chol_smw_update_(gibbs_support_, /*update_L=*/false); + // Sigma is refreshed so the next row's Schur extraction is exact; + // chol(K) is rebuilt once in do_one_gibbs_step after the sweep. + apply_rank2_chol_smw_update_(gibbs_support_); vf1_[i] = 0.0; vf2_[i] = 0.0; @@ -650,11 +644,10 @@ void GGMModel::do_one_gibbs_step(int /*iteration*/) { for (size_t i = 0; i < p_; ++i) { update_row_block_gibbs(i); } - // chol(K) was deferred during the sweep -- the row draws read only Sigma, - // so the 2p per-row Givens passes were skipped. One O(p^3) factorisation - // here restores chol(K), inv chol(K), the log-det, and an exact Sigma - // (clearing the SMW drift the sweep accumulates), ready for the - // edge-indicator between-step that reads them next. + // The row draws read only Sigma. One O(p^3) factorisation here restores + // chol(K), inv chol(K), the log-det, and an exact Sigma (clearing the + // SMW drift the sweep accumulates), ready for the edge-indicator + // between-step that reads them next. refresh_cholesky(); refresh_cauchy_omega_(); } @@ -709,26 +702,11 @@ double GGMModel::update_diagonal_parameter(size_t i) { void GGMModel::cholesky_update_after_diag(double omega_ii_old, size_t i) { - - double delta = omega_ii_old - precision_proposal_(i, i); - - bool s = delta > 0; - vf1_(i) = std::sqrt(std::abs(delta)); - - bool ok = true; - if (s) - ok = cholesky_downdate(cholesky_of_precision_, vf1_); - else - cholesky_update(cholesky_of_precision_, vf1_); - - // reset for next iteration - vf1_(i) = 0.0; - - if (!ok) { - refresh_cholesky(); - return; - } - log_det_precision_ = cholesky_helpers::get_log_det(cholesky_of_precision_); + // log|K| advances by the rank-1 determinant-lemma ratio against the + // pre-update covariance; chol(K) is not maintained per accept (rebuilt + // once per sweep by refresh_cholesky(), which also resets the log-det). + log_det_precision_ += cholesky_helpers::log_det_ratio_diag_kernel( + covariance_matrix_, i, omega_ii_old, precision_proposal_(i, i)); // SMW rank-1 update of covariance_matrix_ = inv(K), O(p^2): // K_new = K_old + alpha e_i e_i^T, alpha = K_new(i,i) - K_old(i,i) @@ -930,9 +908,11 @@ void GGMModel::do_one_metropolis_step(int iteration) { metropolis_adapter_->update(index_mask, accept_prob, iteration); } - // Catch SMW-accumulated drift in covariance_matrix_ over a long chain. - // O(p^2); the refresh path is only taken when drift exceeds tolerance. - check_and_refresh_if_drift_(); + // chol(K) and the log-det were deferred during the sweep (the moves read + // Sigma and the incrementally maintained log-det only). One factorisation + // here restores the factor, an exact Sigma, and an exact log-det for the + // between-step and sample recording that follow. + refresh_cholesky(); } void GGMModel::init_metropolis_adaptation(const WarmupSchedule& schedule) { @@ -957,8 +937,8 @@ void GGMModel::update_edge_indicators() { update_edge_indicator_parameter_pair(i, j); } } - // SMW drift check; same rationale as the end-of-Metropolis-step path. - check_and_refresh_if_drift_(); + // Same rationale as the end-of-Metropolis-step refresh. + refresh_cholesky(); } void GGMModel::update_edge_indicator_conjugate(size_t i, size_t j) { @@ -1111,19 +1091,8 @@ void GGMModel::tune_proposal_sd(int iteration, const WarmupSchedule& schedule) { // Invalidate gradient cache after MH updates invalidate_gradient_cache(); - // SMW drift check; same rationale as the end-of-Metropolis-step path. - check_and_refresh_if_drift_(); -} - -void GGMModel::check_and_refresh_if_drift_() { - // diag(Sigma * K) should be ones. Max abs deviation in O(p^2) via the - // elementwise product: K is symmetric, so sum(Sigma.row(i) % K.row(i)) - // equals (Sigma * K)(i, i). - arma::vec d = arma::sum(covariance_matrix_ % precision_matrix_, 1) - 1.0; - double drift = arma::abs(d).max(); - if (!std::isfinite(drift) || drift > kCovDriftTol_) { - refresh_cholesky(); - } + // Same rationale as the end-of-Metropolis-step refresh. + refresh_cholesky(); } bool GGMModel::sigma_rows_consistent_(const arma::uvec& rows) const { diff --git a/src/models/ggm/ggm_model.h b/src/models/ggm/ggm_model.h index ecd8e4a7..7b10958f 100644 --- a/src/models/ggm/ggm_model.h +++ b/src/models/ggm/ggm_model.h @@ -797,10 +797,12 @@ class GGMModel : public BaseModel { /** - * Update the Cholesky factor after changing an off-diagonal element. + * Refresh Sigma and log|K| after changing an off-diagonal element. * - * Decomposes the rank-2 change into two rank-1 updates and - * recomputes the inverse Cholesky factor and covariance matrix. + * Advances log_det_precision_ by the rank-2 determinant-lemma ratio, + * encodes the accepted (i,j)/(j,j) change as the vf1/vf2 rank-2 pair, + * and delegates the covariance update to apply_rank2_chol_smw_update_. + * chol(K) is not maintained per accept (rebuilt once per sweep). * * @param omega_ij_old Previous value of omega(i,j) * @param omega_jj_old Previous value of omega(j,j) @@ -810,14 +812,13 @@ class GGMModel : public BaseModel { void cholesky_update_after_edge(double omega_ij_old, double omega_jj_old, size_t i, size_t j); /** - * Apply a symmetric rank-2 update to K, refresh chol(K), and refresh Sigma. + * Apply a symmetric rank-2 update to Sigma for a K update already written. * - * Given vf1_, vf2_ of length p, this carries out - * K_new = K_old + vf1 vf2^T + vf2 vf1^T - * chol(K) <- Givens update + downdate on u1 = (vf1+vf2)/sqrt2, u2 = (vf1-vf2)/sqrt2 - * Sigma <- Sherman-Morrison-Woodbury rank-2 update (O(p^2)), with a - * fallback to refresh_cholesky() when the downdate fails or - * the 2x2 capacitance is near-singular. + * Given vf1_, vf2_ of length p, this refreshes + * Sigma <- Sherman-Morrison-Woodbury rank-2 update (O(p^2)) for + * K_new = K_old + vf1 vf2^T + vf2 vf1^T, + * with a fallback to refresh_cholesky() when the 2x2 capacitance is + * near-singular. * * Inputs are taken from the model's vf1_, vf2_ scratch members so callers * can populate them in-place without an extra copy. The helper does not @@ -831,45 +832,29 @@ class GGMModel : public BaseModel { * touches only those columns, O(p |support|) instead of a dense O(p^2) * gemv, with a dense fallback when the support is near-full. * - * `update_L` controls whether chol(K) is advanced. The edge accept needs - * it true (the between-step reads chol(K)/log-det immediately). The - * row-block Gibbs sweep passes false: chol(K) is never read between rows - * (the row draw reads only Sigma), so the per-row Givens passes are - * skipped and chol(K) is rebuilt once via refresh_cholesky() at the end - * of the sweep. Sigma is always maintained so the next row's Schur - * extraction is exact. - * - * Sigma is maintained incrementally, so floating-point error accumulates - * across accepts; check_and_refresh_if_drift_() bounds it once per sweep. + * chol(K) and log|K| are NOT maintained here. The accept paths advance + * log_det_precision_ by the determinant-lemma ratio before calling this + * helper, nothing reads chol(K) between accepts within a sweep, and every + * sweep (Metropolis, tuning, edge-indicator, row-Gibbs) ends with + * refresh_cholesky(), which rebuilds the factor and resets Sigma and the + * log-det exactly -- bounding the SMW floating-point drift per sweep. */ - void apply_rank2_chol_smw_update_(const arma::uvec& support, - bool update_L = true); + void apply_rank2_chol_smw_update_(const arma::uvec& support); /** - * Update the Cholesky factor after changing a diagonal element. + * Refresh Sigma and log|K| after changing a diagonal element. * - * Applies a rank-1 Givens update to chol(K) and a Sherman-Morrison - * rank-1 update (O(p^2)) to the covariance matrix, with a fallback to - * refresh_cholesky() when the downdate fails or the scalar capacitance - * is near-singular. + * Advances log_det_precision_ by the rank-1 determinant-lemma ratio and + * applies a Sherman-Morrison rank-1 update (O(p^2)) to the covariance + * matrix, with a fallback to refresh_cholesky() when the scalar + * capacitance is near-singular. chol(K) is not maintained per accept + * (see apply_rank2_chol_smw_update_). * * @param omega_ii_old Previous value of omega(i,i) * @param i Diagonal index */ void cholesky_update_after_diag(double omega_ii_old, size_t i); - /** - * Refresh all factors if the SMW-maintained covariance has drifted. - * - * Computes max_i |diag(Sigma K) - 1| in O(p^2) and calls - * refresh_cholesky() when it exceeds kCovDriftTol_. Called once per - * Metropolis sweep, proposal-sd tuning sweep, and edge-indicator sweep. - */ - void check_and_refresh_if_drift_(); - - /** Tolerance on max_i |diag(Sigma K) - 1| before a full factor refresh. */ - static constexpr double kCovDriftTol_ = 1e-8; - /** * Check (Sigma K)(r, r) = 1 on the given rows. * @@ -892,12 +877,10 @@ class GGMModel : public BaseModel { static constexpr double kSigmaProbeTol_ = 1e-6; /** - * Recompute Cholesky and its inverse from the precision matrix. - * - * Used as a fallback when accumulated rank-1 updates/downdates - * cause numerical drift that makes the triangular inverse fail. - * Resets both cholesky_of_precision_ and inv_cholesky_of_precision_ - * from precision_matrix_, then recomputes covariance_matrix_. + * Recompute chol(K), its inverse, Sigma, and log|K| from the precision + * matrix. Called once at the end of every sweep (the accept paths + * maintain only Sigma and the log-det incrementally) and when a + * per-accept probe or SMW capacitance guard detects a drifted Sigma. */ void refresh_cholesky(); From 08a0172dea47522abe23ce4f729502ab3b42e00e Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 14:22:11 +0200 Subject: [PATCH 02/10] perf(platform): route remaining warm-path arma::exp through ARMA_MY_EXP The Z-ratio tilt-grid characteristic function and the SBM cluster draw were the last per-sweep arma::exp sites outside the platform macros. No change on macOS/Linux; Windows builds avoid UCRT exp here. Verification: zratio-engine and sbm-correction suites green. --- src/models/ggm/zratio_engine.cpp | 2 +- src/priors/sbm_edge_prior.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/models/ggm/zratio_engine.cpp b/src/models/ggm/zratio_engine.cpp index 932ad165..93d31174 100644 --- a/src/models/ggm/zratio_engine.cpp +++ b/src/models/ggm/zratio_engine.cpp @@ -428,7 +428,7 @@ bool ZRatioEngine::inner_reference_(const arma::mat& k_blk, const arma::uvec& si for (arma::uword k = 0; k < u.n_elem; ++k) { logphi += -0.5 * arma::log1p(u[k] * tg2); } - arma::vec phi = arma::exp(logphi); + arma::vec phi = ARMA_MY_EXP(logphi); fN = arma::accu(wt_ % ihat_ % phi); gG = arma::accu(wt_ % ghat_ % phi); return std::isfinite(w) && w > 0.0 && std::isfinite(fN) && diff --git a/src/priors/sbm_edge_prior.cpp b/src/priors/sbm_edge_prior.cpp index ec984e12..cf98a814 100644 --- a/src/priors/sbm_edge_prior.cpp +++ b/src/priors/sbm_edge_prior.cpp @@ -669,7 +669,7 @@ double corrected_log_marginal_mfm_sbm(const arma::uvec& cluster_assign, static arma::uword sample_cluster_log(const arma::vec& log_weights, SafeRNG& rng) { double mx = log_weights.max(); - arma::vec w = arma::exp(log_weights - mx); + arma::vec w = ARMA_MY_EXP(log_weights - mx); double u = runif(rng) * arma::accu(w); double cum = 0; for(arma::uword c = 0; c < w.n_elem; c++) { From 1803db7990f6da778e3c11665ec54551fa553adc Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 14:15:34 +0200 Subject: [PATCH 03/10] perf(bgmCompare): avoid per-eval copies and scratch matrices in the main component log_pseudoposterior_main_component bound the per-group count and Blume-Capel sufficient-statistic matrices by value on every evaluation; bind them by const reference, matching the other component functions. Replace the zero-filled V x max_categories and V x V scratch matrices, of which only one row and one column were used, with the group main-effect vector already computed and a length-V pairwise-weight vector hoisted out of the group loop. Verified: tests/testthat/test-bgmCompare.R passes (40 passed, 0 failed). --- .../bgmCompare/bgmCompare_logp_and_grad.cpp | 37 +++++++------------ 1 file changed, 14 insertions(+), 23 deletions(-) diff --git a/src/models/bgmCompare/bgmCompare_logp_and_grad.cpp b/src/models/bgmCompare/bgmCompare_logp_and_grad.cpp index 430d6287..5eabc900 100644 --- a/src/models/bgmCompare/bgmCompare_logp_and_grad.cpp +++ b/src/models/bgmCompare/bgmCompare_logp_and_grad.cpp @@ -835,16 +835,15 @@ double log_pseudoposterior_main_component( } const int num_variables = observations.n_cols; - const int max_num_categories = num_categories.max(); double log_pp = 0.0; + // group-specific pairwise weights of `variable` with all other variables + arma::vec pairwise_col(num_variables, arma::fill::zeros); + // --- per group --- for (int group = 0; group < num_groups; ++group) { - const arma::imat counts_per_category = counts_per_category_group[group]; - const arma::imat blume_capel_stats = blume_capel_stats_group[group]; - - arma::mat main_group(num_variables, max_num_categories, arma::fill::zeros); - arma::mat pairwise_group(num_variables, num_variables, arma::fill::zeros); + const arma::imat& counts_per_category = counts_per_category_group[group]; + const arma::imat& blume_capel_stats = blume_capel_stats_group[group]; const arma::vec proj_g = projection.row(group).t(); // length = num_groups-1 @@ -853,27 +852,21 @@ double log_pseudoposterior_main_component( variable, num_groups, main_effects, main_effect_indices, proj_g ); - // store into row v - main_group(variable, arma::span(0, me.n_elem - 1)) = me.t(); - - // upper triangle incl. base value; mirror to keep symmetry + // pairwise weights with the other variables; entry `variable` stays zero for (int u = 0; u < num_variables; u++) { if(u == variable) continue; - double w = compute_group_pairwise_effects( + pairwise_col(u) = compute_group_pairwise_effects( variable, u, num_groups, pairwise_effects, pairwise_effect_indices, inclusion_indicator, proj_g ); - pairwise_group(variable, u) = w; - pairwise_group(u, variable) = w; } // ---- data contribution pseudolikelihood (linear terms) ---- if (is_ordinal_variable(variable)) { - const double val = main_group(variable, category); - log_pp += static_cast(counts_per_category(category, variable)) * val; + log_pp += static_cast(counts_per_category(category, variable)) * + me(category); } else { - log_pp += static_cast(blume_capel_stats(par, variable)) * - main_group(variable, par); + log_pp += static_cast(blume_capel_stats(par, variable)) * me(par); } // ---- data contribution pseudolikelihood (quadratic terms) ---- @@ -882,7 +875,7 @@ double log_pseudoposterior_main_component( const arma::mat obs = arma::conv_to::from(observations.rows(r0, r1)); // ---- pseudolikelihood normalizing constants (per variable) ---- - const arma::vec rest_score = obs * pairwise_group.col(variable); + const arma::vec rest_score = obs * pairwise_col; const int num_cats = num_categories(variable); // bound to stabilize exp; clamp at 0 (the reference-category exponent) so @@ -892,14 +885,12 @@ double log_pseudoposterior_main_component( arma::vec denom(rest_score.n_elem, arma::fill::zeros); if (is_ordinal_variable(variable)) { - arma::vec main_eff = main_group.row(variable).cols(0, num_cats - 1).t(); denom = compute_denom_ordinal( - rest_score, main_eff, bound + rest_score, me, bound ); } else { - // linear/quadratic main effects from main_group - const double lin_effect = main_group(variable, 0); - const double quad_effect = main_group(variable, 1); + const double lin_effect = me(0); + const double quad_effect = me(1); const int ref = baseline_category(variable); denom = compute_denom_blume_capel( From 3a8c8eba592bcf68ed0a6e98698382c1fc847f7e Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 14:21:24 +0200 Subject: [PATCH 04/10] perf(bgmCompare): incremental sufficient-statistic updates in missing-data imputation impute_missing_bgmcompare rebuilt the V x V group pairwise-effect matrix per missing cell although the parameters are constant across the call; build it once per group before the imputation loop. Mutate the per-group count and Blume-Capel statistics in place instead of copying the matrix out and back per changed cell. Replace the full O(n_g * V^2) crossproduct rebuild of the group pairwise statistics with the O(V) rank-1 update of row and column `variable`, with the diagonal corrected by delta^2. Verified: tests/testthat/test-bgmCompare.R passes with NOT_CRAN=true (42 passed, 0 failed); a seeded bgmCompare fit with na_action = "impute" (adaptive-metropolis, 2 chains) produces raw_samples identical() to the pre-change build. --- src/models/bgmCompare/bgmCompare_sampler.cpp | 65 ++++++++++++-------- 1 file changed, 38 insertions(+), 27 deletions(-) diff --git a/src/models/bgmCompare/bgmCompare_sampler.cpp b/src/models/bgmCompare/bgmCompare_sampler.cpp index 3e7f7e78..3b4ca60e 100644 --- a/src/models/bgmCompare/bgmCompare_sampler.cpp +++ b/src/models/bgmCompare/bgmCompare_sampler.cpp @@ -58,9 +58,10 @@ // // Notes: // - The function updates both raw data and sufficient statistics in-place. -// - Group-specific pairwise effects are recomputed per missing entry. -// - For efficiency, you may consider incremental updates to `pairwise_stats` -// instead of full recomputation (`obs.t() * obs`) after each change. +// - Group-specific pairwise effects are built once per group; parameters +// are constant across the imputation pass. +// - `pairwise_stats` is updated incrementally: a changed cell only alters +// row and column `variable` of the group crossproduct. void impute_missing_bgmcompare( const arma::mat& main_effects, const arma::mat& pairwise_effects, @@ -89,6 +90,24 @@ void impute_missing_bgmcompare( double exponent, cumsum, u; int score, person, variable, new_value, old_value, group; + // Group-specific pairwise effect matrices; the parameters are constant + // across the imputation pass, so one build per group suffices. + std::vector group_pairwise_effects(num_groups); + for(int g = 0; g < num_groups; g++) { + const arma::vec proj_g = projection.row(g).t(); + group_pairwise_effects[g].zeros(num_variables, num_variables); + for(int v1 = 0; v1 < num_variables-1; v1++) { + for(int v2 = v1 + 1; v2 < num_variables; v2++) { + double w = compute_group_pairwise_effects( + v1, v2, num_groups, pairwise_effects, pairwise_effect_indices, + inclusion_indicator, proj_g + ); + group_pairwise_effects[g](v1, v2) = w; + group_pairwise_effects[g](v2, v1) = w; + } + } + } + //Impute missing data for(int missing = 0; missing < num_missings; missing++) { // Identify the observation to impute @@ -101,21 +120,8 @@ void impute_missing_bgmcompare( arma::vec group_main_effects = compute_group_main_effects( variable, num_groups, main_effects, main_effect_indices, proj_g); - // Generate a new observation based on the model - arma::mat group_pairwise_effects(num_variables, num_variables, arma::fill::zeros); - for(int v1 = 0; v1 < num_variables-1; v1++) { - for(int v2 = v1 + 1; v2 < num_variables; v2++) { - double w = compute_group_pairwise_effects( - v1, v2, num_groups, pairwise_effects, pairwise_effect_indices, - inclusion_indicator, proj_g - ); - group_pairwise_effects(v1, v2) = w; - group_pairwise_effects(v2, v1) = w; - } - } - double rest_score = - arma::as_scalar(observations.row(person) * group_pairwise_effects.col(variable)); + arma::as_scalar(observations.row(person) * group_pairwise_effects[group].col(variable)); if(is_ordinal_variable[variable] == true) { // For regular binary or ordinal variables @@ -159,27 +165,32 @@ void impute_missing_bgmcompare( // Update sufficient statistics for main effects if(is_ordinal_variable[variable] == true) { - arma::imat counts_per_category_group = counts_per_category[group]; + arma::imat& counts_per_category_group = counts_per_category[group]; if(old_value > 0) counts_per_category_group(old_value-1, variable)--; if(new_value > 0) counts_per_category_group(new_value-1, variable)++; - counts_per_category[group] = counts_per_category_group; } else { - arma::imat blume_capel_stats_group = blume_capel_stats[group]; + arma::imat& blume_capel_stats_group = blume_capel_stats[group]; blume_capel_stats_group(0, variable) -= old_value; blume_capel_stats_group(0, variable) += new_value; blume_capel_stats_group(1, variable) -= old_value * old_value; blume_capel_stats_group(1, variable) += new_value * new_value; - blume_capel_stats[group] = blume_capel_stats_group; } - // Update sufficient statistics for pairwise effects - const int r0 = group_indices(group, 0); - const int r1 = group_indices(group, 1); - arma::mat obs = arma::conv_to::from(observations.rows(r0, r1)); - arma::mat pairwise_stats_group = obs.t() * obs; // crossprod - pairwise_stats[group] = pairwise_stats_group; + // Update sufficient statistics for pairwise effects. In the group + // crossproduct X.t() * X only row and column `variable` depend on the + // changed cell: with delta = new_value - old_value and x the person's + // row (which already holds new_value), entry (u, variable) gains + // delta * x(u), symmetrically; the diagonal gain 2 * delta * new_value + // - delta^2 equals new_value^2 - old_value^2. + const double delta = static_cast(new_value - old_value); + const arma::rowvec obs_row = + arma::conv_to::from(observations.row(person)); + arma::mat& pairwise_stats_group = pairwise_stats[group]; + pairwise_stats_group.col(variable) += delta * obs_row.t(); + pairwise_stats_group.row(variable) += delta * obs_row; + pairwise_stats_group(variable, variable) -= delta * delta; } } return; From b612651ad956f02a35d73989d0d4a5931d5148d4 Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 14:24:46 +0200 Subject: [PATCH 05/10] perf(bgmCompare): reuse per-group double observation matrices in the pair component log_pseudoposterior_pair_component converted two integer observation columns to double on every evaluation. Both adaptive-Metropolis pairwise sweeps already hold per-group double observation matrices (obs_double_groups); pass those down and index the needed column directly instead of converting per call. Verified: tests/testthat/test-bgmCompare.R passes with NOT_CRAN=true (42 passed, 0 failed). --- .../bgmCompare/bgmCompare_logp_and_grad.cpp | 19 ++++++------------- .../bgmCompare/bgmCompare_logp_and_grad.h | 4 ++-- src/models/bgmCompare/bgmCompare_sampler.cpp | 4 ++-- 3 files changed, 10 insertions(+), 17 deletions(-) diff --git a/src/models/bgmCompare/bgmCompare_logp_and_grad.cpp b/src/models/bgmCompare/bgmCompare_logp_and_grad.cpp index 5eabc900..38b52ace 100644 --- a/src/models/bgmCompare/bgmCompare_logp_and_grad.cpp +++ b/src/models/bgmCompare/bgmCompare_logp_and_grad.cpp @@ -950,8 +950,7 @@ double log_pseudoposterior_main_component( // - main_effect_indices: Index ranges [row_start, row_end] for each variable. // - pairwise_effect_indices: Lookup table mapping (var1, var2) to row in pairwise_effects. // - projection: Group projection matrix (num_groups × (num_groups − 1)). -// - observations: Observation matrix (persons × variables). -// - group_indices: Row ranges [start, end] for each group in observations. +// - obs_double_groups: Per-group observation matrices converted to double. // - num_categories: Number of categories per variable. // - pairwise_stats_group: Per-group pairwise sufficient statistics. // - residual_matrices: Per-group residual matrices (persons × variables). @@ -978,8 +977,7 @@ double log_pseudoposterior_pair_component( const arma::imat& main_effect_indices, const arma::imat& pairwise_effect_indices, const arma::mat& projection, - const arma::imat& observations, - const arma::imat& group_indices, + const std::vector& obs_double_groups, const arma::ivec& num_categories, const std::vector& pairwise_stats_group, const std::vector& residual_matrices, @@ -998,7 +996,7 @@ double log_pseudoposterior_pair_component( return 0.0; } - const int num_variables = observations.n_cols; + const int num_variables = obs_double_groups[0].n_cols; const int max_num_categories = num_categories.max(); double log_pp = 0.0; int idx = pairwise_effect_indices(variable1, variable2); @@ -1033,19 +1031,14 @@ double log_pseudoposterior_pair_component( } // ---- pseudolikelihood normalizing constants (using residual matrix + delta) ---- - const int r0 = group_indices(group, 0); - const int r1 = group_indices(group, 1); - - // Pre-convert observation columns for the two variables (needed for delta adjustment) - const arma::vec obs_var1 = arma::conv_to::from(observations.col(variable1).rows(r0, r1)); - const arma::vec obs_var2 = arma::conv_to::from(observations.col(variable2).rows(r0, r1)); + const arma::mat& obs_g = obs_double_groups[group]; for (int v : {variable1, variable2}) { const int num_cats = num_categories(v); - const arma::vec& obs_other = (v == variable1) ? obs_var2 : obs_var1; + const int other = (v == variable1) ? variable2 : variable1; // Use residual_matrix with delta adjustment: O(n) instead of O(n*p) - arma::vec rest_score = residual_matrices[group].col(v) + obs_other * delta_g; + arma::vec rest_score = residual_matrices[group].col(v) + obs_g.col(other) * delta_g; // bound to stabilize exp; clamp at 0 so exp cannot overflow (the // Blume-Capel branch overwrites bound with its own max). diff --git a/src/models/bgmCompare/bgmCompare_logp_and_grad.h b/src/models/bgmCompare/bgmCompare_logp_and_grad.h index 44712018..ca00e246 100644 --- a/src/models/bgmCompare/bgmCompare_logp_and_grad.h +++ b/src/models/bgmCompare/bgmCompare_logp_and_grad.h @@ -211,6 +211,7 @@ double log_pseudoposterior_main_component( * Uses pre-computed residual matrices adjusted by delta to avoid full * recomputation. Used by element-wise Metropolis updates. * + * @param obs_double_groups Per-group observation matrices converted to double * @param residual_matrices Pre-computed residual matrices per group * @param variable1 First variable index * @param variable2 Second variable index @@ -224,8 +225,7 @@ double log_pseudoposterior_pair_component( const arma::imat& main_effect_indices, const arma::imat& pairwise_effect_indices, const arma::mat& projection, - const arma::imat& observations, - const arma::imat& group_indices, + const std::vector& obs_double_groups, const arma::ivec& num_categories, const std::vector& pairwise_stats_group, const std::vector& residual_matrices, diff --git a/src/models/bgmCompare/bgmCompare_sampler.cpp b/src/models/bgmCompare/bgmCompare_sampler.cpp index 3b4ca60e..f72b5220 100644 --- a/src/models/bgmCompare/bgmCompare_sampler.cpp +++ b/src/models/bgmCompare/bgmCompare_sampler.cpp @@ -446,7 +446,7 @@ void update_pairwise_effects_metropolis_bgmcompare ( double delta = theta - current; return log_pseudoposterior_pair_component( main_effects, pairwise_effects, main_effect_indices, - pairwise_effect_indices, projection, observations, group_indices, + pairwise_effect_indices, projection, obs_double_groups, num_categories, pairwise_stats, residual_matrices, num_groups, inclusion_indicator, is_ordinal_variable, baseline_category, var1, var2, h, delta, @@ -1002,7 +1002,7 @@ void tune_proposal_sd_bgmcompare( return log_pseudoposterior_pair_component( main_effects, pairwise_effects, main_effect_indices, pairwise_effect_indices, - projection, observations, group_indices, + projection, obs_double_groups, num_categories, pairwise_stats, residual_matrices, num_groups, inclusion_indicator, is_ordinal_variable, baseline_category, v1, v2, h, delta, From 47522f26289e80ecf461d561f43a90c476a0fe68 Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 14:33:50 +0200 Subject: [PATCH 06/10] perf(mixed): rank-2 quadratic-form shortcut for the AM GGM ratios log_ggm_ratio_edge/_diag evaluate the quadratic-form difference through the low-rank structure of the proposed covariance change (O(nq + q^2) dot products) instead of rebuilding the proposed conditional mean via an O(nq^2) GEMM and both n x q quadratic forms from scratch. The mean and cross moves (update_continuous_mean, update_pairwise_cross, update_edge_indicator_cross) evaluate the GGM term as the cached value plus a rank-1 quadratic-form delta, so conditional_mean_ is only touched on accept and the per-proposal n x q save/restore copies are gone. cond_mean_scratch_ is removed with its last consumers. Verification: fast mixed suites pass (test-mixed-mrf-likelihood 30, test-mixed-gradient 8, test-mixed-edge-indicators 3, test-mixed-theta-space 18). Same-seed 8+8 AM regression (iter 52, with/without selection): indicator paths bitwise identical; draws agree to 7.8e-16 over a short horizon and to 1.8e-7 over 104 iterations (Robbins-Monro feedback amplifies last-bit ln_alpha rounding differences; no acceptance decision flips). AM 8+8 selection fit (iter 1000, load_all -O0): 20.5 s -> 17.5 s. --- src/models/mixed/mixed_mrf_metropolis.cpp | 182 ++++++++++++++-------- src/models/mixed/mixed_mrf_model.cpp | 1 - src/models/mixed/mixed_mrf_model.h | 17 +- 3 files changed, 130 insertions(+), 70 deletions(-) diff --git a/src/models/mixed/mixed_mrf_metropolis.cpp b/src/models/mixed/mixed_mrf_metropolis.cpp index df8376ef..b6e22b82 100644 --- a/src/models/mixed/mixed_mrf_metropolis.cpp +++ b/src/models/mixed/mixed_mrf_metropolis.cpp @@ -60,8 +60,9 @@ double MixedMRFModel::update_main_effect(int s, int c, std::optional rm_ // update_continuous_mean // ============================================================================= // MH update for one continuous mean parameter main_effects_continuous_(j). -// The accept/reject uses log_conditional_ggm() + Normal(0, 1) prior. -// Must save/restore conditional_mean_ around the proposal. +// The accept/reject uses the cached GGM value plus a rank-1 quadratic-form +// delta, and a Normal(0, 1) prior. conditional_mean_ is only touched on +// accept (column-j shift by delta). // ============================================================================= double MixedMRFModel::update_continuous_mean(int j, std::optional rm_weight) { @@ -75,12 +76,22 @@ double MixedMRFModel::update_continuous_mean(int j, std::optional rm_wei // Proposed state: μ_j moves column j of the conditional mean and the // rest-score offsets 2 A_xy μ; everything else is unchanged. - matvec_col_j_scratch_ = conditional_mean_.col(j); + // + // GGM part: ΔM = delta · 1 e_j' with K unchanged, so + // quad_prop - quad_curr = -2 delta (1'D) K[:,j] + n delta² K_jj, + // read off the residual column sums (D = Y - conditional mean). + arma::rowvec resid_colsum = arma::sum(continuous_observations_, 0) + - arma::sum(conditional_mean_, 0); + double quad_delta = + -2.0 * delta * arma::dot(resid_colsum, + -2.0 * pairwise_effects_continuous_.col(j)) + + static_cast(n_) * delta * delta + * (-2.0 * pairwise_effects_continuous_(j, j)); + double ggm_prop = ll_ggm_cache_ - quad_delta / 2.0; + main_effects_continuous_(j) = proposed; - conditional_mean_.col(j) += delta; cross_bias_prop_ = cross_bias_ + (2.0 * delta) * pairwise_effects_cross_.col(j); - double ggm_prop = log_conditional_ggm(); for(size_t s = 0; s < p_; ++s) ll_marginal_prop_(s) = log_marginal_omrf_from( s, marginal_matvec_, marginal_interactions_(s, s), cross_bias_prop_(s)); @@ -92,8 +103,8 @@ double MixedMRFModel::update_continuous_mean(int j, std::optional rm_wei if(MY_LOG(runif(rng_)) >= ln_alpha) { main_effects_continuous_(j) = current_val; // reject - conditional_mean_.col(j) = matvec_col_j_scratch_; } else { + conditional_mean_.col(j) += delta; ll_ggm_cache_ = ggm_prop; ll_marginal_cache_ = ll_marginal_prop_; cross_bias_ = cross_bias_prop_; @@ -195,12 +206,15 @@ double MixedMRFModel::precision_constrained_diagonal(double x) const { // ============================================================================= // log_ggm_ratio_edge // ============================================================================= -// Log-likelihood ratio for a rank-2 off-diagonal precision change using the -// matrix determinant lemma for the log-det part and Woodbury for the -// quadratic-form part. Assumes precision_proposal_ is filled. -// -// TODO: replace the O(npq + nq²) quadratic-form computation with -// an O(nq) rank-2 shortcut. +// Log-likelihood ratio for a rank-2 off-diagonal precision change: matrix +// determinant lemma for the log-det part, Woodbury for the proposed +// covariance, and a rank-2 expansion of the quadratic-form difference. +// With ΔK = vf1 vf2' + vf2 vf1' and ΔΣ = -(w1 s2' + w2 s1'), the conditional +// mean moves by ΔM = a1 s2' + a2 s1' (a_k = -2 X A_xy w_k) and +// quad_prop - quad_curr = -2 tr(ΔM' D K) + tr(K ΔM' ΔM) + tr(ΔK D_p' D_p), +// D = Y - conditional mean, D_p = D - ΔM: O(nq + q²) dot products instead of +// the O(nq²) from-scratch quadratic forms. Assumes precision_proposal_ is +// filled. // ============================================================================= double MixedMRFModel::log_det_ratio_yy_edge(int i, int j) const { @@ -225,17 +239,14 @@ double MixedMRFModel::log_ggm_ratio_edge(int i, int j, arma::mat& cov_prop_out) size_t ui = static_cast(i); size_t uj = static_cast(j); - // Current precision (positive-definite) - arma::mat precision_curr = -2.0 * pairwise_effects_continuous_; - // --- Log-determinant ratio via matrix determinant lemma --- // ΔΩ has 3 nonzero entries: (i,j), (j,i), (j,j). // Ui = old - new off-diag, Uj = (old - new diag) / 2. The same Ui/Uj also // drive the Woodbury covariance update below, so they are kept here; the // log-det ratio itself is the canonical rank-2 det-lemma in // log_det_ratio_yy_edge (recomputes the identical Ui/Uj internally). - double Ui = precision_curr(ui, uj) - precision_proposal_(ui, uj); - double Uj = (precision_curr(uj, uj) - precision_proposal_(uj, uj)) / 2.0; + double Ui = -2.0 * pairwise_effects_continuous_(ui, uj) - precision_proposal_(ui, uj); + double Uj = (-2.0 * pairwise_effects_continuous_(uj, uj) - precision_proposal_(uj, uj)) / 2.0; double logdet_ratio = log_det_ratio_yy_edge(i, j); @@ -243,8 +254,10 @@ double MixedMRFModel::log_ggm_ratio_edge(int i, int j, arma::mat& cov_prop_out) // ΔΩ = vf1 vf2' + vf2 vf1' where vf1 = [0,...,-1,...] (j-th), // vf2 = [0,...,Ui,...,Uj,...] (i-th and j-th). // s1 = Σ vf1 = -Σ[:,j], s2 = Σ vf2 = Ui*Σ[:,i] + Uj*Σ[:,j] - arma::vec s1 = -covariance_continuous_.col(uj); - arma::vec s2 = Ui * covariance_continuous_.col(ui) + Uj * covariance_continuous_.col(uj); + cont_s1_ = -covariance_continuous_.col(uj); + cont_s2_ = Ui * covariance_continuous_.col(ui) + Uj * covariance_continuous_.col(uj); + const arma::vec& s1 = cont_s1_; + const arma::vec& s2 = cont_s2_; // 2×2 core matrix T = I + [vf2,vf1]' [s1,s2] // T = [1 + vf2's1, vf2's2; vf1's1, 1 + vf1's2] @@ -267,21 +280,44 @@ double MixedMRFModel::log_ggm_ratio_edge(int i, int j, arma::mat& cov_prop_out) arma::vec w2 = inv_t12 * s1 + inv_t22 * s2; // coefficient for s1' row arma::mat cov_prop = covariance_continuous_ - w1 * s2.t() - w2 * s1.t(); - // --- Proposed conditional mean --- - // M' = μ_y' + 2 X A_xy Σ', with X A_xy read from the sweep cache. - arma::mat cond_mean_prop = 2.0 * cross_matvec_ * cov_prop; - cond_mean_prop.each_row() += main_effects_continuous_.t(); - - // --- Quadratic form difference --- - arma::mat D_curr = continuous_observations_ - conditional_mean_; - arma::mat D_prop = continuous_observations_ - cond_mean_prop; - - double quad_curr = arma::accu((D_curr * precision_curr) % D_curr); - double quad_prop = arma::accu((D_prop * precision_proposal_) % D_prop); + // --- Quadratic-form difference through the rank-2 structure --- + // ΔΣ = -(w1 s2' + w2 s1'), so ΔM = 2 X A_xy ΔΣ = a1 s2' + a2 s1' with + // a_k = -2 (X A_xy) w_k read off the sweep cache. + cont_a1_ = -2.0 * (cross_matvec_ * w1); + cont_a2_ = -2.0 * (cross_matvec_ * w2); + resid_scratch_ = continuous_observations_ - conditional_mean_; + const arma::mat& D = resid_scratch_; + + // K s and D' a contractions (K = -2 A_yy) + arma::vec k1 = -2.0 * (pairwise_effects_continuous_ * s1); + arma::vec k2 = -2.0 * (pairwise_effects_continuous_ * s2); + arma::vec g1 = D.t() * cont_a1_; + arma::vec g2 = D.t() * cont_a2_; + + // -2 tr(ΔM' D K) + double quad_lin = -2.0 * (arma::dot(g1, k2) + arma::dot(g2, k1)); + + // tr(K ΔM' ΔM) + double a11 = arma::dot(cont_a1_, cont_a1_); + double a12 = arma::dot(cont_a1_, cont_a2_); + double a22 = arma::dot(cont_a2_, cont_a2_); + double quad_sq = a11 * arma::dot(s2, k2) + + a12 * (arma::dot(s1, k2) + arma::dot(s2, k1)) + + a22 * arma::dot(s1, k1); + + // tr(ΔΩ D_p' D_p) = 2 (D_p vf2)'(D_p vf1) with D_p = D - ΔM + double s2_vf2 = Ui * s2(ui) + Uj * s2(uj); + double s1_vf2 = Ui * s1(ui) + Uj * s1(uj); + arma::vec dp_vf1 = -D.col(uj) + s2(uj) * cont_a1_ + s1(uj) * cont_a2_; + arma::vec dp_vf2 = Ui * D.col(ui) + Uj * D.col(uj) + - s2_vf2 * cont_a1_ - s1_vf2 * cont_a2_; + double quad_dk = 2.0 * arma::dot(dp_vf2, dp_vf1); + + double quad_delta = quad_lin + quad_sq + quad_dk; double n = static_cast(n_); cov_prop_out = std::move(cov_prop); - return n / 2.0 * logdet_ratio - (quad_prop - quad_curr) / 2.0; + return n / 2.0 * logdet_ratio - quad_delta / 2.0; } @@ -289,7 +325,8 @@ double MixedMRFModel::log_ggm_ratio_edge(int i, int j, arma::mat& cov_prop_out) // log_ggm_ratio_diag // ============================================================================= // Log-likelihood ratio for a rank-1 diagonal precision change. -// Same structure as log_ggm_ratio_edge but simpler (Ui = 0). +// Same structure as log_ggm_ratio_edge but simpler (Ui = 0): the +// quadratic-form difference contracts through ΔΣ = c s s' in O(nq). // ============================================================================= double MixedMRFModel::log_ggm_ratio_diag(int i, arma::mat& cov_prop_out) const { @@ -308,26 +345,34 @@ double MixedMRFModel::log_ggm_ratio_diag(int i, arma::mat& cov_prop_out) const { // --- Proposed covariance via Sherman-Morrison (rank-1 special case) --- // ΔΩ = -2Uj * e_i e_i', so Σ' = Σ + 2Uj * Σ[:,i] Σ[i,:]' / (1 - 2Uj * Σ(i,i)) - arma::vec s = covariance_continuous_.col(ui); + cont_s1_ = covariance_continuous_.col(ui); + const arma::vec& s = cont_s1_; double denom = 1.0 - 2.0 * Uj * covariance_continuous_(ui, ui); - arma::mat cov_prop = covariance_continuous_ + (2.0 * Uj / denom) * s * s.t(); + double coef = 2.0 * Uj / denom; + arma::mat cov_prop = covariance_continuous_ + coef * s * s.t(); + + // --- Quadratic-form difference through the rank-1 structure --- + // ΔΣ = coef * s s', so ΔM = 2 X A_xy ΔΣ = a s' with a = 2 coef (X A_xy) s; + // quad_prop - quad_curr + // = -2 (D'a)·(K s) + (a·a) (s'K s) - 2Uj ||D_p[:,i]||², + // D = Y - conditional mean, D_p[:,i] = D[:,i] - s(i) a. + cont_a1_ = (2.0 * coef) * (cross_matvec_ * s); + resid_scratch_ = continuous_observations_ - conditional_mean_; + const arma::mat& D = resid_scratch_; - // --- Proposed conditional mean --- - arma::mat cond_mean_prop = 2.0 * cross_matvec_ * cov_prop; - cond_mean_prop.each_row() += main_effects_continuous_.t(); + arma::vec k = -2.0 * (pairwise_effects_continuous_ * s); + arma::vec g = D.t() * cont_a1_; - // --- Quadratic form difference --- - arma::mat D_curr = continuous_observations_ - conditional_mean_; - arma::mat D_prop = continuous_observations_ - cond_mean_prop; + double quad_lin = -2.0 * arma::dot(g, k); + double quad_sq = arma::dot(cont_a1_, cont_a1_) * arma::dot(s, k); + arma::vec dp_col = D.col(ui) - s(ui) * cont_a1_; + double quad_dk = -2.0 * Uj * arma::dot(dp_col, dp_col); - // Precision for quadratic form: only diagonal changed - arma::mat precision_curr = -2.0 * pairwise_effects_continuous_; - double quad_curr = arma::accu((D_curr * precision_curr) % D_curr); - double quad_prop = arma::accu((D_prop * precision_proposal_) % D_prop); + double quad_delta = quad_lin + quad_sq + quad_dk; double n = static_cast(n_); cov_prop_out = std::move(cov_prop); - return n / 2.0 * logdet_ratio - (quad_prop - quad_curr) / 2.0; + return n / 2.0 * logdet_ratio - quad_delta / 2.0; } @@ -566,8 +611,8 @@ double MixedMRFModel::update_pairwise_effects_continuous_diag(int i, std::option // update_pairwise_cross // ============================================================================= // MH update for one cross-type interaction pairwise_effects_cross_(i, j). -// Acceptance: sum_s log_marginal_omrf(s) + log_conditional_ggm() + Cauchy prior. -// Must save/restore conditional_mean_ and marginal_interactions_ around the proposal. +// Acceptance: sum_s log_marginal_omrf(s) + a rank-1 GGM quadratic-form delta +// + Cauchy prior. Parameters and conditional_mean_ are only touched on accept. // ============================================================================= double MixedMRFModel::update_pairwise_cross(int i, int j, std::optional rm_weight) { @@ -596,12 +641,18 @@ double MixedMRFModel::update_pairwise_cross(int i, int j, std::optional cross_bias_prop_ = cross_bias_; cross_bias_prop_(i) += 2.0 * delta * main_effects_continuous_(j); - cond_mean_scratch_ = conditional_mean_; - conditional_mean_ += (2.0 * delta) * discrete_observations_dbl_.col(i) - * covariance_continuous_.row(j); - pairwise_effects_cross_(i, j) = proposed; + // GGM part: the conditional-mean shift is ΔM_y = a v' with + // a = 2 delta x_i, v = Σ[:,j]; K is unchanged, so + // quad_prop - quad_curr = -2 (D'a)·(K v) + (a·a) (v'K v), + // D = Y - conditional mean. + arma::vec a = (2.0 * delta) * discrete_observations_dbl_.col(i); + arma::vec kv = -2.0 * (pairwise_effects_continuous_ * covariance_continuous_.col(j)); + resid_scratch_ = continuous_observations_ - conditional_mean_; + arma::vec g = resid_scratch_.t() * a; + double quad_delta = -2.0 * arma::dot(g, kv) + + arma::dot(a, a) * arma::dot(covariance_continuous_.col(j), kv); + double ggm_prop = ll_ggm_cache_ - quad_delta / 2.0; - double ggm_prop = log_conditional_ggm(); for(size_t s = 0; s < p_; ++s) ll_marginal_prop_(s) = log_marginal_omrf_from( s, marginal_matvec_prop_, mdiag_prop_(s), cross_bias_prop_(s)); @@ -611,10 +662,8 @@ double MixedMRFModel::update_pairwise_cross(int i, int j, std::optional double ln_alpha = ll_prop - ll_curr; - if(MY_LOG(runif(rng_)) >= ln_alpha) { - pairwise_effects_cross_(i, j) = current_val; // reject - std::swap(conditional_mean_, cond_mean_scratch_); - } else { + if(MY_LOG(runif(rng_)) < ln_alpha) { + pairwise_effects_cross_(i, j) = proposed; recompute_marginal_interactions(); recompute_am_caches(); } @@ -843,8 +892,9 @@ void MixedMRFModel::update_edge_indicator_cross(int i, int j) { } // --- Likelihood ratio --- - // Same rank-2 M update and rank-1 conditional-mean shift as - // update_pairwise_cross, with delta = k_prop - k_curr. + // Same rank-2 M update and rank-1 GGM quadratic-form delta as + // update_pairwise_cross, with delta = k_prop - k_curr. Parameters and + // conditional_mean_ are only touched on accept. double delta = k_prop - k_curr; double ll_curr = ll_ggm_cache_ + arma::accu(ll_marginal_cache_); @@ -861,22 +911,20 @@ void MixedMRFModel::update_edge_indicator_cross(int i, int j) { cross_bias_prop_ = cross_bias_; cross_bias_prop_(i) += 2.0 * delta * main_effects_continuous_(j); - cond_mean_scratch_ = conditional_mean_; - conditional_mean_ += (2.0 * delta) * discrete_observations_dbl_.col(i) - * covariance_continuous_.row(j); - pairwise_effects_cross_(i, j) = k_prop; + arma::vec a = (2.0 * delta) * discrete_observations_dbl_.col(i); + arma::vec kv = -2.0 * (pairwise_effects_continuous_ * covariance_continuous_.col(j)); + resid_scratch_ = continuous_observations_ - conditional_mean_; + arma::vec g = resid_scratch_.t() * a; + double quad_delta = -2.0 * arma::dot(g, kv) + + arma::dot(a, a) * arma::dot(covariance_continuous_.col(j), kv); + double ggm_prop = ll_ggm_cache_ - quad_delta / 2.0; - double ggm_prop = log_conditional_ggm(); for(size_t s = 0; s < p_; ++s) ll_marginal_prop_(s) = log_marginal_omrf_from( s, marginal_matvec_prop_, mdiag_prop_(s), cross_bias_prop_(s)); double ll_prop = ggm_prop + arma::accu(ll_marginal_prop_); - // Restore; the accept branch re-applies - pairwise_effects_cross_(i, j) = k_curr; - std::swap(conditional_mean_, cond_mean_scratch_); - double ln_alpha = ll_prop - ll_curr; if(g_prop == 1) { diff --git a/src/models/mixed/mixed_mrf_model.cpp b/src/models/mixed/mixed_mrf_model.cpp index bb7b01a8..9e19160f 100644 --- a/src/models/mixed/mixed_mrf_model.cpp +++ b/src/models/mixed/mixed_mrf_model.cpp @@ -361,7 +361,6 @@ void MixedMRFModel::recompute_am_caches() { cross_bias_prop_.set_size(p_); matvec_col_i_scratch_.set_size(n_); matvec_col_j_scratch_.set_size(n_); - cond_mean_scratch_.set_size(n_, q_); } for(size_t s = 0; s < p_; ++s) ll_marginal_cache_(s) = log_marginal_omrf_cached(s); diff --git a/src/models/mixed/mixed_mrf_model.h b/src/models/mixed/mixed_mrf_model.h index 721195fc..318c397c 100644 --- a/src/models/mixed/mixed_mrf_model.h +++ b/src/models/mixed/mixed_mrf_model.h @@ -450,7 +450,15 @@ class MixedMRFModel : public BaseModel { arma::vec cross_bias_prop_; ///< p proposed rest-score offsets arma::vec matvec_col_i_scratch_; ///< n saved matvec column (exact reject restore) arma::vec matvec_col_j_scratch_; ///< n saved matvec column (exact reject restore) - arma::mat cond_mean_scratch_; ///< n x q saved conditional mean (exact reject restore) + + // Low-rank GGM-ratio scratch. Filled by log_ggm_ratio_edge/_diag with the + // factors of the proposed conditional-mean change ΔM = a1 s2' + a2 s1' + // (rank 1: a1 s1'). Mutable because the ratio evaluations are const. + mutable arma::mat resid_scratch_; ///< n x q residual Y − conditional mean + mutable arma::vec cont_s1_; ///< q Σ-image factor of ΔΣ + mutable arma::vec cont_s2_; ///< q Σ-image factor of ΔΣ + mutable arma::vec cont_a1_; ///< n conditional-mean delta coefficient + mutable arma::vec cont_a2_; ///< n conditional-mean delta coefficient // Rank-1 Cholesky update workspace std::array cont_constants_{}; ///< Reparameterization constants @@ -609,7 +617,9 @@ class MixedMRFModel : public BaseModel { * Assumes precision_proposal_ is already filled by the caller. Writes the * proposed covariance Σ' (computed via Woodbury) to cov_prop_out so callers * can use it to recompute marginal_interactions_ for the OMRF likelihood - * ratio at the proposed Kyy. + * ratio at the proposed Kyy. The quadratic-form difference is evaluated + * through the rank-2 structure of ΔΣ in O(nq + q²); the factors of the + * conditional-mean change are left in cont_a1_/cont_a2_/cont_s1_/cont_s2_. */ double log_ggm_ratio_edge(int i, int j, arma::mat& cov_prop_out) const; @@ -617,6 +627,9 @@ class MixedMRFModel : public BaseModel { * Log-likelihood ratio for a proposed diagonal precision change (rank-1). * Assumes precision_proposal_ is already filled by the caller. Writes the * proposed covariance Σ' (computed via Sherman-Morrison) to cov_prop_out. + * The quadratic-form difference is evaluated through the rank-1 structure + * of ΔΣ in O(nq); the factors of the conditional-mean change are left in + * cont_a1_/cont_s1_. */ double log_ggm_ratio_diag(int i, arma::mat& cov_prop_out) const; From 23564e29afc8f74d4c3101818995fbf6e9590c51 Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 14:40:51 +0200 Subject: [PATCH 07/10] perf(mixed): adopt the proposal buffers on accepted AM Kyy/cross moves Accepted Kyy (off-diagonal, diagonal, edge-indicator) and cross (parameter, edge-indicator) moves adopt the caches the ratio evaluation already computed instead of rebuilding them via recompute_marginal_interactions + recompute_am_caches (an O(np^2) GEMM, all p OMRF denominators, and log_conditional_ggm per accept): - adopt_kyy_proposal_caches swaps marginal_matvec_/ll_marginal_cache_ with the proposal scratch, applies 2 A_xy dSigma A_xy' to marginal_interactions_/cross_term_ via the E = dSigma A_xy' buffer cached by omrf_ratio_for_covariance_change, shifts conditional_mean_ by the low-rank factors of the accepted ratio, and advances ll_ggm_cache_ by the ratio value. - adopt_cross_proposal_caches applies the rank-2 row/column-i update to marginal_interactions_/cross_term_, rank-1-updates cross_matvec_.col(j) and conditional_mean_, and copies the proposed bias and marginals. - The diagonal of marginal_interactions_ is overwritten with mdiag_prop_ so it matches the adopted per-variable marginals. - cholesky_update_after_precision_edge/_diag report whether the factors advanced incrementally; the rebuild fallback keeps the full recompute_am_caches, since the proposal buffers no longer match a from-scratch decomposition. - recompute_am_caches now refreshes marginal_interactions_ first, so the sweep-top full refresh (do_one_metropolis_step, sweep_within_model_mh, update_edge_indicators) bounds the floating-point drift of the low-rank cache updates to one sweep. Verification: fast mixed suites pass (30/8/3/18 assertions). Same-seed 8+8 AM regression (iter 52, with/without selection): indicator paths bitwise identical; draws agree to 6.5e-16 over a short horizon and to 2.4e-7 over 104 iterations (Robbins-Monro amplification of last-bit ln_alpha rounding; no acceptance decision flips). AM 8+8 selection fit (iter 1000, load_all -O0): 17.5 s -> 13.5 s (20.5 s at the audit base). --- src/models/mixed/mixed_mrf_likelihoods.cpp | 5 +- src/models/mixed/mixed_mrf_metropolis.cpp | 57 +++++++++++++--------- src/models/mixed/mixed_mrf_model.cpp | 57 ++++++++++++++++++++++ src/models/mixed/mixed_mrf_model.h | 45 ++++++++++++++--- 4 files changed, 134 insertions(+), 30 deletions(-) diff --git a/src/models/mixed/mixed_mrf_likelihoods.cpp b/src/models/mixed/mixed_mrf_likelihoods.cpp index ccc6d43c..0676d290 100644 --- a/src/models/mixed/mixed_mrf_likelihoods.cpp +++ b/src/models/mixed/mixed_mrf_likelihoods.cpp @@ -129,7 +129,10 @@ double MixedMRFModel::log_conditional_ggm() const { double MixedMRFModel::omrf_ratio_for_covariance_change(const arma::mat& cov_prop) { arma::mat delta_sigma = cov_prop - covariance_continuous_; - arma::mat E = delta_sigma * pairwise_effects_cross_.t(); // q x p + // E = ΔΣ A_xy' is kept in cross_delta_scratch_ for the accept path + // (adopt_kyy_proposal_caches applies 2 A_xy E to M and the cross term). + cross_delta_scratch_ = delta_sigma * pairwise_effects_cross_.t(); // q x p + const arma::mat& E = cross_delta_scratch_; marginal_matvec_prop_ = marginal_matvec_ + 2.0 * (cross_matvec_ * E); for(size_t s = 0; s < p_; ++s) { diff --git a/src/models/mixed/mixed_mrf_metropolis.cpp b/src/models/mixed/mixed_mrf_metropolis.cpp index b6e22b82..d98d84a3 100644 --- a/src/models/mixed/mixed_mrf_metropolis.cpp +++ b/src/models/mixed/mixed_mrf_metropolis.cpp @@ -384,7 +384,7 @@ double MixedMRFModel::log_ggm_ratio_diag(int i, arma::mat& cov_prop_out) const { // Then recomputes inv_cholesky_of_precision_ and covariance_continuous_. // ============================================================================= -void MixedMRFModel::cholesky_update_after_precision_edge( +bool MixedMRFModel::cholesky_update_after_precision_edge( double old_ij, double old_jj, int i, int j) { cont_v2_[0] = old_ij - precision_proposal_(i, j); @@ -406,8 +406,9 @@ void MixedMRFModel::cholesky_update_after_precision_edge( // decomposition from scratch (mirrors GGMModel's drift-guard). // pairwise_effects_continuous_ already holds the accepted value here, so // the rebuild reconstructs the accepted state. - if (down_ok && - arma::inv(inv_cholesky_of_precision_, arma::trimatu(cholesky_of_precision_))) { + bool incremental = down_ok && + arma::inv(inv_cholesky_of_precision_, arma::trimatu(cholesky_of_precision_)); + if (incremental) { covariance_continuous_ = inv_cholesky_of_precision_ * inv_cholesky_of_precision_.t(); log_det_precision_ = cholesky_helpers::get_log_det(cholesky_of_precision_); } else { @@ -418,6 +419,8 @@ void MixedMRFModel::cholesky_update_after_precision_edge( cont_vf1_[j] = 0.0; cont_vf2_[i] = 0.0; cont_vf2_[j] = 0.0; + + return incremental; } @@ -427,7 +430,7 @@ void MixedMRFModel::cholesky_update_after_precision_edge( // Rank-1 Cholesky update after accepting a diagonal precision change. // ============================================================================= -void MixedMRFModel::cholesky_update_after_precision_diag(double old_ii, int i) { +bool MixedMRFModel::cholesky_update_after_precision_diag(double old_ii, int i) { double delta = old_ii - precision_proposal_(i, i); bool downdate = delta > 0.0; @@ -442,8 +445,9 @@ void MixedMRFModel::cholesky_update_after_precision_diag(double old_ii, int i) { // Update the inverse Cholesky; fall back to a full rebuild if the // downdate lost positive definiteness or on drift (mirrors GGMModel's // drift-guard). - if (down_ok && - arma::inv(inv_cholesky_of_precision_, arma::trimatu(cholesky_of_precision_))) { + bool incremental = down_ok && + arma::inv(inv_cholesky_of_precision_, arma::trimatu(cholesky_of_precision_)); + if (incremental) { covariance_continuous_ = inv_cholesky_of_precision_ * inv_cholesky_of_precision_.t(); log_det_precision_ = cholesky_helpers::get_log_det(cholesky_of_precision_); } else { @@ -451,6 +455,8 @@ void MixedMRFModel::cholesky_update_after_precision_diag(double old_ii, int i) { } cont_vf1_[i] = 0.0; + + return incremental; } @@ -492,7 +498,8 @@ double MixedMRFModel::update_pairwise_effects_continuous_offdiag(int i, int j, s precision_proposal_(j, j) = theta_prop_jj; arma::mat cov_prop; - double ln_alpha = log_ggm_ratio_edge(i, j, cov_prop); + double ggm_ratio = log_ggm_ratio_edge(i, j, cov_prop); + double ln_alpha = ggm_ratio; // Determinant-tilt prior on the Kyy block: |Kyy|^delta contributes // delta * (log|Kyy_prop| - log|Kyy_curr|) @@ -529,9 +536,11 @@ double MixedMRFModel::update_pairwise_effects_continuous_offdiag(int i, int j, s pairwise_effects_continuous_(j, i) = -0.5 * theta_prop_ij; pairwise_effects_continuous_(j, j) = -0.5 * theta_prop_jj; - cholesky_update_after_precision_edge(old_theta_ij, old_theta_jj, i, j); - recompute_marginal_interactions(); - recompute_am_caches(); + if (cholesky_update_after_precision_edge(old_theta_ij, old_theta_jj, i, j)) { + adopt_kyy_proposal_caches(ggm_ratio, /*rank2=*/true); + } else { + recompute_am_caches(); + } } if (rm_weight) { @@ -569,7 +578,8 @@ double MixedMRFModel::update_pairwise_effects_continuous_diag(int i, std::option precision_proposal_(i, i) = theta_ii_prop; arma::mat cov_prop; - double ln_alpha = log_ggm_ratio_diag(i, cov_prop); + double ggm_ratio = log_ggm_ratio_diag(i, cov_prop); + double ln_alpha = ggm_ratio; // Determinant-tilt prior: rank-1 lemma, O(1) via the cached covariance. if (determinant_tilt_yy_ != 0.0) { @@ -594,9 +604,11 @@ double MixedMRFModel::update_pairwise_effects_continuous_diag(int i, std::option // Store: pairwise_effects_continuous_ = -1/2 * precision pairwise_effects_continuous_(i, i) = -0.5 * theta_ii_prop; - cholesky_update_after_precision_diag(old_theta_ii, i); - recompute_marginal_interactions(); - recompute_am_caches(); + if (cholesky_update_after_precision_diag(old_theta_ii, i)) { + adopt_kyy_proposal_caches(ggm_ratio, /*rank2=*/false); + } else { + recompute_am_caches(); + } } if (rm_weight) { @@ -664,8 +676,7 @@ double MixedMRFModel::update_pairwise_cross(int i, int j, std::optional if(MY_LOG(runif(rng_)) < ln_alpha) { pairwise_effects_cross_(i, j) = proposed; - recompute_marginal_interactions(); - recompute_am_caches(); + adopt_cross_proposal_caches(i, j, delta, u, ggm_prop); } if (rm_weight) { @@ -790,7 +801,8 @@ void MixedMRFModel::update_edge_indicator_continuous(int i, int j) { // --- Likelihood ratio --- arma::mat cov_prop; - double ln_alpha = log_ggm_ratio_edge(i, j, cov_prop); + double ggm_ratio = log_ggm_ratio_edge(i, j, cov_prop); + double ln_alpha = ggm_ratio; // Determinant-tilt prior: see update_pairwise_effects_continuous_offdiag. if (determinant_tilt_yy_ != 0.0) { @@ -862,9 +874,11 @@ void MixedMRFModel::update_edge_indicator_continuous(int i, int j) { set_gyy(i, j, g_prop); constraint_dirty_ = true; - cholesky_update_after_precision_edge(old_theta_ij, old_theta_jj, i, j); - recompute_marginal_interactions(); - recompute_am_caches(); + if (cholesky_update_after_precision_edge(old_theta_ij, old_theta_jj, i, j)) { + adopt_kyy_proposal_caches(ggm_ratio, /*rank2=*/true); + } else { + recompute_am_caches(); + } } } @@ -945,7 +959,6 @@ void MixedMRFModel::update_edge_indicator_cross(int i, int j) { pairwise_effects_cross_(i, j) = k_prop; set_gxy(i, j, g_prop); constraint_dirty_ = true; - recompute_marginal_interactions(); - recompute_am_caches(); + adopt_cross_proposal_caches(i, j, delta, u, ggm_prop); } } diff --git a/src/models/mixed/mixed_mrf_model.cpp b/src/models/mixed/mixed_mrf_model.cpp index 9e19160f..3ea2cba3 100644 --- a/src/models/mixed/mixed_mrf_model.cpp +++ b/src/models/mixed/mixed_mrf_model.cpp @@ -6,6 +6,7 @@ // gradient, and within-model MH update bodies live in the sibling translation units // (mixed_mrf_likelihoods.cpp, mixed_mrf_gradient.cpp, mixed_mrf_metropolis.cpp). #include +#include #include "models/mixed/mixed_mrf_model.h" #include "mcmc/execution/chain_result.h" #include "math/explog_macros.h" @@ -348,6 +349,12 @@ void MixedMRFModel::refresh_marginal_interactions_entry(int i, int j) { } void MixedMRFModel::recompute_am_caches() { + // Full refresh: the accept paths maintain marginal_interactions_, + // cross_term_, the matvecs, conditional_mean_, and the log-likelihood + // values by low-rank updates within a sweep; recomputing everything from + // the parameters here bounds the accumulated floating-point drift to one + // sweep. + recompute_marginal_interactions(); marginal_matvec_ = discrete_observations_dbl_ * marginal_interactions_; cross_matvec_ = discrete_observations_dbl_ * pairwise_effects_cross_; cross_bias_ = 2.0 * pairwise_effects_cross_ * main_effects_continuous_; @@ -373,6 +380,56 @@ void MixedMRFModel::recompute_conditional_mean_from_cross_matvec() { conditional_mean_.each_row() += main_effects_continuous_.t(); } +void MixedMRFModel::adopt_kyy_proposal_caches(double ggm_ratio, bool rank2) { + // The proposal scratch holds the accepted per-variable state. + std::swap(marginal_matvec_, marginal_matvec_prop_); + std::swap(ll_marginal_cache_, ll_marginal_prop_); + + // M and the cross term move by 2 A_xy ΔΣ A_xy' = 2 A_xy E, with + // E = ΔΣ A_xy' cached by omrf_ratio_for_covariance_change. The diagonal + // is overwritten with mdiag_prop_ so it matches the values the adopted + // per-variable marginals were computed with. + arma::mat m_delta = 2.0 * (pairwise_effects_cross_ * cross_delta_scratch_); + cross_term_ += m_delta; + marginal_interactions_ += m_delta; + marginal_interactions_.diag() = mdiag_prop_; + cross_term_.diag() = mdiag_prop_ - pairwise_effects_discrete_.diag(); + + // Conditional mean: ΔM = a1 s2' + a2 s1' (rank 1: a1 s1'), the same + // factors the accepted GGM ratio was evaluated with. + if(rank2) { + conditional_mean_ += cont_a1_ * cont_s2_.t() + cont_a2_ * cont_s1_.t(); + } else { + conditional_mean_ += cont_a1_ * cont_s1_.t(); + } + ll_ggm_cache_ += ggm_ratio; +} + +void MixedMRFModel::adopt_cross_proposal_caches( + int i, int j, double delta, const arma::vec& u, double ggm_prop) +{ + std::swap(marginal_matvec_, marginal_matvec_prop_); + std::swap(ll_marginal_cache_, ll_marginal_prop_); + cross_bias_ = cross_bias_prop_; + ll_ggm_cache_ = ggm_prop; + + // ΔM = 2δ (e_i u' + u e_i') + 2δ² Σ_jj e_i e_i' touches row and column i + // of M and the cross term; the (i, i) entry is overwritten with + // mdiag_prop_(i), matching the adopted per-variable marginals. + marginal_interactions_.row(i) += (2.0 * delta) * u.t(); + marginal_interactions_.col(i) += (2.0 * delta) * u; + marginal_interactions_(i, i) = mdiag_prop_(i); + cross_term_.row(i) += (2.0 * delta) * u.t(); + cross_term_.col(i) += (2.0 * delta) * u; + cross_term_(i, i) = mdiag_prop_(i) - pairwise_effects_discrete_(i, i); + + // X · A_xy moves in column j only; the conditional mean by the rank-1 + // shift the accepted GGM ratio was evaluated with. + cross_matvec_.col(j) += delta * discrete_observations_dbl_.col(i); + conditional_mean_ += (2.0 * delta) * discrete_observations_dbl_.col(i) + * covariance_continuous_.row(j); +} + // ============================================================================= // Constraint structure (RATTLE) diff --git a/src/models/mixed/mixed_mrf_model.h b/src/models/mixed/mixed_mrf_model.h index 318c397c..444a8b80 100644 --- a/src/models/mixed/mixed_mrf_model.h +++ b/src/models/mixed/mixed_mrf_model.h @@ -426,6 +426,9 @@ class MixedMRFModel : public BaseModel { arma::mat inv_cholesky_of_precision_; ///< q x q R^{-1} (upper triangular) arma::mat covariance_continuous_; ///< q x q Σ = Precision^{-1} double log_det_precision_; ///< log|Precision| + // marginal_interactions_, cross_term_, and conditional_mean_ are updated + // by low-rank formulas on the AM accept paths and refreshed in full at + // each sweep top (recompute_am_caches) and by set_vectorized_parameters. arma::mat marginal_interactions_; ///< p x p marginal PL interaction matrix arma::mat cross_term_; ///< p x p cached 2 A_xy Σ A_xy' (marginal PL cross term) arma::mat conditional_mean_; ///< n x q conditional mean @@ -448,6 +451,7 @@ class MixedMRFModel : public BaseModel { arma::vec mdiag_prop_; ///< p proposed diag(M') arma::vec ll_marginal_prop_; ///< p proposed per-variable marginals arma::vec cross_bias_prop_; ///< p proposed rest-score offsets + arma::mat cross_delta_scratch_; ///< q x p ΔΣ · A_xy' from the last covariance-change ratio arma::vec matvec_col_i_scratch_; ///< n saved matvec column (exact reject restore) arma::vec matvec_col_j_scratch_; ///< n saved matvec column (exact reject restore) @@ -541,9 +545,30 @@ class MixedMRFModel : public BaseModel { /** Refresh marginal_interactions_(i,j)/(j,i) from pairwise_effects_discrete_ and the cached cross_term_. Valid only while pairwise_effects_cross_ and covariance_continuous_ are unchanged since the last cross_term_ refresh. */ void refresh_marginal_interactions_entry(int i, int j); - /** Rebuild all AM sweep caches (matvecs, cross bias, per-variable marginals, GGM value). */ + /** Rebuild all AM sweep caches in full (marginal interactions, matvecs, + cross bias, per-variable marginals, GGM value). Called at the top of + each MH/indicator sweep, it resets the floating-point drift the + low-rank accept-path updates accumulate within a sweep. */ void recompute_am_caches(); + /** Adopt the proposal buffers as the current AM caches after an accepted + Kyy precision move: swaps marginal_matvec_/ll_marginal_cache_ with the + proposal scratch, applies the 2 A_xy ΔΣ A_xy' update to + marginal_interactions_/cross_term_ via cross_delta_scratch_, shifts + conditional_mean_ by the low-rank factors in cont_a*_/cont_s*_, and + advances ll_ggm_cache_ by ggm_ratio. rank2 selects the edge (rank-2) + or diagonal (rank-1) conditional-mean shift. */ + void adopt_kyy_proposal_caches(double ggm_ratio, bool rank2); + + /** Adopt the proposal buffers as the current AM caches after an accepted + A_xy(i, j) move of size delta: swaps the marginal scratch, applies the + rank-2 row/column-i update to marginal_interactions_/cross_term_ using + u = A_xy Σ[:,j] (evaluated at the pre-accept A_xy), rank-1-updates + cross_matvec_.col(j) and conditional_mean_, and sets ll_ggm_cache_ to + ggm_prop. */ + void adopt_cross_proposal_caches(int i, int j, double delta, + const arma::vec& u, double ggm_prop); + /** Recompute conditional_mean_ as 2 · cross_matvec_ · Σ + μ_y' (requires a fresh cross_matvec_). */ void recompute_conditional_mean_from_cross_matvec(); @@ -593,8 +618,9 @@ class MixedMRFModel : public BaseModel { double log_conditional_ggm() const; /** OMRF part of an MH ratio for a proposed covariance Σ'. Fills the - proposal scratch (matvec, diag, per-variable marginals) and returns - Σ_s L'_s − Σ_s L_s. Members are not mutated. */ + proposal scratch (matvec, diag, per-variable marginals) plus + cross_delta_scratch_ = ΔΣ A_xy' and returns Σ_s L'_s − Σ_s L_s. + The current-state caches are not mutated. */ double omrf_ratio_for_covariance_change(const arma::mat& cov_prop); // ========================================================================= @@ -650,11 +676,16 @@ class MixedMRFModel : public BaseModel { */ double log_det_ratio_yy_diag(int i) const; - /** Rank-1 Cholesky update after accepting an off-diagonal precision change. */ - void cholesky_update_after_precision_edge(double old_ij, double old_jj, int i, int j); + /** Rank-1 Cholesky update after accepting an off-diagonal precision + change. Returns true when the factors advanced incrementally; false + when the decomposition was rebuilt from scratch, in which case the + proposal buffers no longer match the state and the caller must rebuild + the AM caches in full. */ + bool cholesky_update_after_precision_edge(double old_ij, double old_jj, int i, int j); - /** Rank-1 Cholesky update after accepting a diagonal precision change. */ - void cholesky_update_after_precision_diag(double old_ii, int i); + /** Rank-1 Cholesky update after accepting a diagonal precision change. + Same return convention as cholesky_update_after_precision_edge. */ + bool cholesky_update_after_precision_diag(double old_ii, int i); // --- Parameter update sweeps --- From 7c1c2c7b7979aa850bd732d6bedc7e6aa626b35c Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 18:59:44 +0200 Subject: [PATCH 08/10] perf(ggm): engine-owned workspace for the NUTS forward map and backward pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit forward_map allocated its entire result per call — Phi, K, and per constrained column the Q/R QR matrices plus the A_q build and .t() temporaries — and theta_gradient_from_phi_bar allocated four q x q / q x m_q matrices per column, totalling O(p^3) doubles per leapfrog step on sparse graphs. The engine now owns the ForwardMapResult and the backward-pass scratch as mutable members reused across calls (zeros(), eye(), and copy assignment keep the existing allocation when dimensions match): forward_map returns a const reference valid until the next call, A_q^T is written directly into the in-place QR working matrix, and P / Phi_bar live in the workspace. Each chain clone owns its engine copy and the sampler is single-threaded within a chain. Verification: same-seed GGM NUTS + selection and mixed NUTS + selection draws are bitwise identical to the previous build (identical() on the full raw samples); ggm-gradient, ggm-nuts, mixed-gradient, mixed-theta-space, and ggm-gibbs suites green. --- src/ggm_gradient_interface.cpp | 2 +- src/models/ggm/ggm_gradient.cpp | 94 ++++++++++++++++++------- src/models/ggm/ggm_gradient.h | 47 ++++++++++++- src/models/ggm/ggm_model.cpp | 2 +- src/models/mixed/mixed_mrf_gradient.cpp | 2 +- src/models/mixed/mixed_mrf_model.cpp | 2 +- 6 files changed, 119 insertions(+), 30 deletions(-) diff --git a/src/ggm_gradient_interface.cpp b/src/ggm_gradient_interface.cpp index 193854a4..577d3dbd 100644 --- a/src/ggm_gradient_interface.cpp +++ b/src/ggm_gradient_interface.cpp @@ -58,7 +58,7 @@ Rcpp::List ggm_test_forward_map( GGMGradientEngine engine; engine.rebuild(cs, 100, dummy_S, ip, dp); - ForwardMapResult fm = engine.forward_map(theta); + const ForwardMapResult& fm = engine.forward_map(theta); return Rcpp::List::create( Rcpp::Named("Phi") = Rcpp::wrap(fm.Phi), diff --git a/src/models/ggm/ggm_gradient.cpp b/src/models/ggm/ggm_gradient.cpp index e3fb3bb6..4c431af4 100644 --- a/src/models/ggm/ggm_gradient.cpp +++ b/src/models/ggm/ggm_gradient.cpp @@ -60,6 +60,24 @@ void GGMGradientEngine::build_Aq( } } +void GGMGradientEngine::fill_Aq_t_( + const arma::mat& Phi, + const ColumnConstraints& col, + size_t q, + arma::mat& Aqt) +{ + // A_q^T (q x m_q) written directly into the QR working matrix: the same + // entries build_Aq produces, without materialising A_q and the .t() + // temporary per column. + Aqt.zeros(q, col.m_q); + for (size_t r = 0; r < col.m_q; ++r) { + size_t i = col.excluded_indices[r]; + for (size_t l = 0; l <= i; ++l) { + Aqt(l, r) = Phi(l, i); + } + } +} + // ===================================================================== // givens_qr // ===================================================================== @@ -83,10 +101,19 @@ void GGMGradientEngine::givens_qr( arma::vec& R_diag, std::vector& rots) { - size_t n = M.n_rows; - size_t m = M.n_cols; R = M; // working copy, will become R - Q.eye(n, n); // accumulate Q + givens_qr_inplace(Q, R, R_diag, rots); +} + +void GGMGradientEngine::givens_qr_inplace( + arma::mat& Q, + arma::mat& R, + arma::vec& R_diag, + std::vector& rots) +{ + size_t n = R.n_rows; + size_t m = R.n_cols; + Q.eye(n, n); // accumulate Q (reuses the allocation at same dims) rots.clear(); for (size_t j = 0; j < m; ++j) { @@ -133,8 +160,12 @@ void GGMGradientEngine::givens_qr( // forward_map // ===================================================================== -ForwardMapResult GGMGradientEngine::forward_map(const arma::vec& theta) const { - ForwardMapResult result; +const ForwardMapResult& GGMGradientEngine::forward_map(const arma::vec& theta) const { + // Fill the engine-owned workspace in place. zeros()/set_size/resize and + // the per-column assignments below reuse the existing allocations when + // the dimensions match, so steady-state calls (fixed graph between + // rebuilds) allocate nothing. + ForwardMapResult& result = fm_ws_; result.Phi.zeros(p_, p_); result.psi.set_size(p_); result.Nq.resize(p_); @@ -143,8 +174,6 @@ ForwardMapResult GGMGradientEngine::forward_map(const arma::vec& theta) const { result.Q_full.resize(p_); result.R_full.resize(p_); - arma::mat Aq_buf; // reusable buffer for A_q - for (size_t q = 0; q < p_; ++q) { const auto& col = structure_->columns[q]; size_t offset = structure_->theta_offsets[q]; @@ -186,19 +215,18 @@ ForwardMapResult GGMGradientEngine::forward_map(const arma::vec& theta) const { // Fully constrained: x_q = 0 // Givens QR for R_diag (Jacobian) and stored rotations if (m_q > 0) { - build_Aq(result.Phi, col, q, Aq_buf); - givens_qr(Aq_buf.t(), - result.Q_full[q], result.R_full[q], - result.R_diag[q], result.givens_rotations[q]); + fill_Aq_t_(result.Phi, col, q, result.R_full[q]); + givens_qr_inplace(result.Q_full[q], result.R_full[q], + result.R_diag[q], + result.givens_rotations[q]); } result.Nq[q].reset(); // x_q stays zero (already zeroed) } else { - // General case: build A_q, Givens QR, null space - build_Aq(result.Phi, col, q, Aq_buf); - givens_qr(Aq_buf.t(), - result.Q_full[q], result.R_full[q], - result.R_diag[q], result.givens_rotations[q]); + // General case: build A_q^T, Givens QR, null space + fill_Aq_t_(result.Phi, col, q, result.R_full[q]); + givens_qr_inplace(result.Q_full[q], result.R_full[q], + result.R_diag[q], result.givens_rotations[q]); // N_q = last d_q columns of Q result.Nq[q] = result.Q_full[q].cols(m_q, q - 1); @@ -258,7 +286,7 @@ std::pair GGMGradientEngine::logp_and_gradient( const arma::vec& theta) const { // --- Forward pass: theta -> Phi, K via null-space constraints --- - ForwardMapResult fm = forward_map(theta); + const ForwardMapResult& fm = forward_map(theta); const arma::mat& Phi = fm.Phi; const arma::mat& K = fm.K; const arma::mat& S = *suf_stat_; @@ -272,10 +300,12 @@ std::pair GGMGradientEngine::logp_and_gradient( arma::vec(theta.n_elem, arma::fill::zeros)}; } - // P = Phi * S — reused for value and gradient. + // P = Phi * S — reused for value and gradient; written into the + // engine workspace so steady-state calls reuse the allocation. // Phi is upper triangular; trimatu dispatches to BLAS dtrmm, // halving the FLOP count vs dense gemm. - arma::mat P = arma::trimatu(Phi) * S; + arma::mat& P = P_ws_; + P = arma::trimatu(Phi) * S; // --- Log-posterior value --- double log_det_K = 2.0 * arma::accu(fm.psi); @@ -317,7 +347,8 @@ std::pair GGMGradientEngine::logp_and_gradient( // Diagonal prior: d/dPhi [log p(K_ii/2)] // = (1/2) * grad(K_ii/2) * dK_ii/dPhi // = (1/2) * grad(K_ii/2) * 2 Phi(:,i) = grad(K_ii/2) * Phi(:,i). - arma::mat Phi_bar = -P; + arma::mat& Phi_bar = Phi_bar_ws_; + Phi_bar = -P; for (size_t i = 0; i < p_; ++i) { double dg = diagonal_prior_->grad(0.5 * K(i, i)); Phi_bar.col(i).head(i + 1) += dg * Phi.col(i).head(i + 1); @@ -409,9 +440,20 @@ void GGMGradientEngine::theta_gradient_from_phi_bar( const auto& rotations = fm.givens_rotations[q]; size_t n_rot = rotations.size(); - // Initialize W_bar (R_bar) and Q_bar with seed adjoints. - arma::mat W_bar(q, m_q, arma::fill::zeros); - arma::mat Q_bar(q, q, arma::fill::zeros); + // Initialize W_bar (R_bar) and Q_bar with seed adjoints. All four + // per-column matrices live in the engine workspace: zeros()/copy + // assignment reuse the existing allocations at steady state, so the + // backward pass allocates nothing per leapfrog step. + if (Wbar_ws_.size() < p_) { + Wbar_ws_.resize(p_); + Qbar_ws_.resize(p_); + Qwork_ws_.resize(p_); + Wwork_ws_.resize(p_); + } + arma::mat& W_bar = Wbar_ws_[q]; + arma::mat& Q_bar = Qbar_ws_[q]; + W_bar.zeros(q, m_q); + Q_bar.zeros(q, q); size_t rank = std::min(m_q, q); const arma::mat& R = fm.R_full[q]; @@ -428,8 +470,10 @@ void GGMGradientEngine::theta_gradient_from_phi_bar( } } - arma::mat Q_work = fm.Q_full[q]; - arma::mat W_work = fm.R_full[q]; + arma::mat& Q_work = Qwork_ws_[q]; + arma::mat& W_work = Wwork_ws_[q]; + Q_work = fm.Q_full[q]; + W_work = fm.R_full[q]; for (size_t k = n_rot; k-- > 0; ) { const auto& rot = rotations[k]; diff --git a/src/models/ggm/ggm_gradient.h b/src/models/ggm/ggm_gradient.h index bde28987..417fee9e 100644 --- a/src/models/ggm/ggm_gradient.h +++ b/src/models/ggm/ggm_gradient.h @@ -100,10 +100,17 @@ class GGMGradientEngine { * computes Givens QR of A_q^T for the null-space basis N_q, sets * x_q = N_q f_q, and accumulates the Jacobian. * + * Returns a reference to the engine-owned workspace, valid until the + * next forward_map() call on this engine. Callers bind it by const + * reference; each chain clone owns its engine copy and the sampler is + * single-threaded within a chain, so reuse across leapfrog steps is + * safe and avoids reallocating the per-column QR matrices + * (O(p^3) doubles per call on sparse graphs). + * * @param theta Parameter vector of length p + |E| * @return ForwardMapResult with Phi, K, log|det J|, and cached Givens data */ - ForwardMapResult forward_map(const arma::vec& theta) const; + const ForwardMapResult& forward_map(const arma::vec& theta) const; /** * Combined log-posterior and gradient evaluation. @@ -162,12 +169,34 @@ class GGMGradientEngine { arma::vec& R_diag, std::vector& rots); + /** + * In-place variant of givens_qr: R holds M on entry and is factored in + * place; Q, R_diag, and rots are reset and filled. Existing allocations + * in Q/R_diag are reused when the dimensions match, so per-leapfrog + * callers do not reallocate the per-column QR matrices. + */ + static void givens_qr_inplace( + arma::mat& Q, + arma::mat& R, + arma::vec& R_diag, + std::vector& rots); + static void build_Aq(const arma::mat& Phi, const ColumnConstraints& col, size_t q, arma::mat& Aq); private: + /** + * Fill Aqt with A_q^T (q x m_q) directly — the entries build_Aq produces, + * without materialising A_q and its .t() temporary. Aqt doubles as the + * in-place QR working matrix. + */ + static void fill_Aq_t_(const arma::mat& Phi, + const ColumnConstraints& col, + size_t q, + arma::mat& Aqt); + const GraphConstraintStructure* structure_ = nullptr; size_t n_ = 0; size_t p_ = 0; @@ -177,4 +206,20 @@ class GGMGradientEngine { // Determinant-tilt exponent: adds delta_ * log|K| to the (unnormalised) // log-prior. delta_ = 0 recovers the untilted target. double delta_ = 0.0; + + // Per-call workspace. Mutable because forward_map/logp_and_gradient are + // logically const; each chain clone owns its own engine copy and the + // sampler is single-threaded within a chain, so reuse across calls is + // safe. fm_ws_ is the object forward_map() returns a reference to; the + // vectors hold the backward pass's per-column scratch (two zero-filled + // and two copied q x q / q x m_q matrices per constrained column), all + // reused via zeros()/copy-assignment, which keep the existing + // allocation when the dimensions match. + mutable ForwardMapResult fm_ws_; + mutable arma::mat P_ws_; + mutable arma::mat Phi_bar_ws_; + mutable std::vector Wbar_ws_; + mutable std::vector Qbar_ws_; + mutable std::vector Qwork_ws_; + mutable std::vector Wwork_ws_; }; diff --git a/src/models/ggm/ggm_model.cpp b/src/models/ggm/ggm_model.cpp index 8aae51e5..9ddff208 100644 --- a/src/models/ggm/ggm_model.cpp +++ b/src/models/ggm/ggm_model.cpp @@ -133,7 +133,7 @@ void GGMModel::set_vectorized_parameters(const arma::vec& parameters) { ensure_constraint_structure(); // Run forward map: theta -> Phi -> K - ForwardMapResult fm = gradient_engine_.forward_map(parameters); + const ForwardMapResult& fm = gradient_engine_.forward_map(parameters); // Update internal state precision_matrix_ = fm.K; diff --git a/src/models/mixed/mixed_mrf_gradient.cpp b/src/models/mixed/mixed_mrf_gradient.cpp index 9b14906b..f84ff000 100644 --- a/src/models/mixed/mixed_mrf_gradient.cpp +++ b/src/models/mixed/mixed_mrf_gradient.cpp @@ -202,7 +202,7 @@ std::pair MixedMRFModel::logp_and_gradient( // which enforces excluded-edge zeros through the null-space bases --- size_t chol_offset = static_cast(chol_grad_offset_); size_t chol_dim = chol_constraint_structure_.active_dim; - ForwardMapResult fm = yy_engine_.forward_map( + const ForwardMapResult& fm = yy_engine_.forward_map( arma::vec(parameters.subvec(chol_offset, chol_offset + chol_dim - 1))); const arma::mat& temp_cholesky = fm.Phi; diff --git a/src/models/mixed/mixed_mrf_model.cpp b/src/models/mixed/mixed_mrf_model.cpp index 3ea2cba3..bcaeba4b 100644 --- a/src/models/mixed/mixed_mrf_model.cpp +++ b/src/models/mixed/mixed_mrf_model.cpp @@ -746,7 +746,7 @@ void MixedMRFModel::set_vectorized_parameters(const arma::vec& params) { // excluded-edge zeros enforced by the null-space parameterization size_t chol_dim = chol_constraint_structure_.active_dim; arma::vec theta_yy = params.subvec(idx, idx + chol_dim - 1); - ForwardMapResult fm = yy_engine_.forward_map(theta_yy); + const ForwardMapResult& fm = yy_engine_.forward_map(theta_yy); cholesky_of_precision_ = fm.Phi; pairwise_effects_continuous_ = -0.5 * fm.K; bool ok = arma::solve(inv_cholesky_of_precision_, From e10364f6f73341101674bff37170af6d2a11ee56 Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 19:05:34 +0200 Subject: [PATCH 09/10] perf(gradient): batch per-variable pairwise GEMVs into per-call GEMMs The OMRF gradient issued one X^T * E GEMV per variable and the mixed gradient two (pw_grad and diff_pw), plus an O(p*q) scalar loop for the cross-contribution. Collect the per-variable expected-score vectors into an n x p buffer and assemble the pairwise, cross, and Theta-bar contributions after the loop from batched GEMMs: one X^T * E_all in the OMRF, two in the mixed gradient, with the cross self- and cross-contributions combined as 4 (diff_pw_all + diff_pw_all^T) * (A_xy Sigma_yy) and the mean gradient as one GEMV. Verification: fixed-seed ordinal NUTS draws (p=10, n=300) are bitwise identical; mixed 6+6 NUTS draws deviate by max 3.1e-10 over a 10+10-iteration horizon (reassociated BLAS accumulation). FD gradient suites pass: test-mixed-gradient.R (8), test-mixed-theta-space.R (18), test-ggm-gradient.R (14). --- src/models/mixed/mixed_mrf_gradient.cpp | 179 +++++++++--------------- src/models/omrf/omrf_model.cpp | 38 ++--- 2 files changed, 87 insertions(+), 130 deletions(-) diff --git a/src/models/mixed/mixed_mrf_gradient.cpp b/src/models/mixed/mixed_mrf_gradient.cpp index f84ff000..aac3a354 100644 --- a/src/models/mixed/mixed_mrf_gradient.cpp +++ b/src/models/mixed/mixed_mrf_gradient.cpp @@ -252,11 +252,17 @@ std::pair MixedMRFModel::logp_and_gradient( arma::mat cross_times_cov; // p x q arma::mat Theta_bar; // p x p marginal-PL coupling for precision gradient cross_times_cov = temp_pairwise_cross * temp_covariance; - Theta_bar = arma::zeros(p_, p_); // ========================================================================= // Part 1: OMRF conditionals // ========================================================================= + // The loop collects per-variable expected scores and residual scalars; + // the pairwise, cross, and Θ̄ contributions are assembled after the loop + // from batched GEMMs instead of per-variable GEMVs. + + arma::mat E_all(n_, p_, arma::fill::none); // expected score per obs + arma::vec diff_diag_all(p_, arma::fill::none); // dot(x_s, x_s) - sum(E_sq) + arma::vec sum_obs_minus_E_all(p_, arma::fill::none); int main_effects_discrete_offset = 0; for(size_t s = 0; s < p_; ++s) { @@ -301,78 +307,17 @@ std::pair MixedMRFModel::logp_and_gradient( // Expected value E_s[c+1|rest] per observation arma::vec weights = arma::regspace(1, C_s); - arma::vec E = logz_out_.probs.cols(1, C_s) * weights; - - // Pairwise discrete gradient: sum_i x_{i,t} * (x_{i,s}+1 - E_s) - // (uses pre-transposed discrete observations for BLAS efficiency) - // Factor 2: chain rule d/dK = 2 × d/dσ - arma::vec pw_grad = discrete_observations_dbl_t_ * E; - for(size_t t = 0; t < p_; ++t) { - if(edge_indicators_(s, t) == 0 || s == t) continue; - int loc = (s < t) ? disc_index_cache_(s, t) : disc_index_cache_(t, s); - grad(loc) -= 2.0 * pw_grad(t); - } - - // Additional pairwise_discrete gradient from Θ_ss in denominator: - // ∂/∂pairwise_effects_discrete_{st} through Θ_ss: zero (∂Θ_ss/∂pairwise_effects_discrete_st = δ_{st}) - // So pairwise_discrete gradient from Θ rest scores is already handled above. - - // Pairwise_cross gradient from marginal OMRF (through Θ): - // ∂marginal_{st}/∂pairwise_effects_cross_{a,j} has two terms: - // = 2 [Σyy pairwise_effects_cross_t']_j δ_{as} + 2 [pairwise_effects_cross_s Σyy]_j δ_{at} - // Self-contribution (a=s): from rest_s → pairwise_effects_cross_s - // Cross-contribution (a=t): from rest_s → pairwise_effects_cross_t for each t≠s + E_all.col(s) = logz_out_.probs.cols(1, C_s) * weights; arma::vec weights_sq = arma::square(weights); arma::vec E_sq = logz_out_.probs.cols(1, C_s) * weights_sq; - arma::vec diff_pw = discrete_observations_dbl_t_ * - (discrete_observations_dbl_.col(s) - E); - diff_pw(s) = 0.0; - - double diff_diag = arma::dot( + diff_diag_all(s) = arma::dot( discrete_observations_dbl_.col(s), discrete_observations_dbl_.col(s)) - arma::accu(E_sq); - double sum_obs_minus_E = arma::accu(discrete_observations_dbl_.col(s)) - arma::accu(E); - - // Accumulate Θ̄ for precision gradient coupling - for(size_t t = 0; t < p_; ++t) { - if(t != s) Theta_bar(s, t) += 2.0 * diff_pw(t); - } - Theta_bar(s, s) += diff_diag; - - // Self-contribution: a = s - // Off-diagonal effective interaction: ∂Θ_{st}/∂pairwise_effects_cross_{s,j} = 2 [Σyy pairwise_effects_cross_t']_j - // Diagonal effective interaction: ∂Θ_{ss}/∂pairwise_effects_cross_{s,j} = 4 [Σyy pairwise_effects_cross_s']_j - // Rest-score bias: ∂(2 pairwise_effects_cross_s μy)/∂pairwise_effects_cross_{s,j} = 2 μy_j - arma::rowvec cross_self = 4.0 * (diff_pw.t() * temp_pairwise_cross) * temp_covariance - + 4.0 * diff_diag * cross_times_cov.row(s) - + 2.0 * sum_obs_minus_E * temp_main_continuous.t(); - - for(size_t j = 0; j < q_; ++j) { - if(edge_indicators_(s, p_ + j) == 0) continue; - int loc = cross_index_cache_(s, j); - grad(loc) += cross_self(j); - } - - // Cross-contribution: a = t, for each t ≠ s - // ∂l_s/∂pairwise_effects_cross_{t,:} = diff_pw(t) * 2 * pairwise_effects_cross_s * Σyy - arma::rowvec V_s = 4.0 * cross_times_cov.row(s); - for(size_t t = 0; t < p_; ++t) { - if(t == s || std::abs(diff_pw(t)) < 1e-300) continue; - for(size_t j = 0; j < q_; ++j) { - if(edge_indicators_(t, p_ + j) == 0) continue; - int loc = cross_index_cache_(t, j); - grad(loc) += diff_pw(t) * V_s(j); - } - } - - // Continuous mean gradient from marginal OMRF: - // ∂l_s/∂main_effects_continuous_j = 2 pairwise_effects_cross_{sj} * sum_i (x_{is} - E_s) - for(size_t j = 0; j < q_; ++j) { - grad(main_effects_continuous_grad_offset_ + j) += 2.0 * temp_pairwise_cross(s, j) * sum_obs_minus_E; - } + sum_obs_minus_E_all(s) = arma::accu(discrete_observations_dbl_.col(s)) + - arma::accu(E_all.col(s)); main_effects_discrete_offset += C_s; } else { @@ -401,67 +346,73 @@ std::pair MixedMRFModel::logp_and_gradient( grad(main_effects_discrete_offset + 1) -= arma::accu(logz_out_.probs * sq_score); // Expected score per person - arma::vec E = logz_out_.probs * score; - - // Pairwise discrete gradient - // Factor 2: chain rule d/dK = 2 × d/dσ - arma::vec pw_grad = discrete_observations_dbl_t_ * E; - for(size_t t = 0; t < p_; ++t) { - if(edge_indicators_(s, t) == 0 || s == t) continue; - int loc = (s < t) ? disc_index_cache_(s, t) : disc_index_cache_(t, s); - grad(loc) -= 2.0 * pw_grad(t); - } + E_all.col(s) = logz_out_.probs * score; - // Pairwise_cross gradient from marginal OMRF (same structure as ordinal) arma::vec E_sq = logz_out_.probs * sq_score; - arma::vec diff_pw = discrete_observations_dbl_t_ * - (discrete_observations_dbl_.col(s) - E); - diff_pw(s) = 0.0; - - double diff_diag = arma::dot( + diff_diag_all(s) = arma::dot( discrete_observations_dbl_.col(s), discrete_observations_dbl_.col(s)) - arma::accu(E_sq); - double sum_obs_minus_E = arma::accu(discrete_observations_dbl_.col(s)) - arma::accu(E); - - // Accumulate Θ̄ for precision gradient coupling - for(size_t t = 0; t < p_; ++t) { - if(t != s) Theta_bar(s, t) += 2.0 * diff_pw(t); - } - Theta_bar(s, s) += diff_diag; - - // Self-contribution: a = s - arma::rowvec cross_self = 4.0 * (diff_pw.t() * temp_pairwise_cross) * temp_covariance - + 4.0 * diff_diag * cross_times_cov.row(s) - + 2.0 * sum_obs_minus_E * temp_main_continuous.t(); + sum_obs_minus_E_all(s) = arma::accu(discrete_observations_dbl_.col(s)) + - arma::accu(E_all.col(s)); - for(size_t j = 0; j < q_; ++j) { - if(edge_indicators_(s, p_ + j) == 0) continue; - int loc = cross_index_cache_(s, j); - grad(loc) += cross_self(j); - } + main_effects_discrete_offset += 2; + } + } - // Cross-contribution: a = t, for each t ≠ s - arma::rowvec V_s = 4.0 * cross_times_cov.row(s); - for(size_t t = 0; t < p_; ++t) { - if(t == s || std::abs(diff_pw(t)) < 1e-300) continue; - for(size_t j = 0; j < q_; ++j) { - if(edge_indicators_(t, p_ + j) == 0) continue; - int loc = cross_index_cache_(t, j); - grad(loc) += diff_pw(t) * V_s(j); - } - } + // --- Batched pairwise/cross gradient assembly --- + // Two X^T * (...) GEMMs replace the 2p per-variable GEMVs: + // pw_grad_all.col(s) = X^T E_s + // diff_pw_all.col(s) = X^T (x_s - E_s), with the (s, s) entry zeroed + arma::mat pw_grad_all = discrete_observations_dbl_t_ * E_all; + arma::mat diff_pw_all = discrete_observations_dbl_t_ + * (discrete_observations_dbl_ - E_all); + diff_pw_all.diag().zeros(); + + // Θ̄ coupling for the precision gradient: Θ̄_{s,t} = 2 diff_pw_s(t), + // diagonal = dot(x_s, x_s) - sum(E_sq) + Theta_bar = 2.0 * diff_pw_all.t(); + Theta_bar.diag() = diff_diag_all; + + // Pairwise discrete gradient: sum_i x_{i,t} * (x_{i,s}+1 - E_s) + // Factor 2: chain rule d/dK = 2 × d/dσ; edge (i, j) collects the + // variable-i and variable-j conditional contributions. + for(size_t i = 0; i + 1 < p_; ++i) { + for(size_t j = i + 1; j < p_; ++j) { + if(edge_indicators_(i, j) == 0) continue; + int loc = disc_index_cache_(i, j); + grad(loc) -= 2.0 * pw_grad_all(j, i); + grad(loc) -= 2.0 * pw_grad_all(i, j); + } + } - // Continuous mean gradient from marginal OMRF - for(size_t j = 0; j < q_; ++j) { - grad(main_effects_continuous_grad_offset_ + j) += 2.0 * temp_pairwise_cross(s, j) * sum_obs_minus_E; - } + // Pairwise_cross gradient from marginal OMRF (through Θ): + // ∂marginal_{st}/∂pairwise_effects_cross_{a,j} has two terms: + // = 2 [Σyy pairwise_effects_cross_t']_j δ_{as} + 2 [pairwise_effects_cross_s Σyy]_j δ_{at} + // Self-contribution (a=s): 4 diff_pw_s^T (A_xy Σyy) row s → diff_pw_all^T term + // Cross-contribution (a=t): 4 diff_pw_s(t) (A_xy Σyy) row s → diff_pw_all term + // Diagonal effective interaction: 4 diff_diag_s [A_xy Σyy] row s + // Rest-score bias: 2 sum_i(x_{is} - E_s) μy^T + arma::mat cross_grad_all = + 4.0 * ((diff_pw_all + diff_pw_all.t()) * cross_times_cov) + + 4.0 * (cross_times_cov.each_col() % diff_diag_all) + + 2.0 * (sum_obs_minus_E_all * temp_main_continuous.t()); - main_effects_discrete_offset += 2; + for(size_t s = 0; s < p_; ++s) { + for(size_t j = 0; j < q_; ++j) { + if(edge_indicators_(s, p_ + j) == 0) continue; + grad(cross_index_cache_(s, j)) += cross_grad_all(s, j); } } + // Continuous mean gradient from marginal OMRF: + // ∂l_s/∂main_effects_continuous_j = 2 pairwise_effects_cross_{sj} * sum_i (x_{is} - E_s) + arma::vec mean_grad_omrf = 2.0 * (temp_pairwise_cross.t() * sum_obs_minus_E_all); + for(size_t j = 0; j < q_; ++j) { + grad(main_effects_continuous_grad_offset_ + j) += mean_grad_omrf(j); + } + // Add numerator contribution to logp from discrete sufficient statistics // (already in grad_obs_cache_ as counts, but logp needs the actual dot-products) main_effects_discrete_offset = 0; diff --git a/src/models/omrf/omrf_model.cpp b/src/models/omrf/omrf_model.cpp index bb9d6449..91ad1f4d 100644 --- a/src/models/omrf/omrf_model.cpp +++ b/src/models/omrf/omrf_model.cpp @@ -704,6 +704,9 @@ std::pair OMRFModel::logp_and_gradient(const arma::vec& param } // ---- Per-variable: joint computation of log-normalizer and gradient ---- + // Expected-score vectors are collected per variable so the pairwise + // gradient reduces to one X^T * E GEMM after the loop. + arma::mat expected_scores(n_, p_, arma::fill::none); int offset = 0; for (int variable = 0; variable < num_variables; variable++) { const int num_cats = num_categories_(variable); @@ -728,15 +731,10 @@ std::pair OMRFModel::logp_and_gradient(const arma::vec& param gradient(offset + cat) -= arma::accu(logz_out_.probs.col(cat + 1)); } - // Pairwise gradient contributions (vectorized using BLAS) + // Expected value E_s[c+1|rest] per observation (pairwise gradient + // is assembled in one GEMM after the loop) arma::vec weights = arma::regspace(1, num_cats); - arma::vec E = logz_out_.probs.cols(1, num_cats) * weights; - arma::vec pw_grad = observations_double_t_ * E; - for (int j = 0; j < num_variables; j++) { - if (edge_indicators_(variable, j) == 0 || variable == j) continue; - int location = (variable < j) ? index_matrix_cache_(variable, j) : index_matrix_cache_(j, variable); - gradient(location) -= 2.0 * pw_grad(j); - } + expected_scores.col(variable) = logz_out_.probs.cols(1, num_cats) * weights; offset += num_cats; } else { const int ref = baseline_category_(variable); @@ -759,18 +757,26 @@ std::pair OMRFModel::logp_and_gradient(const arma::vec& param gradient(offset) -= arma::accu(logz_out_.probs * score); gradient(offset + 1) -= arma::accu(logz_out_.probs * sq_score); - // Pairwise gradient contributions (vectorized using BLAS) - arma::vec E = logz_out_.probs * score; - arma::vec pw_grad = observations_double_t_ * E; - for (int j = 0; j < num_variables; j++) { - if (edge_indicators_(variable, j) == 0 || variable == j) continue; - int location = (variable < j) ? index_matrix_cache_(variable, j) : index_matrix_cache_(j, variable); - gradient(location) -= 2.0 * pw_grad(j); - } + // Expected score per observation (pairwise gradient is assembled + // in one GEMM after the loop) + expected_scores.col(variable) = logz_out_.probs * score; offset += 2; } } + // ---- Pairwise gradient contributions ---- + // One X^T * E GEMM for all variables; edge (i, j) collects the + // variable-i and variable-j conditional contributions. + arma::mat pw_grad_all = observations_double_t_ * expected_scores; + for (int i = 0; i < num_variables - 1; i++) { + for (int j = i + 1; j < num_variables; j++) { + if (edge_indicators_(i, j) == 0) continue; + int location = index_matrix_cache_(i, j); + gradient(location) -= 2.0 * pw_grad_all(j, i); + gradient(location) -= 2.0 * pw_grad_all(i, j); + } + } + // ---- Priors: gradient contributions ---- offset = 0; for (int variable = 0; variable < num_variables; variable++) { From 51701f7f767c06b2a6dfac246e906f76bf696aac Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 9 Jul 2026 19:07:07 +0200 Subject: [PATCH 10/10] perf(mixed): fold the numerator logp pass into the Part-1 gradient loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The numerator loop recomputed the per-variable rest score verbatim from X_marginal — a second full O(np) pass with p n-vector allocations per leapfrog step. Accumulate the quadratic self-term, dot(x_s, rest), and the main-effect sufficient-statistic sums inside the Part-1 loop, which already holds the rest vector, and drop the second pass. Verification: fixed-seed ordinal NUTS draws remain bitwise identical; mixed 6+6 NUTS draws deviate by max 1.3e-9 from the pre-branch baseline over a 10+10-iteration horizon (logp accumulation order). FD gradient suites pass: test-mixed-gradient.R (8), test-mixed-theta-space.R (18), test-ggm-gradient.R (14). --- src/models/mixed/mixed_mrf_gradient.cpp | 44 ++++++++++--------------- 1 file changed, 17 insertions(+), 27 deletions(-) diff --git a/src/models/mixed/mixed_mrf_gradient.cpp b/src/models/mixed/mixed_mrf_gradient.cpp index aac3a354..b29f9bb2 100644 --- a/src/models/mixed/mixed_mrf_gradient.cpp +++ b/src/models/mixed/mixed_mrf_gradient.cpp @@ -276,6 +276,23 @@ std::pair MixedMRFModel::logp_and_gradient( - discrete_observations_dbl_.col(s) * precision_ss) + cross_bias(s); + // Numerator contribution to logp from discrete sufficient statistics + // (already in grad_obs_cache_ as counts, but logp needs the actual + // dot-products). Marginal self-interaction quadratic contribution, + // dot(x_s, rest), and the main-effect sums. + logp += precision_ss * arma::dot( + discrete_observations_dbl_.col(s), + discrete_observations_dbl_.col(s)); + logp += arma::dot(discrete_observations_dbl_.col(s), rest); + if(is_ordinal_variable_(s)) { + for(int c = 1; c <= C_s; ++c) { + logp += static_cast(counts_per_category_(c, s)) * temp_main_discrete(s, c - 1); + } + } else { + logp += temp_main_discrete(s, 0) * static_cast(blume_capel_stats_(0, s)) + + temp_main_discrete(s, 1) * static_cast(blume_capel_stats_(1, s)); + } + if(is_ordinal_variable_(s)) { arma::vec main_param = temp_main_discrete.row(s).cols(0, C_s - 1).t(); @@ -413,33 +430,6 @@ std::pair MixedMRFModel::logp_and_gradient( grad(main_effects_continuous_grad_offset_ + j) += mean_grad_omrf(j); } - // Add numerator contribution to logp from discrete sufficient statistics - // (already in grad_obs_cache_ as counts, but logp needs the actual dot-products) - main_effects_discrete_offset = 0; - for(size_t s = 0; s < p_; ++s) { - int C_s = num_categories_(s); - arma::vec rest; - double precision_ss = temp_marginal(s, s); - rest = 2.0 * (X_marginal.col(s) - - discrete_observations_dbl_.col(s) * precision_ss) - + cross_bias(s); - // Marginal self-interaction quadratic contribution - logp += precision_ss * arma::dot( - discrete_observations_dbl_.col(s), - discrete_observations_dbl_.col(s)); - // Numerator: dot(x_s, rest) + main-effect sums - logp += arma::dot(discrete_observations_dbl_.col(s), rest); - - if(is_ordinal_variable_(s)) { - for(int c = 1; c <= C_s; ++c) { - logp += static_cast(counts_per_category_(c, s)) * temp_main_discrete(s, c - 1); - } - } else { - logp += temp_main_discrete(s, 0) * static_cast(blume_capel_stats_(0, s)) - + temp_main_discrete(s, 1) * static_cast(blume_capel_stats_(1, s)); - } - } - // ========================================================================= // Part 2: GGM conditional log-likelihood and gradients // =========================================================================