From e9b99ae7c30559a1ce52745dc604058eaf38bb1c Mon Sep 17 00:00:00 2001 From: Maarten Marsman Date: Thu, 4 Jun 2026 17:10:34 +0200 Subject: [PATCH] refactor(ggm): share active<->full index traversal via for_each_active_full_pair get_full_vectorized_parameters (scatter active theta -> full zero-padded vector) and get_active_inv_mass (gather full-dim inverse mass -> active layout) open-coded the same per-column index arithmetic in inverse directions. Add GraphConstraintStructure::for_each_active_full_pair(fn), which visits each active entry as (active_index, full_index); the scatter and gather become one-line lambdas over it. Behavior-preserving: identical (active,full) index pairs in identical order; the assignments are unchanged. Verified bitwise: NUTS GGM fits with mass adaptation (edge-selection exercising included_indices, and no-edge) digest-identical pre/post. --- src/models/ggm/ggm_model.cpp | 38 ++++++--------------- src/models/ggm/graph_constraint_structure.h | 22 ++++++++++++ 2 files changed, 32 insertions(+), 28 deletions(-) diff --git a/src/models/ggm/ggm_model.cpp b/src/models/ggm/ggm_model.cpp index 501d7564..e42f0244 100644 --- a/src/models/ggm/ggm_model.cpp +++ b/src/models/ggm/ggm_model.cpp @@ -89,24 +89,11 @@ arma::vec GGMModel::get_full_vectorized_parameters() const { const auto& cs = constraint_structure_; arma::vec full(cs.full_dim, arma::fill::zeros); - for (size_t q = 0; q < p_; ++q) { - const auto& col = cs.columns[q]; - size_t active_offset = cs.theta_offsets[q]; - size_t full_offset = cs.full_theta_offsets[q]; - - // Copy f_q entries into their matching slots in the full vector. - // In the full vector, column q has q slots for off-diagonal + 1 for diagonal. - // The included indices map to specific positions. - for (size_t k = 0; k < col.d_q; ++k) { - // The k-th included index maps to position included_indices[k] in - // the column's off-diagonal block - size_t full_pos = full_offset + col.included_indices[k]; - full(full_pos) = theta_(active_offset + k); - } - - // psi_q is at the end of the column's block in both layouts - full(cs.full_psi_offset(q)) = theta_(active_offset + col.d_q); - } + // Scatter each active theta entry into its slot in the full (zero-padded) + // vector; excluded off-diagonal slots stay zero. + cs.for_each_active_full_pair([&](size_t active_idx, size_t full_idx) { + full(full_idx) = theta_(active_idx); + }); return full; } @@ -160,17 +147,12 @@ arma::vec GGMModel::get_active_inv_mass() const { // the theta entry at that slot directly, so the active inverse mass is // simply the included-slot subset — no N_q rotation needed. if (inv_mass_.n_elem == cs.full_dim) { + // Gather the full-dim inverse mass back into the active layout — the + // exact inverse of the get_full_vectorized_parameters scatter. arma::vec active(cs.active_dim); - for (size_t q = 0; q < p_; ++q) { - const auto& col = cs.columns[q]; - size_t active_offset = cs.theta_offsets[q]; - size_t full_offset = cs.full_theta_offsets[q]; - for (size_t k = 0; k < col.d_q; ++k) { - active(active_offset + k) = - inv_mass_(full_offset + col.included_indices[k]); - } - active(cs.psi_offset(q)) = inv_mass_(cs.full_psi_offset(q)); - } + cs.for_each_active_full_pair([&](size_t active_idx, size_t full_idx) { + active(active_idx) = inv_mass_(full_idx); + }); return active; } diff --git a/src/models/ggm/graph_constraint_structure.h b/src/models/ggm/graph_constraint_structure.h index 4e6f7abc..7c7ebf44 100644 --- a/src/models/ggm/graph_constraint_structure.h +++ b/src/models/ggm/graph_constraint_structure.h @@ -107,4 +107,26 @@ struct GraphConstraintStructure { size_t full_psi_offset(size_t q) const { return full_theta_offsets[q] + q; } + + /** + * Visit every active theta entry, invoking fn(active_index, full_index) + * with its slot in the active (packed) vector and the matching slot in the + * full (zero-padded) vector. Per column q: the d_q free off-diagonal + * entries (active theta_offsets[q]+k -> full full_theta_offsets[q]+ + * included_indices[k]), then psi_q (active psi_offset(q) -> full + * full_psi_offset(q)). The active<->full scatter and gather share this one + * traversal instead of duplicating the index arithmetic. + */ + template + void for_each_active_full_pair(Fn&& fn) const { + for (size_t q = 0; q < p; ++q) { + const auto& col = columns[q]; + size_t active_offset = theta_offsets[q]; + size_t full_offset = full_theta_offsets[q]; + for (size_t k = 0; k < col.d_q; ++k) { + fn(active_offset + k, full_offset + col.included_indices[k]); + } + fn(psi_offset(q), full_psi_offset(q)); + } + } };