added some tests

Author: antonkulaga

docs/DONOR_LEVEL_SPLIT.md

@@ -0,0 +1,211 @@
+# Donor-Level Train/Test Split Implementation
+
+## Overview
+
+Modified the preprocessing pipeline to implement **donor-level train/test splitting** instead of cell-level splitting. This prevents data leakage of subject-specific signatures that could cause models to learn how to identify subjects rather than predict age.
+
+## Problem Statement
+
+Previously, the train/test split was performed at the **cell level** with age stratification. This meant that cells from the same donor could appear in both train and test sets. This creates **data leakage** because:
+
+1. Each donor has unique biological signatures (genetics, environment, lifestyle, etc.)
+2. The model can learn these donor-specific patterns
+3. During evaluation, the model might recognize cells from the same donor seen during training
+4. This leads to inflated performance metrics that don't reflect true generalization
+
+## Solution
+
+The new implementation performs **donor-level stratified splitting**:
+
+1. **Group cells by donor** (using `donor_id`, `donor`, `subject_id`, or `individual` columns)
+2. **Compute representative age** for each donor (median age of their cells)
+3. **Split donors** (not cells) into train/test using stratified sampling by donor age
+4. **Assign all cells** from each donor to the same split (train OR test, never both)
+
+### Key Features
+
+- **Donor-level stratification**: Maintains age distribution at the donor level
+- **Automatic fallback**: Falls back to cell-level split if no donor column is found
+- **Multiple donor column support**: Tries `donor_id`, `donor`, `subject_id`, `individual`
+- **Comprehensive logging**: Logs donor statistics and split quality
+- **Memory efficient**: Uses Polars lazy API and streaming where possible
+
+## Files Modified
+
+### 1. `src/cell2sentence4longevity/preprocessing/h5ad_converter.py`
+
+**Function**: `convert_h5ad_to_train_test()`
+
+**Changes**:
+- Added donor column detection logic (lines 1342-1343)
+- Implemented donor-level stratified split (lines 1383-1442)
+- Computes median age per donor for stratification (lines 1395-1402)
+- Uses sklearn stratified split on donors, then joins back to cells (lines 1412-1432)
+- Falls back to cell-level split if no donor column found (lines 1345-1382)
+- Added comprehensive logging for donor statistics
+
+**Key Logic**:
+```python
+# Identify donor column
+donor_cols = ['donor_id', 'donor', 'subject_id', 'individual']
+donor_col = next((col for col in donor_cols if col in chunk_df.columns), None)
+
+if donor_col:
+    # Donor-level split
+    donor_ages = chunk_df.group_by(donor_col).agg([
+        pl.col('age').median().alias('donor_age')
+    ])
+    
+    # Stratify donors by age
+    train_donors = stratified_split_of_donors(...)
+    test_donors = ...
+    
+    # Join back to get all cells for each donor
+    train_chunk = chunk_df.join(train_donors, on=donor_col, how='inner')
+    test_chunk = chunk_df.join(test_donors, on=donor_col, how='inner')
+```
+
+### 2. `src/cell2sentence4longevity/preprocessing/train_test_split.py`
+
+**Function**: `create_train_test_split()`
+
+**Changes**:
+- Updated docstring to reflect donor-level splitting (lines 25-29)
+- Added donor column detection for lazy datasets (lines 168-203)
+- Implemented donor-level stratified split using sklearn (lines 212-272)
+- Computes donor statistics (total donors, cells per donor) (lines 189-203)
+- Uses sklearn's `train_test_split` with donor age stratification (lines 240-265)
+- Filters lazy dataset using `is_in` with donor lists (lines 268-272)
+- Falls back to cell-level split if no donor column (lines 274-287)
+
+**Key Logic**:
+```python
+# Compute representative age for each donor
+donor_ages = (
+    lazy_dataset
+    .group_by(donor_col)
+    .agg([
+        pl.col('age').median().alias('donor_age'),
+        pl.len().alias('cell_count')
+    ])
+).collect()
+
+# Stratified split of donors by age
+train_donor_ids, test_donor_ids = sklearn_split(
+    donor_ids_array,
+    test_size=test_size,
+    random_state=random_state,
+    stratify=np.round(donor_ages_array, 1)
+)
+
+# Filter cells by donor assignment
+lazy_train = lazy_dataset.filter(pl.col(donor_col).is_in(train_donor_series))
+lazy_test = lazy_dataset.filter(pl.col(donor_col).is_in(test_donor_series))
+```
+
+### 3. `tests/test_donor_split.py` (New File)
+
+**Purpose**: Comprehensive tests for donor-level split validation
+
+**Tests**:
+1. **`test_donor_split_no_leakage()`**: Verifies no donor appears in both train and test
+   - Creates synthetic data with 10 donors, 50 cells each
+   - Runs donor-level split
+   - Validates no overlap between train and test donors
+   - Verifies all cells from same donor are in same split
+
+2. **`test_fallback_to_cell_level_split_without_donor_column()`**: Verifies fallback behavior
+   - Creates data without donor column
+   - Runs split
+   - Validates cell-level split works correctly
+
+## Benefits
+
+1. **Prevents data leakage**: No donor-specific signatures can be learned and exploited
+2. **True generalization**: Model must learn age-related patterns, not donor identification
+3. **Better evaluation**: Test performance reflects true ability to predict age for unseen subjects
+4. **Maintains age distribution**: Stratification ensures similar age distributions in train/test
+5. **Robust fallback**: Automatically handles datasets without donor information
+6. **Production ready**: Comprehensive logging and error handling
+
+## Usage
+
+No changes required to existing code! The donor-level split is **automatically applied** when:
+- A donor column exists (`donor_id`, `donor`, `subject_id`, or `individual`)
+- Age stratification is enabled (default: `stratify_by_age=True`)
+
+If no donor column is found, it gracefully falls back to cell-level split with a warning.
+
+### Example
+
+```python
+from cell2sentence4longevity.preprocessing import convert_h5ad_to_train_test
+
+# Automatically uses donor-level split if donor column exists
+convert_h5ad_to_train_test(
+    h5ad_path=Path("data.h5ad"),
+    output_dir=Path("output"),
+    test_size=0.05,
+    stratify_by_age=True,  # Stratifies at donor level!
+    random_state=42
+)
+```
+
+## Validation
+
+Run the tests to verify donor-level split works correctly:
+
+```bash
+uv run pytest tests/test_donor_split.py -v
+```
+
+Expected output:
+```
+✓ Donor-level split validation passed:
+  - Train donors: 8
+  - Test donors: 2
+  - Train cells: 400
+  - Test cells: 100
+  - No donor leakage detected!
+```
+
+## Logging
+
+The implementation logs comprehensive donor statistics:
+
+```json
+{
+  "message_type": "using_donor_level_split",
+  "donor_col": "donor_id",
+  "unique_donors_in_chunk": 47
+}
+
+{
+  "message_type": "donor_statistics",
+  "total_donors": 150,
+  "total_cells": 45000,
+  "avg_cells_per_donor": 300.0
+}
+
+{
+  "message_type": "donor_split_complete",
+  "train_donors": 142,
+  "test_donors": 8
+}
+```
+
+## Future Improvements
+
+Potential enhancements (not implemented yet):
+
+1. **Multiple stratification factors**: Stratify by both age and tissue type
+2. **Minimum donor requirements**: Ensure test set has minimum number of donors per age bin
+3. **Donor metadata export**: Export donor-level statistics for analysis
+4. **Cross-validation support**: K-fold donor-level cross-validation
+
+## References
+
+- Original issue: "cells from the same donor do not appear in both train and test"
+- Related concept: Group-aware splitting in machine learning
+- Similar to: Patient-level splitting in medical ML, speaker-level splitting in audio ML
+

tests/test_donor_split.py

@@ -0,0 +1,175 @@
+"""Tests for donor-level train/test split to prevent data leakage."""
+
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+import pytest
+
+
+def test_donor_split_no_leakage():
+    """Test that donor-level split ensures no donor appears in both train and test."""
+    
+    # Create synthetic test data with multiple donors and cells
+    np.random.seed(42)
+    n_donors = 10
+    cells_per_donor = 50
+    
+    # Generate donor data
+    donors = [f"donor_{i}" for i in range(n_donors)]
+    donor_ages = np.random.uniform(20, 80, n_donors)
+    
+    # Generate cell data
+    data = []
+    for donor_id, donor_age in zip(donors, donor_ages):
+        for _ in range(cells_per_donor):
+            # Add some variation in cell age (but centered on donor age)
+            cell_age = donor_age + np.random.normal(0, 0.1)
+            data.append({
+                'donor_id': donor_id,
+                'age': cell_age,
+                'gene_sentence_2000': 'gene1 gene2 gene3',
+                'cell_type': f'type_{np.random.randint(1, 5)}'
+            })
+    
+    df = pl.DataFrame(data)
+    
+    # Write to temporary parquet files (simulating input chunks)
+    with tempfile.TemporaryDirectory() as tmpdir:
+        chunks_dir = Path(tmpdir) / 'chunks'
+        chunks_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Write in chunks
+        chunk_size = 100
+        n_chunks = (len(df) + chunk_size - 1) // chunk_size
+        for i in range(n_chunks):
+            start_idx = i * chunk_size
+            end_idx = min(start_idx + chunk_size, len(df))
+            chunk = df.slice(start_idx, end_idx - start_idx)
+            chunk.write_parquet(chunks_dir / f"chunk_{i:04d}.parquet")
+        
+        # Run the donor-level split using our modified function
+        from cell2sentence4longevity.preprocessing.train_test_split import create_train_test_split
+        
+        output_dir = Path(tmpdir) / 'output'
+        create_train_test_split(
+            parquet_dir=chunks_dir,
+            output_dir=output_dir,
+            dataset_name='test_dataset',
+            test_size=0.2,
+            random_state=42,
+            chunk_size=100,
+            compression='zstd',
+            compression_level=3,
+            use_pyarrow=True
+        )
+        
+        # Load train and test sets
+        train_dir = output_dir / 'test_dataset' / 'train'
+        test_dir = output_dir / 'test_dataset' / 'test'
+        
+        train_df = pl.scan_parquet(train_dir / '*.parquet').collect()
+        test_df = pl.scan_parquet(test_dir / '*.parquet').collect()
+        
+        # Verify no donor appears in both train and test
+        train_donors = set(train_df['donor_id'].unique().to_list())
+        test_donors = set(test_df['donor_id'].unique().to_list())
+        
+        overlap = train_donors & test_donors
+        assert len(overlap) == 0, f"Found {len(overlap)} donors in both train and test: {overlap}"
+        
+        # Verify all donors are accounted for
+        all_original_donors = set(df['donor_id'].unique().to_list())
+        all_split_donors = train_donors | test_donors
+        assert all_original_donors == all_split_donors, "Some donors were lost during split"
+        
+        # Verify test size is approximately correct at donor level
+        donor_test_ratio = len(test_donors) / (len(train_donors) + len(test_donors))
+        assert 0.1 <= donor_test_ratio <= 0.3, f"Test donor ratio {donor_test_ratio:.2f} is not close to target 0.2"
+        
+        # Verify all cells from same donor are in same split
+        for donor in all_original_donors:
+            donor_train_count = train_df.filter(pl.col('donor_id') == donor).height
+            donor_test_count = test_df.filter(pl.col('donor_id') == donor).height
+            assert (donor_train_count > 0 and donor_test_count == 0) or \
+                   (donor_train_count == 0 and donor_test_count > 0), \
+                   f"Donor {donor} has cells in both train ({donor_train_count}) and test ({donor_test_count})"
+        
+        print(f"✓ Donor-level split validation passed:")
+        print(f"  - Train donors: {len(train_donors)}")
+        print(f"  - Test donors: {len(test_donors)}")
+        print(f"  - Train cells: {train_df.height}")
+        print(f"  - Test cells: {test_df.height}")
+        print(f"  - No donor leakage detected!")
+
+
+def test_fallback_to_cell_level_split_without_donor_column():
+    """Test that the split falls back to cell-level when no donor column exists."""
+    
+    # Create synthetic test data WITHOUT donor column
+    np.random.seed(42)
+    n_cells = 500
+    
+    data = []
+    for i in range(n_cells):
+        data.append({
+            'age': np.random.uniform(20, 80),
+            'gene_sentence_2000': 'gene1 gene2 gene3',
+            'cell_type': f'type_{np.random.randint(1, 5)}'
+        })
+    
+    df = pl.DataFrame(data)
+    
+    # Write to temporary parquet files
+    with tempfile.TemporaryDirectory() as tmpdir:
+        chunks_dir = Path(tmpdir) / 'chunks'
+        chunks_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Write in chunks
+        chunk_size = 100
+        n_chunks = (len(df) + chunk_size - 1) // chunk_size
+        for i in range(n_chunks):
+            start_idx = i * chunk_size
+            end_idx = min(start_idx + chunk_size, len(df))
+            chunk = df.slice(start_idx, end_idx - start_idx)
+            chunk.write_parquet(chunks_dir / f"chunk_{i:04d}.parquet")
+        
+        # Run the split (should fall back to cell-level)
+        from cell2sentence4longevity.preprocessing.train_test_split import create_train_test_split
+        
+        output_dir = Path(tmpdir) / 'output'
+        create_train_test_split(
+            parquet_dir=chunks_dir,
+            output_dir=output_dir,
+            dataset_name='test_dataset',
+            test_size=0.2,
+            random_state=42,
+            chunk_size=100,
+            compression='zstd',
+            compression_level=3,
+            use_pyarrow=True
+        )
+        
+        # Load train and test sets
+        train_dir = output_dir / 'test_dataset' / 'train'
+        test_dir = output_dir / 'test_dataset' / 'test'
+        
+        train_df = pl.scan_parquet(train_dir / '*.parquet').collect()
+        test_df = pl.scan_parquet(test_dir / '*.parquet').collect()
+        
+        # Verify split happened (basic sanity check)
+        total_cells = train_df.height + test_df.height
+        test_ratio = test_df.height / total_cells
+        
+        assert 0.15 <= test_ratio <= 0.25, f"Test ratio {test_ratio:.2f} is not close to target 0.2"
+        
+        print(f"✓ Cell-level split validation passed (fallback when no donor column):")
+        print(f"  - Train cells: {train_df.height}")
+        print(f"  - Test cells: {test_df.height}")
+        print(f"  - Test ratio: {test_ratio:.2%}")
+
+
+if __name__ == '__main__':
+    pytest.main([__file__, '-v'])
+