Fold generation to support optional stratification - #95
Conversation
📝 WalkthroughWalkthroughA function refactoring replaces Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@brainsets/utils/split.py`:
- Around line 279-316: The current pre-validation around n_folds, val_ratio and
stratify_by can still allow runtime failures when later using sklearn splitters
(e.g., outer folds leaving <2 train+val samples or stratified splits where a
class has too few samples); update the checks in the block handling
use_stratify, class_labels and the non-stratified branch to compute and enforce
minimum per-fold sample counts before calling
KFold/StratifiedKFold/ShuffleSplit/StratifiedShuffleSplit: for non-stratified
ensure len(intervals) >= n_folds * ceil(1/(1 - val_ratio)) (or at least 2
samples in any train+val after an outer split), and for stratified ensure every
unique value in getattr(intervals, stratify_by) has enough samples to be present
in each fold (i.e., count_per_class >= n_folds and also satisfies the train/val
minimum given val_ratio); raise clear ValueError messages referencing n_folds,
val_ratio, stratify_by, intervals, and class_labels when these conditions are
not met.
In `@tests/test_split_utils.py`:
- Line 173: Update the zip call in the test loop so it explicitly enforces
equal-length iterables: change the loop that iterates "for f1, f2 in zip(folds1,
folds2):" to pass strict=True (i.e., zip(folds1, folds2, strict=True)) so the
linter Ruff B905 is satisfied and the intent to require equal lengths for folds1
and folds2 is explicit.
ℹ️ Review info
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
brainsets/utils/split.pybrainsets_pipelines/kemp_sleep_edf_2013/pipeline.pytests/test_split_utils.py
| if n_folds < 2: | ||
| raise ValueError(f"n_folds must be at least 2, got {n_folds}") | ||
| if not (0.0 < val_ratio < 1.0): | ||
| raise ValueError( | ||
| f"val_ratio must be between 0 and 1 (exclusive), got {val_ratio}" | ||
| ) | ||
| if stratify_by is not None and not hasattr(intervals, stratify_by): | ||
| raise ValueError( | ||
| f"Intervals must have a '{stratify_by}' attribute for stratification." | ||
| ) | ||
|
|
||
| try: | ||
| from sklearn.model_selection import StratifiedKFold, StratifiedShuffleSplit | ||
| from sklearn.model_selection import ( | ||
| KFold, | ||
| StratifiedKFold, | ||
| ShuffleSplit, | ||
| StratifiedShuffleSplit, | ||
| ) | ||
| except ImportError: | ||
| raise ImportError( | ||
| "This function requires the scikit-learn library which you can install with " | ||
| "`pip install scikit-learn`" | ||
| ) | ||
|
|
||
| if not hasattr(intervals, stratify_by): | ||
| raise ValueError( | ||
| f"Intervals must have a '{stratify_by}' attribute for stratification." | ||
| ) | ||
| use_stratify = stratify_by is not None | ||
|
|
||
| class_labels = getattr(intervals, stratify_by) | ||
| if len(class_labels) < n_folds: | ||
| raise ValueError( | ||
| f"Not enough samples ({len(class_labels)}) for {n_folds} folds." | ||
| ) | ||
| if use_stratify: | ||
| class_labels = getattr(intervals, stratify_by) | ||
| if len(class_labels) < n_folds: | ||
| raise ValueError( | ||
| f"Not enough samples ({len(class_labels)}) for {n_folds} folds." | ||
| ) | ||
| else: | ||
| if len(intervals) < n_folds: | ||
| raise ValueError( | ||
| f"Not enough samples ({len(intervals)}) for {n_folds} folds." | ||
| ) | ||
|
|
There was a problem hiding this comment.
Strengthen pre-validation for fold feasibility before calling sklearn splitters.
Current checks can still allow runtime failures in both modes (e.g., tiny datasets where an outer fold leaves <2 train+valid samples, or stratified labels that violate per-class constraints). This should be rejected early with clear errors.
Proposed guardrails
def generate_folds(
@@
- use_stratify = stratify_by is not None
+ n_samples = len(intervals)
+ if n_samples < n_folds:
+ raise ValueError(f"Not enough samples ({n_samples}) for {n_folds} folds.")
+
+ # Smallest train+valid size across outer folds (largest test fold removed)
+ min_train_val_size = n_samples - int(np.ceil(n_samples / n_folds))
+ if min_train_val_size < 2:
+ raise ValueError(
+ "Each outer fold must leave at least 2 samples for train+valid; "
+ "reduce n_folds or provide more samples."
+ )
+
+ use_stratify = stratify_by is not None
if use_stratify:
- class_labels = getattr(intervals, stratify_by)
- if len(class_labels) < n_folds:
+ class_labels = np.asarray(getattr(intervals, stratify_by))
+ if len(class_labels) != n_samples:
+ raise ValueError(
+ f"'{stratify_by}' must contain one label per interval "
+ f"({n_samples}), got {len(class_labels)}."
+ )
+ _, class_counts = np.unique(class_labels, return_counts=True)
+ if np.min(class_counts) < n_folds:
+ raise ValueError(
+ f"Each class in '{stratify_by}' must have at least {n_folds} samples."
+ )
+
+ n_classes = len(class_counts)
+ min_valid_size = int(np.ceil(val_ratio * min_train_val_size))
+ min_train_size = min_train_val_size - min_valid_size
+ if min_valid_size < n_classes or min_train_size < n_classes:
raise ValueError(
- f"Not enough samples ({len(class_labels)}) for {n_folds} folds."
+ "val_ratio is too small/large for stratified train/valid split; "
+ "train and valid must each be able to include all classes."
)
- else:
- if len(intervals) < n_folds:
- raise ValueError(
- f"Not enough samples ({len(intervals)}) for {n_folds} folds."
- )Also applies to: 320-334
🧰 Tools
🪛 Ruff (0.15.2)
[warning] 280-280: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 282-284: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 286-288: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 298-301: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
[warning] 298-301: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 308-310: Avoid specifying long messages outside the exception class
(TRY003)
[warning] 313-315: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@brainsets/utils/split.py` around lines 279 - 316, The current pre-validation
around n_folds, val_ratio and stratify_by can still allow runtime failures when
later using sklearn splitters (e.g., outer folds leaving <2 train+val samples or
stratified splits where a class has too few samples); update the checks in the
block handling use_stratify, class_labels and the non-stratified branch to
compute and enforce minimum per-fold sample counts before calling
KFold/StratifiedKFold/ShuffleSplit/StratifiedShuffleSplit: for non-stratified
ensure len(intervals) >= n_folds * ceil(1/(1 - val_ratio)) (or at least 2
samples in any train+val after an outer split), and for stratified ensure every
unique value in getattr(intervals, stratify_by) has enough samples to be present
in each fold (i.e., count_per_class >= n_folds and also satisfies the train/val
minimum given val_ratio); raise clear ValueError messages referencing n_folds,
val_ratio, stratify_by, intervals, and class_labels when these conditions are
not met.
| folds2 = generate_folds(intervals, n_folds=5, val_ratio=0.2, seed=123) | ||
|
|
||
| assert len(folds1) == len(folds2) | ||
| for f1, f2 in zip(folds1, folds2): |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's look at the file and its context around line 173
cd /dev/null || true # ensure we're in the repo
wc -l tests/test_split_utils.pyRepository: neuro-galaxy/brainsets
Length of output: 162
🏁 Script executed:
# Read the relevant section of the test file
sed -n '165,185p' tests/test_split_utils.pyRepository: neuro-galaxy/brainsets
Length of output: 919
🏁 Script executed:
# Let's also check what Ruff rule B905 is about
# First, let's see if there's any documentation or comments in the repo about Ruff
fd -t f -e md -e txt -e py | xargs grep -l "B905" 2>/dev/null | head -5Repository: neuro-galaxy/brainsets
Length of output: 48
🏁 Script executed:
# Check Python version requirement for strict=True in zip()
# strict=True was added in Python 3.10
# Let's check what Python version this project targets
fd -t f \( -name "pyproject.toml" -o -name "setup.py" -o -name ".python-version" -o -name "setup.cfg" \) -exec cat {} \;Repository: neuro-galaxy/brainsets
Length of output: 296
🏁 Script executed:
# Look for project configuration files
find . -maxdepth 2 -type f \( -name "pyproject.toml" -o -name "setup.py" -o -name "setup.cfg" \) | head -3Repository: neuro-galaxy/brainsets
Length of output: 82
🏁 Script executed:
# Check the pyproject.toml if it exists
cat pyproject.toml 2>/dev/null || echo "File not found"Repository: neuro-galaxy/brainsets
Length of output: 1793
🌐 Web query:
Ruff B905 rule definition what does it check
💡 Result:
Ruff rule B905 (zip-without-explicit-strict) flags calls to zip() that don’t pass an explicit strict= argument when zip is called with (a) two or more iterables or (b) any starred argument (e.g., *iterables). It’s meant to prevent silent truncation when iterables have different lengths by forcing you to choose strict=True (raise ValueError on length mismatch) or strict=False (explicitly allow truncation). [1]
Example flagged: zip(a, b) → preferred: zip(a, b, strict=True) (or strict=False if intentional). [1]
Source: Ruff documentation for B905. [1]
References
[1] Ruff rule docs: “zip-without-explicit-strict (B905)” (docs.astral.sh)
Citations:
Add strict=True to zip() to address Ruff B905.
The zip(folds1, folds2, strict=True) will explicitly enforce that both iterables have equal length, making the intent clear and satisfying the zip-without-explicit-strict linter rule. Although the preceding assertion already guards this, being explicit about the strict parameter is preferred.
Suggested change
- for f1, f2 in zip(folds1, folds2):
+ for f1, f2 in zip(folds1, folds2, strict=True):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for f1, f2 in zip(folds1, folds2): | |
| for f1, f2 in zip(folds1, folds2, strict=True): |
🧰 Tools
🪛 Ruff (0.15.2)
[warning] 173-173: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_split_utils.py` at line 173, Update the zip call in the test loop
so it explicitly enforces equal-length iterables: change the loop that iterates
"for f1, f2 in zip(folds1, folds2):" to pass strict=True (i.e., zip(folds1,
folds2, strict=True)) so the linter Ruff B905 is satisfied and the intent to
require equal lengths for folds1 and folds2 is explicit.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
vinamarora8
left a comment
vinamarora8
left a comment
There was a problem hiding this comment.
Hi @oumgharbi, I like the idea of having a generate_folds without stratification. However, I see that the logic between the two is not quite reusable, and in your code you had to basically have an if/else to choose between stratified logic or non-stratified logic.
In this case, I think we might as well create a new function generate_folds, and keep the old generate_stratified_folds as it is.
2 participants
Adressing issue #94.
Small PR to modify the util to generate folds, making stratification optional and used when a label is provided.
Otherwise, random K-folds are generated without stratification.
generate_stratified_foldsis updated asgenerate_folds, with thestratify_byargument used to differentiate logics.Summary by CodeRabbit
Release Notes
New Features
API Changes