ml4t-diagnostic 0.1.0a1__py3-none-any.whl

ml4t/diagnostic/splitters/cpcv/partitioning.py

@@ -0,0 +1,263 @@
+"""Group partitioning strategies for CPCV.
+
+This module handles partitioning the timeline into groups:
+- Contiguous partitioning (equal-sized time slices)
+- Session-aligned partitioning (respects trading session boundaries)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+if TYPE_CHECKING:
+    import pandas as pd
+    import polars as pl
+
+
+def create_contiguous_partitions(
+    n_samples: int,
+    n_groups: int,
+) -> list[tuple[int, int]]:
+    """Create boundaries for contiguous groups.
+
+    Partitions n_samples into n_groups approximately equal-sized groups.
+    Earlier groups get extra samples when n_samples is not evenly divisible.
+
+    Parameters
+    ----------
+    n_samples : int
+        Total number of samples.
+    n_groups : int
+        Number of groups to create.
+
+    Returns
+    -------
+    boundaries : list of tuple
+        List of (start_idx, end_idx) for each group.
+        end_idx is exclusive (standard Python convention).
+
+    Raises
+    ------
+    ValueError
+        If boundaries don't satisfy CPCV invariants.
+
+    Examples
+    --------
+    >>> create_contiguous_partitions(100, 5)
+    [(0, 20), (20, 40), (40, 60), (60, 80), (80, 100)]
+
+    >>> create_contiguous_partitions(103, 5)
+    [(0, 21), (21, 42), (42, 62), (62, 82), (82, 103)]
+    """
+    base_size = n_samples // n_groups
+    remainder = n_samples % n_groups
+
+    boundaries = []
+    current_start = 0
+
+    for i in range(n_groups):
+        # Add extra sample to first 'remainder' groups
+        group_size = base_size + (1 if i < remainder else 0)
+        group_end = current_start + group_size
+
+        boundaries.append((current_start, group_end))
+        current_start = group_end
+
+    # Validate invariants
+    validate_contiguous_partitions(boundaries, n_samples)
+
+    return boundaries
+
+
+def validate_contiguous_partitions(
+    boundaries: list[tuple[int, int]],
+    n_samples: int,
+) -> None:
+    """Validate CPCV group boundary invariants.
+
+    Ensures:
+    1. All samples are covered (no gaps)
+    2. No overlap between groups
+    3. Groups are contiguous
+
+    Parameters
+    ----------
+    boundaries : list of tuple
+        List of (start_idx, end_idx) for each group.
+    n_samples : int
+        Total number of samples.
+
+    Raises
+    ------
+    ValueError
+        If any invariant is violated.
+    """
+    if not boundaries:
+        raise ValueError("CPCV invariant violated: no group boundaries created")
+
+    # Check first boundary starts at 0
+    if boundaries[0][0] != 0:
+        raise ValueError(
+            f"CPCV invariant violated: first group must start at 0, got {boundaries[0][0]}"
+        )
+
+    # Check last boundary ends at n_samples
+    if boundaries[-1][1] != n_samples:
+        raise ValueError(
+            f"CPCV invariant violated: last group must end at {n_samples}, got {boundaries[-1][1]}"
+        )
+
+    # Check contiguity (each group starts where previous ended)
+    for i in range(1, len(boundaries)):
+        prev_end = boundaries[i - 1][1]
+        curr_start = boundaries[i][0]
+        if curr_start != prev_end:
+            raise ValueError(
+                f"CPCV invariant violated: gap between group {i - 1} (ends at {prev_end}) "
+                f"and group {i} (starts at {curr_start})"
+            )
+
+    # Check each group is non-empty
+    for i, (start, end) in enumerate(boundaries):
+        if end <= start:
+            raise ValueError(
+                f"CPCV invariant violated: group {i} is empty or invalid (start={start}, end={end})"
+            )
+
+
+def create_session_partitions(
+    X: pl.DataFrame | pd.DataFrame,
+    session_col: str,
+    n_groups: int,
+    session_to_indices_fn: Callable[
+        [pl.DataFrame | pd.DataFrame, str],
+        tuple[list[Any], dict[Any, NDArray[np.intp]]],
+    ],
+) -> list[NDArray[np.intp]]:
+    """Create exact index arrays per group, aligned to session boundaries.
+
+    Unlike contiguous partitioning which returns (start, end) ranges,
+    this method returns EXACT index arrays for each group. This is critical
+    for correct behavior with non-contiguous or interleaved data.
+
+    Parameters
+    ----------
+    X : DataFrame
+        Data with session column.
+    session_col : str
+        Name of column containing session identifiers.
+    n_groups : int
+        Number of groups to create.
+    session_to_indices_fn : callable
+        Function that returns (ordered_sessions, session_to_indices_dict).
+        Typically from BaseSplitter._session_to_indices.
+
+    Returns
+    -------
+    group_indices : list of np.ndarray
+        List of numpy arrays containing exact row indices for each group.
+        Each array contains the indices for all rows belonging to sessions
+        in that group.
+
+    Raises
+    ------
+    ValueError
+        If not enough sessions for the requested number of groups.
+
+    Notes
+    -----
+    The key difference from contiguous partitioning is that we track
+    exact indices rather than (start, end) boundaries. This prevents
+    incorrect index ranges when data is interleaved by asset within sessions.
+    """
+    # Get session -> indices mapping
+    ordered_sessions, session_to_indices = session_to_indices_fn(X, session_col)
+    n_sessions = len(ordered_sessions)
+
+    if n_sessions < n_groups:
+        raise ValueError(
+            f"Not enough sessions ({n_sessions}) for {n_groups} groups. "
+            f"Need at least {n_groups} sessions."
+        )
+
+    # Partition sessions into groups
+    base_sessions_per_group = n_sessions // n_groups
+    remainder = n_sessions % n_groups
+
+    group_indices_list = []
+    current_session_idx = 0
+
+    for i in range(n_groups):
+        # Add extra session to first 'remainder' groups
+        sessions_in_group = base_sessions_per_group + (1 if i < remainder else 0)
+        session_group_end = current_session_idx + sessions_in_group
+
+        # Get sessions for this group
+        group_sessions = ordered_sessions[current_session_idx:session_group_end]
+
+        # Collect EXACT indices for sessions in this group
+        indices_arrays = [session_to_indices[s] for s in group_sessions]
+        if indices_arrays:
+            group_indices = np.concatenate(indices_arrays)
+            # Sort for predictable ordering
+            group_indices = np.sort(group_indices)
+        else:
+            group_indices = np.array([], dtype=np.intp)
+
+        group_indices_list.append(group_indices)
+        current_session_idx = session_group_end
+
+    return group_indices_list
+
+
+def boundaries_to_indices(
+    boundaries: list[tuple[int, int]],
+    groups: tuple[int, ...],
+) -> NDArray[np.intp]:
+    """Convert group boundaries to flat index array for selected groups.
+
+    Parameters
+    ----------
+    boundaries : list of tuple
+        List of (start_idx, end_idx) for each group.
+    groups : tuple of int
+        Which groups to include.
+
+    Returns
+    -------
+    indices : np.ndarray
+        Sorted array of indices for selected groups.
+    """
+    # Use numpy concatenation instead of Python list extend for performance
+    ranges = [np.arange(boundaries[g][0], boundaries[g][1], dtype=np.intp) for g in groups]
+    if not ranges:
+        return np.array([], dtype=np.intp)
+    return np.concatenate(ranges)
+
+
+def exact_indices_to_array(
+    group_indices_list: list[NDArray[np.intp]],
+    groups: tuple[int, ...],
+) -> NDArray[np.intp]:
+    """Concatenate exact index arrays for selected groups.
+
+    Parameters
+    ----------
+    group_indices_list : list of np.ndarray
+        List of exact index arrays for each group.
+    groups : tuple of int
+        Which groups to include.
+
+    Returns
+    -------
+    indices : np.ndarray
+        Sorted array of indices for selected groups.
+    """
+    arrays = [group_indices_list[g] for g in groups]
+    if not arrays or all(len(a) == 0 for a in arrays):
+        return np.array([], dtype=np.intp)
+    return np.sort(np.concatenate(arrays))

ml4t/diagnostic/splitters/cpcv/purge_engine.py

@@ -0,0 +1,379 @@
+"""Purging engine for CPCV.
+
+This module implements the core purging and embargo logic:
+- Mask-based purging (efficient for large datasets)
+- Single-asset and multi-asset purging strategies
+- Segment-based purging for temporal coherence
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ml4t.diagnostic.core.purging import apply_purging_and_embargo
+from ml4t.diagnostic.splitters.cpcv.windows import (
+    find_contiguous_segments,
+    timestamp_window_from_indices,
+)
+from ml4t.diagnostic.splitters.utils import convert_indices_to_timestamps
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+def apply_single_asset_purging(
+    train_indices: NDArray[np.intp],
+    test_group_indices: tuple[int, ...],
+    group_boundaries: list[tuple[int, int]],
+    n_samples: int,
+    timestamps: pd.DatetimeIndex | None,
+    label_horizon: int | pd.Timedelta,
+    embargo_size: int | pd.Timedelta | None,
+    embargo_pct: float | None,
+    group_indices_list: list[NDArray[np.intp]] | None = None,
+) -> NDArray[np.intp]:
+    """Apply purging for single-asset data.
+
+    For each test group, removes training samples that would cause
+    look-ahead bias due to label overlap or temporal proximity.
+
+    Parameters
+    ----------
+    train_indices : ndarray
+        Initial training indices.
+    test_group_indices : tuple of int
+        Indices of groups used for testing.
+    group_boundaries : list of tuple
+        Boundaries (start, end) for each group.
+    n_samples : int
+        Total number of samples.
+    timestamps : pd.DatetimeIndex, optional
+        Timestamps for time-based purging.
+    label_horizon : int or pd.Timedelta
+        Forward-looking period of labels.
+    embargo_size : int or pd.Timedelta, optional
+        Buffer period after test set.
+    embargo_pct : float, optional
+        Embargo as percentage of samples.
+    group_indices_list : list of ndarray, optional
+        Exact indices per group (for session-aligned mode).
+
+    Returns
+    -------
+    clean_indices : ndarray
+        Training indices after purging.
+    """
+    for test_group_idx in test_group_indices:
+        # Compute purge window bounds
+        if group_indices_list is not None and timestamps is not None:
+            # Session-aligned mode: use actual timestamps from test indices
+            test_indices = group_indices_list[test_group_idx]
+            window = timestamp_window_from_indices(test_indices, timestamps)
+            if window is None:
+                # Empty test group - skip purging for this group
+                continue
+            test_start_time = window.start
+            test_end_time = window.end_exclusive
+        else:
+            # Standard mode: use boundaries
+            test_start_idx, test_end_idx = group_boundaries[test_group_idx]
+            test_start_time, test_end_time = convert_indices_to_timestamps(
+                test_start_idx,
+                test_end_idx,
+                timestamps,
+            )
+
+        # Apply purging and embargo for this test group
+        train_indices = apply_purging_and_embargo(
+            train_indices=train_indices,
+            test_start=test_start_time,
+            test_end=test_end_time,
+            label_horizon=label_horizon,
+            embargo_size=embargo_size,
+            embargo_pct=embargo_pct,
+            n_samples=n_samples,
+            timestamps=timestamps,
+        )
+
+    return train_indices
+
+
+def apply_multi_asset_purging(
+    train_indices: NDArray[np.intp],
+    test_group_indices: tuple[int, ...],
+    group_boundaries: list[tuple[int, int]],
+    n_samples: int,
+    timestamps: pd.DatetimeIndex | None,
+    groups_array: NDArray[Any],
+    label_horizon: int | pd.Timedelta,
+    embargo_size: int | pd.Timedelta | None,
+    embargo_pct: float | None,
+    group_indices_list: list[NDArray[np.intp]] | None = None,
+) -> NDArray[np.intp]:
+    """Apply purging for multi-asset data with per-asset isolation.
+
+    This method correctly handles non-contiguous test groups by applying
+    purging for each contiguous segment of test data separately per asset.
+
+    Parameters
+    ----------
+    train_indices : ndarray
+        Initial training indices.
+    test_group_indices : tuple of int
+        Indices of groups used for testing.
+    group_boundaries : list of tuple
+        Boundaries (start, end) for each group.
+    n_samples : int
+        Total number of samples.
+    timestamps : pd.DatetimeIndex, optional
+        Timestamps for time-based purging.
+    groups_array : ndarray
+        Asset labels for each sample.
+    label_horizon : int or pd.Timedelta
+        Forward-looking period of labels.
+    embargo_size : int or pd.Timedelta, optional
+        Buffer period after test set.
+    embargo_pct : float, optional
+        Embargo as percentage of samples.
+    group_indices_list : list of ndarray, optional
+        Exact indices per group (for session-aligned mode).
+
+    Returns
+    -------
+    clean_indices : ndarray
+        Training indices after per-asset purging.
+    """
+    if len(groups_array) != n_samples:
+        raise ValueError(
+            f"groups length ({len(groups_array)}) must match number of samples ({n_samples})",
+        )
+
+    # Prepare test groups data for contiguous segment detection
+    test_groups_data = prepare_test_groups_data(
+        test_group_indices, group_boundaries, group_indices_list
+    )
+
+    # Apply purging per asset
+    final_train_indices: list[int] = []
+    unique_assets = np.unique(groups_array)
+
+    for asset_id in unique_assets:
+        # Process this asset's training data with purging
+        asset_train = process_asset_purging(
+            asset_id=asset_id,
+            groups_array=groups_array,
+            train_indices=train_indices,
+            test_groups_data=test_groups_data,
+            n_samples=n_samples,
+            timestamps=timestamps,
+            label_horizon=label_horizon,
+            embargo_size=embargo_size,
+            embargo_pct=embargo_pct,
+            group_indices_list=group_indices_list,
+        )
+        final_train_indices.extend(asset_train)
+
+    # Sort for deterministic output
+    return np.sort(np.array(final_train_indices, dtype=np.intp))
+
+
+def prepare_test_groups_data(
+    test_group_indices: tuple[int, ...],
+    group_boundaries: list[tuple[int, int]],
+    group_indices_list: list[NDArray[np.intp]] | None = None,
+) -> list[tuple[int, int, int, NDArray[np.intp] | None]]:
+    """Prepare and sort test groups data for contiguous segment detection.
+
+    Parameters
+    ----------
+    test_group_indices : tuple of int
+        Which groups are used for testing.
+    group_boundaries : list of tuple
+        Boundaries (start, end) for each group.
+    group_indices_list : list of ndarray, optional
+        Exact indices per group (for session-aligned mode).
+
+    Returns
+    -------
+    test_groups_data : list of tuple
+        Sorted list of (group_idx, start_idx, end_idx, exact_indices).
+        In session-aligned mode, exact_indices contains the actual row indices;
+        otherwise it's None.
+    """
+    test_groups_data: list[tuple[int, int, int, NDArray[np.intp] | None]] = []
+    for test_group_idx in test_group_indices:
+        test_start_idx, test_end_idx = group_boundaries[test_group_idx]
+        exact_indices = (
+            group_indices_list[test_group_idx] if group_indices_list is not None else None
+        )
+        test_groups_data.append((test_group_idx, test_start_idx, test_end_idx, exact_indices))
+
+    # Sort test groups by start index to identify contiguous segments
+    test_groups_data.sort(key=lambda x: x[1])
+    return test_groups_data
+
+
+def process_asset_purging(
+    asset_id: Any,
+    groups_array: NDArray[Any],
+    train_indices: NDArray[np.intp],
+    test_groups_data: list[tuple[int, int, int, NDArray[np.intp] | None]],
+    n_samples: int,
+    timestamps: pd.DatetimeIndex | None,
+    label_horizon: int | pd.Timedelta,
+    embargo_size: int | pd.Timedelta | None,
+    embargo_pct: float | None,
+    group_indices_list: list[NDArray[np.intp]] | None = None,
+) -> list[int]:
+    """Process purging for a single asset across all test segments.
+
+    Parameters
+    ----------
+    asset_id : any
+        Identifier for this asset.
+    groups_array : ndarray
+        Asset labels for all samples.
+    train_indices : ndarray
+        Candidate training indices.
+    test_groups_data : list of tuple
+        Test group information from prepare_test_groups_data.
+    n_samples : int
+        Total number of samples.
+    timestamps : pd.DatetimeIndex, optional
+        Timestamps for time-based purging.
+    label_horizon : int or pd.Timedelta
+        Forward-looking period of labels.
+    embargo_size : int or pd.Timedelta, optional
+        Buffer period after test set.
+    embargo_pct : float, optional
+        Embargo as percentage of samples.
+    group_indices_list : list of ndarray, optional
+        Exact indices per group (for session-aligned mode).
+
+    Returns
+    -------
+    clean_indices : list of int
+        Training indices for this asset after purging.
+    """
+    # Find indices for this asset
+    asset_mask = groups_array == asset_id
+    asset_indices = np.where(asset_mask)[0]
+
+    # Get train indices for this asset
+    asset_train_indices = np.intersect1d(train_indices, asset_indices)
+
+    if len(asset_train_indices) == 0:
+        return []
+
+    # Find contiguous segments of test groups for this asset
+    contiguous_segments = find_contiguous_segments(
+        test_groups_data,
+        asset_indices,
+    )
+
+    # If no test data for this asset, keep all training data
+    if not contiguous_segments:
+        return asset_train_indices.tolist()
+
+    # Apply purging for each contiguous segment
+    return apply_segment_purging(
+        asset_train_indices=asset_train_indices,
+        contiguous_segments=contiguous_segments,
+        n_samples=n_samples,
+        timestamps=timestamps,
+        label_horizon=label_horizon,
+        embargo_size=embargo_size,
+        embargo_pct=embargo_pct,
+        group_indices_list=group_indices_list,
+    )
+
+
+def apply_segment_purging(
+    asset_train_indices: NDArray[np.intp],
+    contiguous_segments: list[list[tuple[int, int, int, NDArray[np.intp]]]],
+    n_samples: int,
+    timestamps: pd.DatetimeIndex | None,
+    label_horizon: int | pd.Timedelta,
+    embargo_size: int | pd.Timedelta | None,
+    embargo_pct: float | None,
+    group_indices_list: list[NDArray[np.intp]] | None = None,
+) -> list[int]:
+    """Apply purging across all contiguous segments for an asset.
+
+    Uses a set-based approach for tracking remaining indices, which is
+    efficient for the iterative purging across segments.
+
+    Parameters
+    ----------
+    asset_train_indices : ndarray
+        Training indices for this asset.
+    contiguous_segments : list of list of tuple
+        Segments from find_contiguous_segments.
+    n_samples : int
+        Total number of samples.
+    timestamps : pd.DatetimeIndex, optional
+        Timestamps for time-based purging.
+    label_horizon : int or pd.Timedelta
+        Forward-looking period of labels.
+    embargo_size : int or pd.Timedelta, optional
+        Buffer period after test set.
+    embargo_pct : float, optional
+        Embargo as percentage of samples.
+    group_indices_list : list of ndarray, optional
+        Exact indices per group (for session-aligned mode).
+
+    Returns
+    -------
+    clean_indices : list of int
+        Sorted training indices after purging all segments.
+    """
+    remaining_train_indices = set(asset_train_indices)
+
+    for segment in contiguous_segments:
+        if not segment:
+            continue
+
+        # Compute purge window bounds
+        if group_indices_list is not None and timestamps is not None:
+            # Session-aligned mode: compute timestamp bounds from actual test indices
+            segment_test_indices = np.concatenate([item[3] for item in segment])
+            window = timestamp_window_from_indices(segment_test_indices, timestamps)
+            if window is None:
+                # Empty test segment - skip purging for this segment
+                continue
+            segment_start_time = window.start
+            segment_end_time = window.end_exclusive
+        else:
+            # Standard mode: use boundaries
+            segment_start_idx = segment[0][1]  # Start of first group in segment
+            segment_end_idx = segment[-1][2]  # End of last group in segment
+            segment_start_time, segment_end_time = convert_indices_to_timestamps(
+                segment_start_idx,
+                segment_end_idx,
+                timestamps,
+            )
+
+        # Apply purging for this contiguous segment
+        remaining_array = np.array(list(remaining_train_indices), dtype=np.intp)
+
+        if len(remaining_array) == 0:
+            break
+
+        clean_segment_train = apply_purging_and_embargo(
+            train_indices=remaining_array,
+            test_start=segment_start_time,
+            test_end=segment_end_time,
+            label_horizon=label_horizon,
+            embargo_size=embargo_size,
+            embargo_pct=embargo_pct,
+            n_samples=n_samples,
+            timestamps=timestamps,
+        )
+
+        # Update remaining indices (remove those that were purged)
+        remaining_train_indices = set(clean_segment_train)
+
+    return sorted(remaining_train_indices)