Function apply_purge_and_embargo

crates/openquant/src/backtesting_engine.rs:515–571

(
    initial_train: &[usize],
    test_indices: &[usize],
    label_spans: &[(NaiveDateTime, NaiveDateTime)],
    pct_embargo: f64,
    n_samples: usize,
)

}

fn apply_purge_and_embargo(
    initial_train: &[usize],
    test_indices: &[usize],
    label_spans: &[(NaiveDateTime, NaiveDateTime)],
    pct_embargo: f64,
    n_samples: usize,
) -> (Vec<usize>, usize, usize) {
    let mut train_mask = vec![false; n_samples];
    let mut initial_train_mask = vec![false; n_samples];
    for idx in initial_train {
        train_mask[*idx] = true;
        initial_train_mask[*idx] = true;
    }

    let mut purged_count = 0;
    for idx in initial_train {
        let mut should_purge = false;
        for test_idx in test_indices {
            if overlaps(label_spans[*idx], label_spans[*test_idx]) {
                should_purge = true;
                break;
            }
        }
        if should_purge {
            train_mask[*idx] = false;
            purged_count += 1;
        }
    }

    let embargo_width = (pct_embargo * n_samples as f64).ceil() as usize;
    let mut embargoed = vec![false; n_samples];
    if embargo_width > 0 {
        for test_idx in test_indices {
            let start = test_idx.saturating_sub(embargo_width);
            let stop = (*test_idx + embargo_width + 1).min(n_samples);
            for idx in start..stop {
                if initial_train_mask[idx] {
                    embargoed[idx] = true;
                }
            }
        }
        for idx in 0..n_samples {
            if embargoed[idx] {
                train_mask[idx] = false;
            }
        }
    }

    let train_indices: Vec<usize> = train_mask
        .iter()
        .enumerate()
        .filter_map(|(idx, keep)| if *keep { Some(idx) } else { None })
        .collect();
    let embargo_count = embargoed.into_iter().filter(|v| *v).count();

    (train_indices, purged_count, embargo_count)
}