metacausal.aggregation.CrossFitSplit¶

class metacausal.aggregation.CrossFitSplit(n_folds=5, stratify=True)[source]¶

Bases: object

Q-fold cross-fitting split.

Every observation gets exactly one out-of-fold (OOF) prediction. Wraps sklearn’s StratifiedKFold (when stratify=True) or KFold.

Parameters:

n_folds (int) – Number of folds (Q). Must be >= 2.
stratify (bool) – If True, stratify folds on treatment assignment T. Requires T to be discrete (binary or categorical). If T appears continuous (more than 10 unique values), a warning is issued and stratification is skipped.

Examples

>>> import numpy as np
>>> from metacausal.aggregation import CrossFitSplit
>>> T = np.array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
>>> fold_spec = CrossFitSplit(n_folds=5).split(T, random_state=0)
>>> fold_spec.n_folds
5
>>> np.sort(np.concatenate(fold_spec.test_indices))  # every index appears once
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

Methods

`__init__`
`split`	Partition n observations into Q folds.

Attributes

`n_folds`
`stratify`

Details

split(T, random_state=None)[source]¶

Partition n observations into Q folds.

Parameters:

T (array of shape (n,)) – Treatment assignments. Used for stratification when stratify=True.
random_state (int or None) – Random seed for reproducibility. Lives on the call, not the object, so the same CrossFitSplit can be reused with different seeds.

Returns:

FoldSpec with n_folds folds. Every index in range(n) appears in
exactly one test fold.

Return type:

FoldSpec

n_folds: int = 5¶

stratify: bool = True¶