[ENH] add a low memory condition for ALESubtraction - #1004
Merged
Conversation
Contributor
Reviewer's GuideIntroduces a low-memory execution path for ALESubtraction by chunking modeled activation (MA) maps into disk-backed CSR blocks, refactors permutation/null computation into reusable helpers, and adds tests to validate numerical equivalence and low-memory behavior. Sequence diagram for ALESubtraction low-memory permutation workflowsequenceDiagram
actor Researcher
participant ALE as ALESubtraction
participant StoreMgr as _managed_pairwise_ma_store
participant Store as _PairwiseMAStore
participant Kernel as KernelTransformer
participant Disk
Researcher->>ALE: fit(dataset1, dataset2)
ALE->>StoreMgr: _managed_pairwise_ma_store(maps_key1, coords_key1, maps_key2, coords_key2)
activate StoreMgr
StoreMgr->>ALE: _prepare_pairwise_ma_maps(...)
alt precomputed_MA_maps_present
ALE->>ALE: _collect_masked_ma_maps for maps_key1, maps_key2
ALE->>ALE: _compute_summarystat_est(ma_maps1)
ALE->>ALE: _compute_summarystat_est(ma_maps2)
ALE-->>StoreMgr: _PairwiseMAStore(in_memory CSR)
else coords_only_and_low_memory_triggered
ALE->>ALE: _estimate_group_ma_bytes(coords_key1)
ALE->>ALE: _estimate_group_ma_bytes(coords_key2)
ALE->>ALE: _should_use_low_memory(combined_nbytes)
ALE-->>StoreMgr: low_memory=True
loop per_group
ALE->>ALE: _determine_chunk_rows(bytes_per_study, available_bytes)
ALE->>ALE: _collect_chunked_ma_maps(coords_key, chunk_rows, prefix)
activate ALE
loop per_chunk
ALE->>Kernel: transform(chunk_coordinates, masker, return_type="sparse")
Kernel-->>ALE: csr_chunk
ALE->>Disk: _csr_to_memmap(csr_chunk, prefix)
Disk-->>ALE: memmapped_csr_chunk, temp_files
end
ALE-->>ALE: compute group_stat via log_sums
deactivate ALE
end
ALE-->>StoreMgr: _PairwiseMAStore(chunked_memmaps)
end
StoreMgr-->>ALE: yield Store
deactivate StoreMgr
ALE->>ALE: diff_ale_values = Store.group1_stat - Store.group2_stat
ALE->>ALE: _run_null_permutations(Store, n_iters, n_cores, diff_ale_values, iter_diff_values)
activate ALE
loop permutations
ALE->>ALE: _iterate_permutation_diffs(Store, n_iters, n_cores)
loop per_iteration (Parallel)
ALE->>ALE: _run_permutation(i_iter, Store)
activate ALE
ALE->>Store: compute_partition_summarystat(row_idx_group1)
Store-->>ALE: iter_grp1_ale_values
ALE->>Store: compute_partition_summarystat(row_idx_group2)
Store-->>ALE: iter_grp2_ale_values
ALE-->>ALE: iter_diff = iter_grp1_ale_values - iter_grp2_ale_values
deactivate ALE
end
ALE-->>ALE: update tail_counts, iter_abs_max, iter_diff_values
end
ALE-->>ALE: p_values, diff_signs = _finalize_alediff_tail_counts(...)
deactivate ALE
ALE->>StoreMgr: context exit
activate StoreMgr
StoreMgr->>Store: close()
Store->>Disk: _close_csr_memmaps, _cleanup_temp_files
Store-->>StoreMgr: closed
deactivate StoreMgr
ALE-->>Researcher: ALESubtractionResults
Class diagram for ALESubtraction low-memory MA-map storageclassDiagram
class ALESubtraction {
+bool_or_str low_memory
+float _low_memory_fraction
+int n_iters
+int n_cores
+object masker
+dict null_distributions_
+_run_permutation(i_iter, ma_store)
+_iterate_permutation_diffs(ma_store, n_iters, n_cores)
+_run_null_permutations(ma_store, n_iters, n_cores, diff_ale_values, iter_diff_values)
+_compute_cluster_nulls(iter_diff_values, voxel_thresh, n_iters)
+_should_use_low_memory(projected_nbytes)
+_estimate_group_ma_bytes(coords_key)
+_determine_chunk_rows(bytes_per_study, available_bytes)
+_collect_chunked_ma_maps(coords_key, chunk_rows, prefix)
+_prepare_pairwise_ma_maps(maps_key1, coords_key1, maps_key2, coords_key2)
+_managed_pairwise_ma_store(maps_key1, coords_key1, maps_key2, coords_key2)
}
class _PairwiseMAStore {
+object group1
+object group2
+np_ndarray group1_stat
+np_ndarray group2_stat
+list temp_files
+int n_group1
+int n_total
+int n_voxels
+compute_partition_summarystat(row_idx)
+close()
}
class _ChunkedCSRGroup {
+list chunks
+np_ndarray row_offsets
+tuple shape
}
class KernelTransformer {
+transform(coordinates, masker, return_type)
}
class Masker {
+inverse_transform(arr)
}
class NiftiImage {
+get_fdata(dtype)
}
ALESubtraction --> _PairwiseMAStore : uses
ALESubtraction --> KernelTransformer : uses
ALESubtraction --> Masker : uses
_PairwiseMAStore --> _ChunkedCSRGroup : group1, group2 may be
_ChunkedCSRGroup --> "*" sp_sparse_csr_matrix : chunks
class sp_sparse_csr_matrix {
+np_ndarray data
+np_ndarray indices
+np_ndarray indptr
+tuple shape
}
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 3 issues, and left some high level feedback:
- In
_compute_partition_ale_summarystat, it may be worth adding explicit shape checks (e.g., matchingn_voxelsand row counts vsn_grp1) so that misuse of mismatched MA groups fails fast with a clear error instead of producing silent misalignment. - The combination of
_estimate_group_ma_bytesand_collect_chunked_ma_mapsrecomputes MA maps for the same studies; if this becomes a bottleneck for large datasets, consider reusing the sample chunk or caching per-study byte estimates to avoid a full extra transform pass.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `_compute_partition_ale_summarystat`, it may be worth adding explicit shape checks (e.g., matching `n_voxels` and row counts vs `n_grp1`) so that misuse of mismatched MA groups fails fast with a clear error instead of producing silent misalignment.
- The combination of `_estimate_group_ma_bytes` and `_collect_chunked_ma_maps` recomputes MA maps for the same studies; if this becomes a bottleneck for large datasets, consider reusing the sample chunk or caching per-study byte estimates to avoid a full extra transform pass.
## Individual Comments
### Comment 1
<location path="nimare/meta/cbma/ale.py" line_range="1037-1045" />
<code_context>
+ return False
+ return projected_nbytes >= (available_bytes * self._low_memory_fraction)
+
+ def _estimate_group_ma_bytes(self, coords_key):
+ """Estimate total CSR bytes and bytes per study for one MA-map group."""
+ coordinates = self.inputs_[coords_key]
+ sample_df = next(
+ _iter_study_id_chunks(
+ coordinates, chunk_rows=min(32, len(np.unique(coordinates["id"].values)))
+ )
+ )
+ sample_ma = require_masked_csr(
+ self.kernel_transformer.transform(
+ sample_df,
</code_context>
<issue_to_address>
**issue:** Handle empty or malformed coordinate groups when estimating MA-map size.
If `self.inputs_[coords_key]` is empty (or all rows fail kernel transformation), `next(_iter_study_id_chunks(...))` will raise `StopIteration` and `require_masked_csr` may also fail. Consider explicitly handling an empty `coordinates` DataFrame (e.g., raise a clear `ValueError`) or catching `StopIteration` and re-raising with a more informative message to avoid a cryptic failure in low-memory mode when a group has no valid studies.
</issue_to_address>
### Comment 2
<location path="nimare/meta/cbma/ale.py" line_range="1102-1106" />
<code_context>
+ )
+ return ma_group, stat_values, temp_files
+
+ def _prepare_pairwise_ma_maps(self, maps_key1, coords_key1, maps_key2, coords_key2):
+ """Collect pairwise MA maps and optionally spill coordinate-generated maps to disk."""
+ temp_files = []
+
+ if maps_key1 in self.inputs_ or maps_key2 in self.inputs_:
+ if self.low_memory is not False:
+ LGR.info(
</code_context>
<issue_to_address>
**question (performance):** Clarify treatment when only one group provides precomputed MA maps.
This condition triggers the precomputed path even when only one side has MA maps, so low‑memory chunking is disabled for the coordinate‑generated side in mixed cases. If the goal is to skip low‑memory only when both groups provide MA maps, switch this to `and` or explicitly document the mixed-input behavior.
</issue_to_address>
### Comment 3
<location path="nimare/tests/test_meta_ale.py" line_range="936-945" />
<code_context>
+def test_ALESubtraction_partitioned_summarystat_matches_combined_path():
</code_context>
<issue_to_address>
**suggestion (testing):** Add coverage for `_compute_partition_ale_summarystat` with chunked MA-map groups
This test only covers the standard CSR path. `_compute_partition_ale_summarystat` also has a `_ChunkedCSRGroup` path that underpins the new low‑memory behavior. Please add a case using a small `_ChunkedCSRGroup` (e.g., split `ma_maps1` or `ma_maps2` into 2 chunks with correct `row_offsets`) and assert it matches the CSR result, so the chunked indexing logic is directly exercised.
Suggested implementation:
```python
def test_ALESubtraction_partitioned_summarystat_matches_combined_path():
"""Partitioned ALE summary stats should match the stacked sparse path.
This also exercises the _ChunkedCSRGroup path in _compute_partition_ale_summarystat
by wrapping one MA-map group in a small, two-chunk _ChunkedCSRGroup and checking
that the result matches the pure-CSR path.
"""
rng = np.random.default_rng(19)
n_grp1 = 5
n_grp2 = 7
n_voxels = 29
# Base MA-map groups as CSR (standard path)
ma_maps1 = sp_sparse.random(
n_grp1,
n_voxels,
density=0.2,
format="csr",
random_state=rng,
)
ma_maps2 = sp_sparse.random(
n_grp2,
n_voxels,
density=0.25,
format="csr",
random_state=rng,
)
# --- Standard CSR path ---
# Depending on the actual signature, adjust arguments so that this call exercises
# the non-chunked CSR path of _compute_partition_ale_summarystat.
group_sizes = np.array([n_grp1, n_grp2], dtype=int)
csr_groups = [ma_maps1, ma_maps2]
csr_summarystat = ale._compute_partition_ale_summarystat(
csr_groups,
group_sizes,
low_memory=False,
)
# --- Chunked CSR path ---
# Split ma_maps1 into two CSR chunks and wrap in _ChunkedCSRGroup to directly
# exercise the chunked indexing logic in _compute_partition_ale_summarystat.
split_idx = n_grp1 // 2
chunks = [
ma_maps1[:split_idx, :],
ma_maps1[split_idx:, :],
]
row_offsets = np.array([0, split_idx, n_grp1], dtype=int)
# Adjust the constructor signature if necessary to match the actual
# _ChunkedCSRGroup implementation.
chunked_grp1 = ale._ChunkedCSRGroup(
chunks=chunks,
row_offsets=row_offsets,
n_rows=n_grp1,
n_cols=n_voxels,
)
chunked_groups = [chunked_grp1, ma_maps2]
chunked_summarystat = ale._compute_partition_ale_summarystat(
chunked_groups,
group_sizes,
low_memory=True,
)
# The chunked-group path should be numerically equivalent to the pure-CSR path.
np.testing.assert_allclose(chunked_summarystat, csr_summarystat)
```
1. Ensure that `ale` and `sp_sparse` are imported in this test module. For example:
- `from nimare.meta import ale`
- `from scipy import sparse as sp_sparse`
2. Adjust the call signature for `ale._compute_partition_ale_summarystat` to match the actual implementation. In many codebases this might instead take a single stacked matrix plus partition boundaries (e.g., `stacked_ma_maps` and `group_bounds`), or additional keyword arguments such as `voxel_idx`, `return_components`, etc.
3. Adjust the `_ChunkedCSRGroup` constructor arguments to match the real class definition. Common patterns are:
- `_ChunkedCSRGroup(chunks, row_offsets)`
- or `_ChunkedCSRGroup(chunks=chunks, row_offsets=row_offsets)` without `n_rows`/`n_cols` if those are inferred.
4. If `_compute_partition_ale_summarystat` returns a tuple (e.g., `(summarystat, extra_info)`), unpack it accordingly and only compare the summary statistic component in the `assert_allclose`.
5. If the test suite uses a different RNG pattern (e.g. `RandomState` instead of `default_rng`), align the random number generator usage with the rest of the file for reproducibility.
</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
+1037
to
+1045
| def _estimate_group_ma_bytes(self, coords_key): | ||
| """Estimate total CSR bytes and bytes per study for one MA-map group.""" | ||
| coordinates = self.inputs_[coords_key] | ||
| sample_df = next( | ||
| _iter_study_id_chunks( | ||
| coordinates, chunk_rows=min(32, len(np.unique(coordinates["id"].values))) | ||
| ) | ||
| ) | ||
| sample_ma = require_masked_csr( |
Contributor
There was a problem hiding this comment.
issue: Handle empty or malformed coordinate groups when estimating MA-map size.
If self.inputs_[coords_key] is empty (or all rows fail kernel transformation), next(_iter_study_id_chunks(...)) will raise StopIteration and require_masked_csr may also fail. Consider explicitly handling an empty coordinates DataFrame (e.g., raise a clear ValueError) or catching StopIteration and re-raising with a more informative message to avoid a cryptic failure in low-memory mode when a group has no valid studies.
Comment on lines
+1102
to
+1106
| def _prepare_pairwise_ma_maps(self, maps_key1, coords_key1, maps_key2, coords_key2): | ||
| """Collect pairwise MA maps and optionally spill coordinate-generated maps to disk.""" | ||
| temp_files = [] | ||
|
|
||
| if maps_key1 in self.inputs_ or maps_key2 in self.inputs_: |
Contributor
There was a problem hiding this comment.
question (performance): Clarify treatment when only one group provides precomputed MA maps.
This condition triggers the precomputed path even when only one side has MA maps, so low‑memory chunking is disabled for the coordinate‑generated side in mixed cases. If the goal is to skip low‑memory only when both groups provide MA maps, switch this to and or explicitly document the mixed-input behavior.
Comment on lines
+936
to
+945
| def test_ALESubtraction_partitioned_summarystat_matches_combined_path(): | ||
| """Partitioned ALE summary stats should match the stacked sparse path.""" | ||
| rng = np.random.default_rng(19) | ||
| n_grp1 = 5 | ||
| n_grp2 = 7 | ||
| n_voxels = 29 | ||
|
|
||
| ma_maps1 = sp_sparse.random( | ||
| n_grp1, | ||
| n_voxels, |
Contributor
There was a problem hiding this comment.
suggestion (testing): Add coverage for _compute_partition_ale_summarystat with chunked MA-map groups
This test only covers the standard CSR path. _compute_partition_ale_summarystat also has a _ChunkedCSRGroup path that underpins the new low‑memory behavior. Please add a case using a small _ChunkedCSRGroup (e.g., split ma_maps1 or ma_maps2 into 2 chunks with correct row_offsets) and assert it matches the CSR result, so the chunked indexing logic is directly exercised.
Suggested implementation:
def test_ALESubtraction_partitioned_summarystat_matches_combined_path():
"""Partitioned ALE summary stats should match the stacked sparse path.
This also exercises the _ChunkedCSRGroup path in _compute_partition_ale_summarystat
by wrapping one MA-map group in a small, two-chunk _ChunkedCSRGroup and checking
that the result matches the pure-CSR path.
"""
rng = np.random.default_rng(19)
n_grp1 = 5
n_grp2 = 7
n_voxels = 29
# Base MA-map groups as CSR (standard path)
ma_maps1 = sp_sparse.random(
n_grp1,
n_voxels,
density=0.2,
format="csr",
random_state=rng,
)
ma_maps2 = sp_sparse.random(
n_grp2,
n_voxels,
density=0.25,
format="csr",
random_state=rng,
)
# --- Standard CSR path ---
# Depending on the actual signature, adjust arguments so that this call exercises
# the non-chunked CSR path of _compute_partition_ale_summarystat.
group_sizes = np.array([n_grp1, n_grp2], dtype=int)
csr_groups = [ma_maps1, ma_maps2]
csr_summarystat = ale._compute_partition_ale_summarystat(
csr_groups,
group_sizes,
low_memory=False,
)
# --- Chunked CSR path ---
# Split ma_maps1 into two CSR chunks and wrap in _ChunkedCSRGroup to directly
# exercise the chunked indexing logic in _compute_partition_ale_summarystat.
split_idx = n_grp1 // 2
chunks = [
ma_maps1[:split_idx, :],
ma_maps1[split_idx:, :],
]
row_offsets = np.array([0, split_idx, n_grp1], dtype=int)
# Adjust the constructor signature if necessary to match the actual
# _ChunkedCSRGroup implementation.
chunked_grp1 = ale._ChunkedCSRGroup(
chunks=chunks,
row_offsets=row_offsets,
n_rows=n_grp1,
n_cols=n_voxels,
)
chunked_groups = [chunked_grp1, ma_maps2]
chunked_summarystat = ale._compute_partition_ale_summarystat(
chunked_groups,
group_sizes,
low_memory=True,
)
# The chunked-group path should be numerically equivalent to the pure-CSR path.
np.testing.assert_allclose(chunked_summarystat, csr_summarystat)- Ensure that
aleandsp_sparseare imported in this test module. For example:from nimare.meta import alefrom scipy import sparse as sp_sparse
- Adjust the call signature for
ale._compute_partition_ale_summarystatto match the actual implementation. In many codebases this might instead take a single stacked matrix plus partition boundaries (e.g.,stacked_ma_mapsandgroup_bounds), or additional keyword arguments such asvoxel_idx,return_components, etc. - Adjust the
_ChunkedCSRGroupconstructor arguments to match the real class definition. Common patterns are:_ChunkedCSRGroup(chunks, row_offsets)- or
_ChunkedCSRGroup(chunks=chunks, row_offsets=row_offsets)withoutn_rows/n_colsif those are inferred.
- If
_compute_partition_ale_summarystatreturns a tuple (e.g.,(summarystat, extra_info)), unpack it accordingly and only compare the summary statistic component in theassert_allclose. - If the test suite uses a different RNG pattern (e.g.
RandomStateinstead ofdefault_rng), align the random number generator usage with the rest of the file for reproducibility.
jdkent
commented
Apr 3, 2026
jdkent
left a comment
jdkent
left a comment
Member
Author
There was a problem hiding this comment.
fix this test for windows.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1004 +/- ##
==========================================
+ Coverage 85.38% 85.81% +0.42%
==========================================
Files 52 52
Lines 9256 9480 +224
==========================================
+ Hits 7903 8135 +232
+ Misses 1353 1345 -8 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
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 # .
The other major blockage in the AWS cloud storage logs is trying to run a large reference studyset (like neurostore) with ALEsubtraction. This pull request lowers the memory cost for running this algorithm, and hopefully helps in some low memory environments.
Changes proposed in this pull request:
Summary by Sourcery
Introduce a low-memory execution path for ALESubtraction and refactor permutation handling to support disk-backed, chunked MA maps while preserving existing statistical behavior.
Enhancements:
Tests: