Adding the KFoldQuantileSplitter

[florisvdf]

Changes

Adding K-fold cross validation support for quantile splitting.

Newly introduced convention for Subsets structure

Since quantile splitting creates test folds that contain samples that never appear in the training data, the structure of the resulting Subsets object is slightly different from the Subsets object created by the KFoldSplitter. The former creates a dictionary with a key train_folds storing all the training sets for each fold, and a key test_folds storing all the test folds.

Existing proteingym-benchmark models do not yet comply with this structure.

Checklist

• I broke the PR down so that it contains a reasonable amount of changes for an effective review
• I performed a self-review of my code. Amongst other things, I have commented my code in hard-to-understand areas.
• I made corresponding changes to the documentation
• I added tests that prove my fix is effective or that my feature works
• I accounted for dependent changes to be merged and published in downstream modules

[codecov]

Codecov Report

❌ Patch coverage is 98.13084% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/proteingym/base/splits.py	98.13%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[hredestig]

Let's think of the api a bit - and also, can you please add some content to the splits notebook as well to demonstrate how kfold-quantile is used

[florisvdf]

Couple of big changes. First of all I addressed the issue of conforming with split creation api and added a demo to the notebook demonstrating the usage of the quantile splitters. Next, I realized that there was a significant problem with split assignment in both quantile splitters, namely that they reasoned over sequences at the individual assay level, rather than at the aggregated dataset level. This caused a number of downstream inconsistent and breaking behaviors. I described my fix below:

Instead of reasoning on the level of per-assay records, both the QuantileSplitter and the KFoldQuantileSplitter reason about variants at the aggregated level (data that results from .to_df(target_names=[target]). This fixes two important issues with the splitting behavior:

1. Reasoning over variants at the individual record level in the dataset yields quantile thresholds and values of top_k that may no longer be valid once the data is aggregated for a given target. This is now fixed by reasoning at the aggregated dataset (result of .to_df) level.
2. With split assignment at the individual record level, one may encounter the scenario where a sequence is assigned to the upper quantile when its property value exceeds the threshold in assay_i, but also assigned to the lower quantile when it is under the threshold in assay_j, creating conflicting AssaySlice objects. Split assignment now also happens after sequences are first aggregated.

The splitters now also raise a ValueError when one attempts to use them to split datasets with a target that combines assays with varying assay variables, since property value thresholds can not be compared across different assay conditions.

One concern worth raising is that the aggregation method used in downstream training jobs must match the default aggregation of to_df , otherwise top_k is no longer valid. Not sure how to validate this, though seems like an very rare edge case.

[karel-w]

We now allow for empty splits, which will fail downstream when we cannot extract data from the split. I would suggest to fail early and raise errors when we obtain an empty split instead.

This also removes some of the edge cases in testing.

[florisvdf]

Thanks a lot for the extensive review. I addressed most of the comments, most importantly no longer allowing for empty slices to be created when the passed target is not present in the assays.

[karel-w]

Lets raise an issue for the validation of QuantileFolds to discuss this point, but in the meanwhile we can move on.

Line 325 is now redundant since we would fail earlier when the slice is empty? Would maybe be a nice test to make sure _count_high_property_variants return correct counts.

[florisvdf]

Thanks, made the final changes. Will raise issue for Quantile validation splits