Fix/issue 397 stratified splitting

[intagliated]

PR: Label-Aware Site Stratification for Training Splits

Linked Issue

Closes #397

Summary

This PR introduces Label-Aware Site Stratification to ensure that rare species (e.g., the "zebra" case) are consistently represented across all data splits (Train, Validation, and Holdout).

Previously, zamba utilized a standard group-based split which prioritized site integrity but ignored label distribution. This resulted in rare classes being dropped from evaluation sets if their corresponding sites were randomly assigned to the training set. This fix implements a Greedy Iterative Stratification algorithm to balance these constraints.

Key Changes

• zamba/data/operations.py: Added generate_stratified_group_split. This core function uses a greedy approach to distribute sites based on their label composition.
• zamba/models/model_manager.py: Updated the training entry point to automatically trigger this stratified approach when a site column is present and no manual split is provided.
• zamba/models/config.py: Updated TrainConfig to ensure seamless integration with existing data-loading parameters.

Implementation Details

The algorithm functions as a multi-objective optimization:

1. Rarity Ranking: Labels are sorted by global frequency, ensuring "at-risk" species are assigned first.
2. Greedy Assignment: Unassigned sites containing rare labels are evaluated for each split.
3. Variance Minimization: Sites are assigned to the split that has the largest negative gap relative to the target ratio ($70/15/15$ by default).

Note

This approach strictly respects the site group constraint. 100% of samples from a single site will always remain in the same split to prevent data leakage.

Verification Results

• Unit Test: Verified via tests/test_issue_397_stratification.py using a skewed dummy dataset.
• CI/CD: Added a GitHub Action workflow to validate the stratification logic across different environments.
• Integrity: Confirmed that no site-level leakage occurs.

Comparative Performance

Feature	Random Site Split (Old)	Stratified Site Split (New)
Site Integrity	✅	✅
Common Label Distribution	✅	✅
Rare Label Distribution	❌	✅