[REF] CBMR internal and public API refactor - #1005
Merged
Merged
Conversation
Profile and refactor the CBMR hot paths while preserving numerical behavior within tolerance. - speed up summary table construction and inference result copying - cache tensorized inputs, log-intensity values, and covariance terms - vectorize multi-contrast GLH paths and add analytic Poisson Fisher info - trim the synthetic CBMR fixture and expand regression coverage - remove stale CBMR/model code, fix lint issues, and support CUDA devices
Contributor
Reviewer's GuideRefactors the CBMR estimator and GLM model stack for experiment-level handling, deterministic/cuda-aware execution, and more efficient tensor/Fisher computations, while tightening coordinate/mask preprocessing and greatly expanding tests for reproducibility and numerical correctness. Sequence diagram for CBMR fitting with tensor input caching and cuda aware executionsequenceDiagram
actor User
participant Estimator as CBMREstimator
participant Model as GeneralLinearModelEstimator
participant Torch as torch
User->>Estimator: instantiate(model, device)
Estimator->>Estimator: check device with _uses_cuda
Estimator->>Torch: cuda.is_available()
Torch-->>Estimator: availability
Estimator->>Model: set device
User->>Estimator: fit(dataset)
Estimator->>Estimator: _preprocess_input(dataset)
Estimator->>Model: init_weights(groups, moderators, spatial_coef_dim, moderators_coef_dim)
Estimator->>Torch: manual_seed(100)
Estimator->>Torch: cuda.manual_seed_all(100) if _uses_cuda
Estimator->>Model: fit(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment)
Model->>Model: _prepare_tensor_inputs(...)
Model->>Model: _optimizer(tensor_inputs)
Model->>Model: extract_optimized_params(coef_spline_bases, moderators_by_group)
Model->>Model: standard_error_estimation(tensor_inputs)
Model->>Model: _cache_tensor_inputs(original_inputs, tensor_inputs)
Model-->>Estimator: maps, tables via summary()
Sequence diagram for CBMR inference with Fisher caching and tensor reusesequenceDiagram
actor User
participant Inference as CBMRInference
participant Result as CBMRResult
participant Estimator as CBMREstimator
participant Model as GeneralLinearModelEstimator
User->>Inference: __init__(device)
Inference->>Inference: _reset_inference_caches()
User->>Inference: fit(result)
Inference->>Inference: _copy_result_for_inference(result)
Inference->>Estimator: set device
Inference->>Model: set device and to(device)
Inference->>Model: _invalidate_tensor_inputs_cache()
Inference->>Inference: build group_reference_dict and moderator_reference_dict
User->>Inference: transform(t_con_groups, t_con_moderators)
Inference->>Inference: _preprocess_t_con_regressor(groups)
Inference->>Inference: _glh_con_group()
Inference->>Inference: _get_group_spatial_covariance(involved_groups)
Inference->>Model: fisher_info_multiple_group_spatial(...)
Model->>Model: _resolve_tensor_inputs(...)
Model-->>Inference: Fisher matrix
Inference->>Inference: cache covariance in _group_spatial_covariance_cache
Inference->>Inference: _chi_square_log_intensity(...)
Inference->>Inference: _glh_con_moderator()
Inference->>Inference: _get_moderator_covariance()
Inference->>Model: fisher_info_multiple_group_moderator(...)
Model-->>Inference: moderator Fisher matrix
Inference->>Inference: cache covariance and variance
Inference-->>User: updated maps and tables for contrasts
Class diagram for refactored CBMR estimator and GLM model stackclassDiagram
class CBMREstimator {
+dict _required_inputs
+string _group_column
+list groups
+list moderators
+string device
+dict inputs_
+__init__(model, penalty, spline_spacing, sample_size_moderator, group_categories, moderators, lr, lr_decay, tol, device)
+_generate_description()
+_get_mask_img(dataset)
+_validate_coordinates()
+_filter_coordinates_to_mask(coordinates, mask_img, mask_data, mask_lookup)
+_build_experiment_annotations(dataset)
+_build_group_moderators(experiment_annotations)
+_build_group_foci(coordinates, ids_by_group, n_mask_voxels)
+_preprocess_input(dataset)
+_fit(dataset)
}
class CBMRInference {
+string device
+result
+list groups
+list moderators
+dict group_reference_dict
+dict moderator_reference_dict
+__init__(device)
+fit(result)
+display()
+create_regular_expressions()
+create_contrast(contrast_name, source)
+transform(t_con_groups, t_con_moderators)
+_copy_result_for_inference(result)
+_reset_inference_caches()
+_get_group_log_intensity(group)
+_get_group_null_log_intensity(group)
+_get_group_spatial_covariance(involved_groups)
+_get_moderator_covariance()
+_glh_con_group()
+_glh_con_moderator()
+_chi_square_log_intensity(m, n_brain_voxel, n_con_group_involved, simp_con_group, contrast_log_intensity, cov_log_intensity)
}
class _CBMRTensorInputs {
+Tensor coef_spline_bases
+object moderators_by_group
+dict foci_per_voxel
+dict foci_per_experiment
+subset(groups) _CBMRTensorInputs
}
class GeneralLinearModelEstimator {
+list groups
+list moderators
+int spatial_coef_dim
+int moderators_coef_dim
+bool penalty
+float lr
+float lr_decay
+float tol
+string device
+dict spatial_regression_coef
+dict spatial_intensity_estimation
+dict spatial_regression_coef_se
+dict log_spatial_intensity_se
+dict spatial_intensity_se
+numpy.ndarray moderators_coef
+dict moderators_effect
+Tensor coef_spline_bases
+object _tensor_inputs_cache
+tuple _tensor_inputs_cache_keys
+__init__(penalty, lr, lr_decay, tol, device)
+init_spatial_weights()
+init_moderator_weights()
+init_weights(groups, moderators, spatial_coef_dim, moderators_coef_dim)
+_invalidate_tensor_inputs_cache()
+_to_numpy_array(array_like)
+_flatten_tensor(tensor)
+_as_float_tensor(array_like)
+_prepare_tensor_inputs(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment)
+_cache_tensor_inputs(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment, tensor_inputs)
+_resolve_tensor_inputs(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment)
+_frame_from_uniform_group_dict(group_values) DataFrame
+_optimizer(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment)
+fit(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment)
+extract_optimized_params(coef_spline_bases, moderators_by_group)
+standard_error_estimation(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment)
+summary() maps tables
+fisher_info_multiple_group_spatial(involved_groups, coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment) numpy.ndarray
+fisher_info_multiple_group_moderator(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment) numpy.ndarray
+firth_penalty(foci_per_voxel, foci_per_experiment, moderators, coef_spline_bases, overdispersion, overdispersion_coef, square_root)
<<abstract>> _log_likelihood_single_group(...)
<<abstract>> _log_likelihood_mult_group(...)
<<abstract>> forward(coef_spline_bases, moderators, foci_per_voxel, foci_per_experiment)
}
class OverdispersionModelEstimator {
+dict overdispersion
+bool square_root
+init_overdispersion_weights()
+init_weights(groups, moderators, spatial_coef_dim, moderators_coef_dim)
+inference_outcome(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment) maps tables
}
class PoissonEstimator {
+fisher_info_multiple_group_spatial(involved_groups, coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment) numpy.ndarray
+fisher_info_multiple_group_moderator(coef_spline_bases, moderators_by_group, foci_per_voxel, foci_per_experiment) numpy.ndarray
+_log_likelihood_single_group(group_spatial_coef, moderators_coef, coef_spline_bases, group_moderators, group_foci_per_voxel, group_foci_per_experiment, device)
+_log_likelihood_mult_group(spatial_coef, moderator_coef, coef_spline_bases, foci_per_voxel, foci_per_experiment, moderators, device)
+forward(coef_spline_bases, moderators, foci_per_voxel, foci_per_experiment)
}
class NegativeBinomialEstimator {
+_log_likelihood_single_group(group_spatial_coef, moderators_coef, group_overdispersion, coef_spline_bases, group_moderators, group_foci_per_voxel, group_foci_per_experiment, device)
+_log_likelihood_mult_group(overdispersion_coef, spatial_coef, coef_spline_bases, foci_per_voxel, foci_per_experiment, moderator_coef, moderators, device)
+forward(coef_spline_bases, moderators, foci_per_voxel, foci_per_experiment)
}
class ClusteredNegativeBinomialEstimator {
+_log_likelihood_single_group(group_spatial_coef, moderators_coef, group_overdispersion, coef_spline_bases, group_moderators, group_foci_per_voxel, group_foci_per_experiment, device)
+_log_likelihood_mult_group(overdispersion_coef, spatial_coef, coef_spline_bases, foci_per_voxel, foci_per_experiment, moderator_coef, moderators, device)
+forward(coef_spline_bases, moderators, foci_per_voxel, foci_per_experiment)
}
CBMREstimator ..> GeneralLinearModelEstimator : uses model
CBMRInference ..> CBMREstimator : uses estimator
CBMRInference ..> GeneralLinearModelEstimator : calls fisher_info_*
GeneralLinearModelEstimator ..> _CBMRTensorInputs : creates
OverdispersionModelEstimator --|> GeneralLinearModelEstimator
PoissonEstimator --|> GeneralLinearModelEstimator
NegativeBinomialEstimator --|> OverdispersionModelEstimator
ClusteredNegativeBinomialEstimator --|> OverdispersionModelEstimator
Flow diagram for refactored CBMR preprocessing and input constructionflowchart TD
A[Start _preprocess_input] --> B[Get masker and mask image using _get_mask_img]
B --> C[Validate coordinate spaces with _validate_coordinates]
C --> D[Compute boolean mask_data from mask_img]
D --> E[Build mask_lookup and count n_mask_voxels]
E --> F[Compute B spline bases X with b_spline_bases using mask_data and spacing]
F --> G[Store coef_spline_bases in inputs_]
G --> H[Filter coordinates to mask with _filter_coordinates_to_mask]
H --> I[Drop helper index from stored coordinates]
I --> J[Build per experiment annotations and ids_by_group with _build_experiment_annotations]
J --> K[Store ids_by_group and groups in estimator]
K --> L[Build moderators_by_group with _build_group_moderators]
L --> M[Build foci_per_voxel and foci_per_experiment with _build_group_foci]
M --> N[Store foci_per_voxel and foci_per_experiment in inputs_]
N --> O[End _preprocess_input]
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- CBMREstimator._fit and GeneralLinearModelEstimator._optimizer both call torch.manual_seed(100), which globally resets the RNG on every fit; consider seeding only once (e.g., in tests or via an explicit parameter) to avoid surprising users who rely on PyTorch randomness elsewhere.
- In CBMRInference._copy_result_for_inference you downcast maps to DEFAULT_FLOAT_DTYPE when they were previously double tensors; if DEFAULT_FLOAT_DTYPE is float32 this silently reduces precision—consider preserving the original dtype or making the cast opt-in.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- CBMREstimator._fit and GeneralLinearModelEstimator._optimizer both call torch.manual_seed(100), which globally resets the RNG on every fit; consider seeding only once (e.g., in tests or via an explicit parameter) to avoid surprising users who rely on PyTorch randomness elsewhere.
- In CBMRInference._copy_result_for_inference you downcast maps to DEFAULT_FLOAT_DTYPE when they were previously double tensors; if DEFAULT_FLOAT_DTYPE is float32 this silently reduces precision—consider preserving the original dtype or making the cast opt-in.
## Individual Comments
### Comment 1
<location path="nimare/tests/test_meta_cbmr.py" line_range="386-395" />
<code_context>
+def test_poisson_analytic_fisher_matches_generic_hessian():
</code_context>
<issue_to_address>
**suggestion (testing):** Use explicit tolerances when comparing analytic and Hessian-based Fisher information matrices.
This test compares analytic and Hessian-based Fisher matrices using the default `assert_allclose` tolerances. Because both paths use numerics/AD, small floating-point differences are expected even when correct. Please set explicit `rtol`/`atol` (e.g., `rtol=1e-6, atol=1e-8`) for both spatial and moderator Fisher checks to reduce flakiness.
Suggested implementation:
```python
np.testing.assert_allclose(
spatial_fisher_analytic,
spatial_fisher_hessian,
rtol=1e-6,
atol=1e-8,
)
```
```python
np.testing.assert_allclose(
moderator_fisher_analytic,
moderator_fisher_hessian,
rtol=1e-6,
atol=1e-8,
)
```
If the test currently uses a different assertion helper (e.g., `numpy.testing.assert_allclose` imported as `assert_allclose`, or `torch.testing.assert_close`), adjust the two SEARCH patterns to match the existing calls and add `rtol=1e-6, atol=1e-8` (or the equivalent keyword argument names for that helper). Make sure you apply these explicit tolerances to *both* the spatial and moderator Fisher comparisons within `test_poisson_analytic_fisher_matches_generic_hessian`.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+386
to
+395
| def test_poisson_analytic_fisher_matches_generic_hessian(): | ||
| """Analytic Poisson Fisher matrices should match the generic Hessian route.""" | ||
| model = models.PoissonEstimator(device="cpu") | ||
| groups = ["A", "B"] | ||
| moderators = ["m1", "m2"] | ||
| model.init_weights( | ||
| groups=groups, | ||
| moderators=moderators, | ||
| spatial_coef_dim=2, | ||
| moderators_coef_dim=2, |
Contributor
There was a problem hiding this comment.
suggestion (testing): Use explicit tolerances when comparing analytic and Hessian-based Fisher information matrices.
This test compares analytic and Hessian-based Fisher matrices using the default assert_allclose tolerances. Because both paths use numerics/AD, small floating-point differences are expected even when correct. Please set explicit rtol/atol (e.g., rtol=1e-6, atol=1e-8) for both spatial and moderator Fisher checks to reduce flakiness.
Suggested implementation:
np.testing.assert_allclose(
spatial_fisher_analytic,
spatial_fisher_hessian,
rtol=1e-6,
atol=1e-8,
) np.testing.assert_allclose(
moderator_fisher_analytic,
moderator_fisher_hessian,
rtol=1e-6,
atol=1e-8,
)If the test currently uses a different assertion helper (e.g., numpy.testing.assert_allclose imported as assert_allclose, or torch.testing.assert_close), adjust the two SEARCH patterns to match the existing calls and add rtol=1e-6, atol=1e-8 (or the equivalent keyword argument names for that helper). Make sure you apply these explicit tolerances to both the spatial and moderator Fisher comparisons within test_poisson_analytic_fisher_matches_generic_hessian.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1005 +/- ##
==========================================
+ Coverage 85.38% 85.90% +0.52%
==========================================
Files 52 52
Lines 9256 9818 +562
==========================================
+ Hits 7903 8434 +531
- Misses 1353 1384 +31 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
jdkent
commented
Apr 7, 2026
jdkent
left a comment
jdkent
left a comment
Member
Author
There was a problem hiding this comment.
I have a couple comments/questions about the implementation.
| return masker, mask_img | ||
|
|
||
| def _validate_coordinates(self): | ||
| """Validate coordinate space metadata consistently with other CBMA estimators.""" |
Member
Author
There was a problem hiding this comment.
This should be a universal utility function or on a base class if this is being checked for most estimators.
Comment on lines
+69
to
+75
| def available_groups(self): | ||
| """Return the fitted groups that can be used in CBMR contrasts.""" | ||
| return list(self.groups) | ||
|
|
||
| def available_moderators(self): | ||
| """Return the fitted moderators that can be used in CBMR contrasts.""" | ||
| return list(self.moderators) |
Member
Author
There was a problem hiding this comment.
these helpers are not helpful, remove
Comment on lines
+80
to
+81
| "groups": self.available_groups(), | ||
| "moderators": self.available_moderators(), |
Member
Author
There was a problem hiding this comment.
use self.groups and self.moderators
Comment on lines
+505
to
+506
| if self.group_categories is None: | ||
| experiment_annotations[self._group_column] = "Default" |
Member
Author
There was a problem hiding this comment.
"Default" should be a global variable stated at the top of the file and referenced here.
Comment on lines
+208
to
+214
| ClusteredNegativeBinomial This method is also an efficient but less | ||
| accurate approach. Clustered NB model is | ||
| "random effect" Poisson model, which asserts | ||
| that the random effects are latent | ||
| characteristics of each experiment, and | ||
| represent a shared effect over the entire brain | ||
| for a given experiment. |
Member
Author
There was a problem hiding this comment.
does this still render appropriately?
| self.model.device = self.device | ||
|
|
||
| # Initialize optimisation parameters | ||
| self.iter = 0 |
Member
Author
There was a problem hiding this comment.
when is self.iter used? should it be a self attribute?
| dataset = normalize_collection(dataset) | ||
| self._collect_inputs(dataset, drop_invalid=drop_invalid) | ||
| self._preprocess_input(dataset) | ||
| maps, tables, description = self._cache(self._fit, func_memory_level=1)(dataset) |
Member
Author
There was a problem hiding this comment.
where is the definition of the _cache function
Comment on lines
+710
to
+712
| torch.manual_seed(100) | ||
| if _uses_cuda(self.device): | ||
| torch.cuda.manual_seed_all(100) |
Member
Author
There was a problem hiding this comment.
is there a general purpose way to set seeds for either cpu/gpu depending on request?
Comment on lines
+1414
to
+1415
| # Build the masked design matrix directly instead of materializing the | ||
| # full sparse Kronecker product and then indexing into it. |
Member
Author
There was a problem hiding this comment.
this comment should just explain the current implementation without referencing the previous implementation.
Member
Author
|
@sourcery-ai review |
Contributor
|
Sorry @jdkent, your pull request is larger than the review limit of 150000 diff characters |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes # .
Changes proposed in this pull request:
infermethod)Summary by Sourcery
Refactor CBMR meta-regression and inference to operate at the experiment level, improve numerical stability and performance (including CUDA support and tensor caching), and strengthen reproducibility and test coverage.
New Features:
Enhancements:
Tests: