DeepLC 4.5.0: keep the multitask prediction matrix as its low-rank factors

[RobbinBouwmeester]

Stacked on #125. Ships as 4.5.0, and fixes the reported out-of-memory failure with no change on the MS2Rescore side.

The failure

MS2Rescore asks DeepLC for the whole matrix for every PSM:

pred_matrix = predict(psm_list, ..., return_matrix=True)

At 6,543 setups that is 25.6 kB per PSM, so a 2.6 M PSM run needs 63 GiB and dies:

DefaultCPUAllocator: not enough memory: you tried to allocate 67731356304 bytes

Change

FactorHead ends with proj(trunk) @ embedding.T * scale + shift, so the (n_peptides, n_tasks) matrix is determined by (n_peptides, rank).

2,587,932 peptides
dense (n, 6543) float32	63.1 GiB
factors (n, 64) float32	0.617 GiB

predict(..., return_matrix=True) returns a FactoredPredictionMatrix holding those factors. It reports the same shape, indexes the same way, and selecting rows alone gives another factored matrix, which is what keeps a per-run slice cheap. Anything the factors cannot answer falls through to the dense array via __getattr__, so out.min() and np.isfinite(out) behave exactly as before. predict_kwargs={"factored": False} forces the dense array.

Second half, without which the first is not enough: MultiHeadCalibration.fit no longer expands a lazy source. Head ranking already walked the heads in blocks and every later step reads single columns, so nothing in fit needed the matrix whole. Block size is now chosen from the row count, so the peak stays near 64 MiB regardless of reference size:

     1,000 rows ->  6543 heads per block =  49.9 MiB
    20,000 rows ->   419 heads per block =  63.9 MiB
   400,000 rows ->    20 heads per block =  61.0 MiB
 2,587,932 rows ->     3 heads per block =  59.2 MiB

That matters because MS2Rescore's default calibration_set_size=None calibrates on every target PSM at q <= 0.01, which can be most of the run.

Verified against the MS2Rescore path

Replaying ms2rescore 4.0.2's feature generator against this build, with every PSM calibrating (the worst case), instrumented to record every densification:

full matrix would be            99.8 MiB  (4000, 6543)
widest block actually built      3.9 MiB
factors held                     2.6 MiB
calibrated 4000/4000, MAE 0.235

no dense peptide x setup matrix at any point, with every PSM calibrating

It still selects the correct head and calibrates correctly.

Why the default, after all

An earlier revision made this opt-in because returning a non-ndarray broke eight tests that call out.min(). The __getattr__ fallback removes that objection: the result answers everything an array does, and is only cheap where it can be. All 246 tests pass with it on by default.

Requires ms2rescore >= 4.0.2

4.0.1 imports deeplc.core._best_correlating_head, removed in 4.3.0. 4.0.2 already guards that import and takes the MultiHeadRidgeCalibration path, which is the one this PR makes lazy. So the deployment fix is: upgrade DeepLC to 4.5.0, and ms2rescore to 4.0.2 if not already there. No code change in either.

Tests

tests/test_factored.py, 16 cases. Seven indexing forms parametrised against the dense matrix from the same model, agreeing to 7.6e-5 min (float32 accumulation order). Plus row slices staying factored, array-method fallback, np.asarray round-trip, is_head_source passthrough, both non-factorable cases (single-task, and a head fine-tuned with add_task), and rank mismatch rejected at construction.

246 pass. ruff check and ruff format --check clean on every file touched.

🤖 Generated with Claude Code