diff --git a/.gitignore b/.gitignore index 1ad3487..5814c41 100644 --- a/.gitignore +++ b/.gitignore @@ -21,6 +21,7 @@ pip-wheel-metadata/ # Test runner and coverage artifacts .pytest_cache/ .pytest_tmp*/ +.tmp_pytest_*/ .coverage .coverage.* coverage.xml diff --git a/AGENTS.md b/AGENTS.md index 57b0fb3..9862773 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -27,8 +27,8 @@ The main CLIs are: - `pepseqpred-labels` - `pepseqpred-predict` - `pepseqpred-preprocess` -- `pepseqpred-train-ffnn` -- `pepseqpred-train-ffnn-optuna` +- `pepseqpred-train` +- `pepseqpred-train-optuna` These map to files in `src/pepseqpred/apps/`. @@ -51,8 +51,8 @@ Batch scripts live in `scripts/hpc/`. These are part of the intended workflow, e - label generation - preprocessing - prediction -- FFNN training -- FFNN Optuna tuning +- model-head training +- model-head Optuna tuning Treat these scripts as first class project interfaces, not throwaway helpers. @@ -181,8 +181,8 @@ Available CLIs: - `pepseqpred-labels` - `pepseqpred-predict` - `pepseqpred-preprocess` -- `pepseqpred-train-ffnn` -- `pepseqpred-train-ffnn-optuna` +- `pepseqpred-train` +- `pepseqpred-train-optuna` ## Testing guidance @@ -250,4 +250,4 @@ For most tasks: ### `tests/` - add targeted coverage for bug fixes -- do not rewrite unrelated fixtures or tests just for style \ No newline at end of file +- do not rewrite unrelated fixtures or tests just for style diff --git a/README.md b/README.md index ce00121..61704b8 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ For lightweight inference usage and API quickstart, use [README.pypi.md](README. PepSeqPred supports two usage profiles: - **PyPI quickstart profile (`pip install pepseqpred`)**: user-facing inference API with bundled pretrained artifacts and artifact-path inference helpers. -- **Repository developer profile (`pip install -e .[dev]`)**: full source tree for preprocessing, embeddings, label generation, FFNN training, Optuna tuning, prediction, evaluation, and HPC orchestration. +- **Repository developer profile (`pip install -e .[dev]`)**: full source tree for preprocessing, embeddings, label generation, model-head training, Optuna tuning, prediction, evaluation, and HPC orchestration. The repository profile is the source of truth for reproducing experiments end-to-end. @@ -52,7 +52,7 @@ The repository profile is the source of truth for reproducing experiments end-to Stage 1 normalize dataset inputs (PV1/CWP/BKP) to a shared training contract Stage 2 generate ESM-2 per-residue embeddings Stage 3 build residue-level label shards -Stage 4 train FFNN (unified n-fold interface, DDP-aware) +Stage 4 train model head (unified n-fold interface, DDP-aware) Stage 5 optional Optuna tuning (DDP-aware) Stage 6 predict residue masks from checkpoint/manifest Stage 7 evaluate residue metrics (+ optional Cocci peptide compare) @@ -206,6 +206,11 @@ pepseqpred-esm \ - embedding index CSV under `/artifacts/*.csv` - optional shard-specific outputs when `--num-shards > 1` +Length feature note: + +- `--seq-len-feature {none,raw,inverse}` controls whether a sequence-length scalar is appended to every residue embedding. +- The default is `none`. `raw` appends `float(seq_len)`; `inverse` appends `1.0 / seq_len`. + ### Stage 3: Build Residue Labels **CLI:** `pepseqpred-labels` (`src/pepseqpred/apps/labels_cli.py`) @@ -236,15 +241,18 @@ pepseqpred-labels \ - label shard `.pt` with protein label tensors and peptide metadata - optional `class_stats` payload when `--calc-pos-weight` is enabled -### Stage 4: Train FFNN +### Stage 4: Train Model Head -**CLI:** `pepseqpred-train-ffnn` (`src/pepseqpred/apps/train_ffnn_cli.py`) +**CLI:** `pepseqpred-train` (`src/pepseqpred/apps/train_cli.py`) **Unified run interface** - `--n-folds 1`: one holdout run per split/train seed pair (uses `--val-frac`) - `--n-folds K` (`K > 1`): K-fold members per split/train seed pair set - `--split-seeds` and `--train-seeds` are paired by index; if both are omitted, both default to `--seed` +- `--model-head ffnn` is the default and preserves existing dense-head behavior +- `--model-head conv1d` adds a local Conv1d feature stack before the dense residue classifier +- `--seq-len-feature {none,raw,inverse}` records whether embeddings include an appended sequence-length feature; use the same mode used during embedding generation **Core modules** @@ -255,18 +263,21 @@ pepseqpred-labels \ **Command (smoke)** ```bash -pepseqpred-train-ffnn \ +pepseqpred-train \ --embedding-dirs data/esm2/artifacts/pts/shard_000 \ --label-shards data/labels/labels_shard_000.pt \ --epochs 1 \ + --model-head ffnn \ --subset 100 \ --save-path data/models/ffnn_smoke \ --results-csv data/models/ffnn_smoke/runs.csv ``` +For a local sequence head, use `--model-head conv1d` and optionally tune `--conv-channels`, `--conv-layers`, `--conv-kernel-size`, and `--conv-dropout`. + **Submit one SLURM training job with multiple datasets (PV1 + CWP + BKP)** -`scripts/hpc/trainffnn.sh` accepts multiple embedding directories and multiple label shards in one call: +`scripts/hpc/train.sh` accepts multiple embedding directories and multiple label shards in one call: - all embedding dirs first - separator `--` @@ -304,7 +315,7 @@ LABEL_SHARDS=( /scratch/$USER/labels/bkp/labels_shard_003.pt ) -sbatch trainffnn.sh "${EMB_DIRS[@]}" -- "${LABEL_SHARDS[@]}" +sbatch train.sh "${EMB_DIRS[@]}" -- "${LABEL_SHARDS[@]}" ``` Notes: @@ -321,7 +332,7 @@ Notes: ### Stage 5: Optuna Tuning (Optional) -**CLI:** `pepseqpred-train-ffnn-optuna` (`src/pepseqpred/apps/train_ffnn_optuna_cli.py`) +**CLI:** `pepseqpred-train-optuna` (`src/pepseqpred/apps/train_optuna_cli.py`) **Core modules** @@ -331,7 +342,7 @@ Notes: **Command (smoke)** ```bash -pepseqpred-train-ffnn-optuna \ +pepseqpred-train-optuna \ --embedding-dirs data/esm2/artifacts/pts/shard_000 \ --label-shards data/labels/labels_shard_000.pt \ --n-trials 2 \ @@ -342,10 +353,11 @@ pepseqpred-train-ffnn-optuna \ **Current Optuna search space** -The current study samples the following hyperparameters per trial: +Optuna is fixed-head per study. `--model-head ffnn` samples the dense-head space. `--model-head conv1d` samples the same dense/optimizer settings plus convolutional head settings. | Hyperparameter (`best_params` key) | Type | Search space (current implementation) | Controlled by | | --- | --- | --- | --- | +| `model_head` | fixed | `ffnn` or `conv1d` | `--model-head` | | `depth` | integer | `[depth_min, depth_max]` | `--depth-min`, `--depth-max` | | `width_step` | categorical | `{16, 32, 64}` | fixed in code | | `base_width` | integer | `[width_min, width_max]` with `step=width_step` | `--width-min`, `--width-max` | @@ -353,6 +365,10 @@ The current study samples the following hyperparameters per trial: | `dropout` | float | `[0.00, 0.25]` | fixed in code | | `use_layer_norm` | categorical | `{True, False}` | fixed in code | | `use_residual` | categorical | `{True, False}` | fixed in code | +| `conv_channels` | categorical | values from `--conv-channel-choices` | conv1d only | +| `conv_layers` | integer | `[conv_layers_min, conv_layers_max]` | conv1d only | +| `conv_kernel_size` | categorical | odd values from `--conv-kernel-size-choices` | conv1d only | +| `conv_dropout` | float | `[conv_dropout_min, conv_dropout_max]` | conv1d only | | `learning_rate` | float (log) | `[lr_min, lr_max]` | `--lr-min`, `--lr-max` | | `weight_decay` | float (log) | `[wd_min, wd_max]` | `--wd-min`, `--wd-max` | | `batch_size` | categorical | values from `--batch-sizes` CSV | `--batch-sizes` | @@ -368,13 +384,14 @@ Not tuned by Optuna in the current setup: - `pos_weight` (fixed for the study via `--pos-weight`, or computed once from label shards if omitted) - split strategy and validation fraction (`--split-type`, `--val-frac`) - sequence windowing (`--window-size`, `--stride`) and data-loader behavior +- sequence-length feature mode (`--seq-len-feature`; defaults to `none`) - trial budget/pruning controls (`--n-trials`, `--epochs`, `--pruner-warmup`, `--timeout-s`) - optimization target metric selection (`--metric`) is user-selected, then maximized by Optuna HPC default override note: - The CLI default for `--batch-sizes` is `32,64,128`. -- The SLURM wrapper `scripts/hpc/trainffnnoptuna.sh` currently overrides this to `256,512,1024` unless changed via env var. +- The SLURM wrapper `scripts/hpc/trainoptuna.sh` currently overrides this to `256,512,1024` unless changed via env var. **Outputs** @@ -410,6 +427,11 @@ pepseqpred-predict \ - FASTA containing binary residue masks +Length feature note: + +- `--seq-len-feature auto` is the default and resolves from checkpoint `model_config.seq_len_feature`. +- A missing checkpoint key means no appended sequence-length feature. Use `--seq-len-feature raw` only when predicting with older or explicit raw-length models. + ### Stage 7: Evaluate **CLI:** `pepseqpred-eval-ffnn` (`src/pepseqpred/apps/evaluate_ffnn_cli.py`) @@ -470,9 +492,9 @@ Bundled pretrained registry currently includes: ### Embedding `.pt` -- tensor shape: `(L, D+1)` +- tensor shape: `(L, D)` by default, or `(L, D+1)` when `--seq-len-feature raw` or `inverse` was used - `L`: residue count -- final feature column stores sequence length +- optional final feature column stores either raw sequence length or inverse sequence length ### Label shard `.pt` @@ -493,6 +515,10 @@ Bundled pretrained registry currently includes: "optim_state_dict": ..., "epoch": int, "config": {...}, + "model_config": { + # seq_len_feature is omitted for no appended length feature + # or set to "raw" / "inverse" when present + }, "best_loss": float, "metrics": {...} } @@ -512,8 +538,8 @@ Bundled pretrained registry currently includes: | `pepseqpred-preprocess` | `apps/preprocess_cli.py` | metadata + z-score preprocessing | | `pepseqpred-esm` | `apps/esm_cli.py` | ESM-2 embedding generation | | `pepseqpred-labels` | `apps/labels_cli.py` | residue label shard generation | -| `pepseqpred-train-ffnn` | `apps/train_ffnn_cli.py` | unified holdout/K-fold FFNN training (`--n-folds`) | -| `pepseqpred-train-ffnn-optuna` | `apps/train_ffnn_optuna_cli.py` | Optuna tuning | +| `pepseqpred-train` | `apps/train_cli.py` | unified holdout/K-fold model-head training (`--n-folds`) | +| `pepseqpred-train-optuna` | `apps/train_optuna_cli.py` | fixed-head Optuna tuning | | `pepseqpred-predict` | `apps/prediction_cli.py` | FASTA inference from checkpoint/manifest | | `pepseqpred-eval-ffnn` | `apps/evaluate_ffnn_cli.py` | residue-level evaluation | @@ -525,8 +551,8 @@ These wrappers are production-facing interfaces and should be treated as first-c | --- | --- | --- | | `generateembeddings.sh` | Embeddings | GPU, array `0-3`, `a100`, `2` CPU/GPU, `8G`/GPU, `01:00:00` | | `generatelabels.sh` | Labels | CPU, `1` CPU, `16G`, `01:00:00` | -| `trainffnn.sh` | Train FFNN | GPU, `4xa100`, `20` CPU, `256G`, `12:00:00` | -| `trainffnnoptuna.sh` | Optuna | GPU, `4xa100`, `20` CPU, `448G`, `48:00:00` | +| `train.sh` | Train model head | GPU, `4xa100`, `20` CPU, `256G`, `12:00:00` | +| `trainoptuna.sh` | Optuna | GPU, `4xa100`, `20` CPU, `448G`, `48:00:00` | | `predictepitope.sh` | Predict | GPU, `a100`, `4` CPU, `32G`, `00:30:00` | | `evaluateffnn.sh` | End-to-end eval pipeline | GPU, `a100`, `8` CPU, `128G`, `04:00:00` | | `evalffnnsweep.sh` | Set-indexed eval batch submitter | wrapper script (calls `evaluateffnn.sh`) | @@ -536,7 +562,7 @@ These wrappers are production-facing interfaces and should be treated as first-c - `evaluateffnn.sh` orchestrates prepare, embed, labels, predict, eval, and peptide compare stages with stage toggles (`RUN_PREP`, `RUN_EMBED`, `RUN_LABELS`, `RUN_PREDICT`, `RUN_EVAL`, `RUN_COMPARE`). - `evaluateffnn.sh` and `evalffnnsweep.sh` depend on `scripts/tools/cocci_eval_pipeline.py`. -- HPC wrappers expect `.pyz` runtime artifacts in the working directory (for example `esm.pyz`, `train_ffnn.pyz`, `predict.pyz`). +- HPC wrappers expect `.pyz` runtime artifacts in the working directory (for example `esm.pyz`, `train.pyz`, `predict.pyz`). ## Zipapp and Tooling (`scripts/tools`) diff --git a/docs/pv1_cwp_bkp_merge_split_and_pos_weight.md b/docs/pv1_cwp_bkp_merge_split_and_pos_weight.md index 66cbfd8..8975f83 100644 --- a/docs/pv1_cwp_bkp_merge_split_and_pos_weight.md +++ b/docs/pv1_cwp_bkp_merge_split_and_pos_weight.md @@ -61,7 +61,7 @@ That is the behavior exercised in `tests/integration/test_prepare_dataset_multis ### 3.1 Base protein universe -`train_ffnn_cli` builds a `ProteinDataset` from provided embedding dirs + label shards, then uses: +`train_cli` builds a `ProteinDataset` from provided embedding dirs + label shards, then uses: - `protein_ids = intersection(embedding_index IDs, label_index IDs)` @@ -119,7 +119,7 @@ There are two connected pieces: - `pos_weight = neg_count / max(1, pos_count)` - For 3-column labels `[Def epitope, Uncertain, Not epitope]`, counts only include residues where `Uncertain == 0`. -2. Train time (`pepseqpred-train-ffnn`) +2. Train time (`pepseqpred-train`) - If `--pos-weight` is provided, that value is used directly. - Otherwise it reads `class_stats` from all provided label shards and recomputes: - `total_neg / max(1, total_pos)` diff --git a/docs/training_quality_findings.md b/docs/training_quality_findings.md new file mode 100644 index 0000000..220419e --- /dev/null +++ b/docs/training_quality_findings.md @@ -0,0 +1,348 @@ +# PepSeqPred Training Quality Findings + +This document summarizes the read-only exploratory review of PepSeqPred training, evaluation, label, and HPC code paths for unexpectedly poor multi-pathogen dataset results. + +No code edits were made during the investigation. The review focused on likely training-quality failure modes rather than style or cosmetic issues. + +## Scope Reviewed + +Primary files inspected: + +- `src/pepseqpred/apps/train_cli.py` +- `src/pepseqpred/apps/train_optuna_cli.py` +- `src/pepseqpred/apps/evaluate_ffnn_cli.py` +- `src/pepseqpred/core/models/ffnn.py` +- `src/pepseqpred/core/train/trainer.py` +- `src/pepseqpred/core/train/threshold.py` +- `src/pepseqpred/core/train/split.py` +- `src/pepseqpred/core/train/metrics.py` +- `src/pepseqpred/core/train/weights.py` +- `src/pepseqpred/core/data/proteindataset.py` +- `src/pepseqpred/core/labels/builder.py` +- `src/pepseqpred/core/preprocess/preparedataset.py` +- `scripts/hpc/train.sh` +- `scripts/hpc/trainoptuna.sh` +- `scripts/hpc/evaluateffnn.sh` +- related tests under `tests/unit`, `tests/integration`, and `tests/e2e` + +## Highest-Risk Findings + +### 1. Label/objective mismatch for residue-level prediction + +Current label generation expands peptide labels across every residue in a peptide alignment window. If a peptide is reactive, every residue in that peptide window becomes a positive residue. + +Evidence: + +- `src/pepseqpred/core/labels/builder.py` + - `_build_labels_for_protein` marks `def_mask[start:stop] = True` for definite epitope peptides. + - The FFNN then trains residue-level BCE on those expanded residue labels. + +Why this can hurt: + +- A reactive peptide means "the peptide contains an epitope signal", not necessarily "every residue in this peptide is epitope". +- This creates dense false-positive residue labels within positive peptides. +- The model is optimized for residue-wise correctness under noisy residue labels, while downstream use may care about peptide/protein regions or sparse true epitopes. +- The issue can compound in multi-pathogen data if peptide lengths, overlap density, or labeling criteria differ by source. + +Planning direction: + +- Add diagnostics comparing peptide-level labels to residue-level expansion density. +- Consider a multiple-instance learning objective, peptide-window objective, or soft/weak residue labels. +- If residue labels remain, consider down-weighting positive residues within reactive peptide spans or using boundary-aware smoothing. + +### 2. Sparse or zero-valid windows still participate in training + +Resolved for optimizer behavior: the dataset still yields windows with no valid labeled residues, but a synchronized step with zero valid residues globally now completes a graph-connected zero backward pass without advancing Adam. A rank with zero local support still takes the optimizer step when another rank has valid residues, using the globally reduced gradient. + +Evidence: + +- `src/pepseqpred/core/data/proteindataset.py` + - `ProteinDataset.__iter__` yields every window from `_iter_windows`. + - The final mask may be all zero when a window contains only uncertain or padded residues. +- `src/pepseqpred/core/train/trainer.py` + - `_batch_step` distinguishes local valid support from global valid support. + - Globally empty steps skip `optimizer.step()` on every rank. + - Exhausted ranks use zero-valid dummy batches so their Adam state stays synchronized while other ranks still have data. + +Remaining concern: + +- Zero-valid windows still consume data-loading and forward/backward compute. +- If some pathogens or proteins have many uncertain/unlabeled regions, they can consume substantial runtime without adding training signal. + +Planning direction: + +- Add logging for zero-valid windows by split, fold, rank, and source. +- Optionally filter zero-valid windows in `ProteinDataset` for training. + +### 3. DDP and batch loss weighting are likely biased by local valid-residue counts + +Resolved: training now backpropagates each rank's summed masked BCE scaled by `world_size / global_valid_residues`. After DDP averages gradients, the result is the globally valid-residue-normalized gradient. + +Evidence: + +- `src/pepseqpred/core/train/trainer.py` + - Every rank reduces its valid-residue count before backward. + - Exhausted ranks remain in a collective-aware loop using minimal zero-valid dummy batches until all ranks are exhausted. + - Training no longer relies on `DDP.join()`, which cannot shadow the custom loss collective. +- `src/pepseqpred/apps/train_cli.py` + - IDs are partitioned across ranks by estimated embedding file size, not by valid positive/negative residue count. + +Remaining concern: + +- Weighted partitioning now affects efficiency and late-step effective batch size rather than gradient correctness. +- Large rank imbalance produces more dummy batches and smaller global batches near the end of an epoch. + +Planning direction: + +- Monitor the per-rank real/dummy batch counts emitted by `train_sync_summary`. +- Consider balancing partitions by estimated window or valid-residue counts if dummy-batch fractions are large. + +Test coverage: + +- Unit tests cover local/global zero-valid behavior, loss scaling, dummy generation, and synchronized iteration. +- A two-process CPU/Gloo integration test uses one versus three rank-local batches across two epochs and compares model parameters and Adam state. + +### 4. Training uses overlapping windows while evaluation uses full proteins + +Training defaults to windowed proteins with overlap, while evaluation defaults to full proteins. + +Evidence: + +- `src/pepseqpred/core/data/proteindataset.py` + - `_iter_windows` supports `window_size` and `stride`. +- `scripts/hpc/train.sh` + - `WINDOW_SIZE=1000` + - `STRIDE=900` +- `src/pepseqpred/apps/evaluate_ffnn_cli.py` + - Evaluation constructs datasets with `window_size=None` and `pad_last_window=False`. + +Why this can hurt: + +- Overlapped residues are duplicated during training and validation. +- Validation threshold selection is based on windowed validation arrays, but final evaluation is on full proteins. +- Overlap duplicates can overweight boundary regions or long proteins. +- If there are very long multi-pathogen proteins, length and overlap can distort training and metrics. + +Planning direction: + +- Run an ablation with `--window-size 0` where feasible. +- Run an ablation with non-overlapping windows, for example `stride == window_size`. +- Add metrics that count unique proteins/residues separately from yielded training residues. + +### 5. Positive class weight is stale or fold-inappropriate + +HPC scripts hard-code a positive class weight. When auto-computed, the weight is computed over all label shards, not the current training split. + +Evidence: + +- `scripts/hpc/train.sh` + - `POS_WEIGHT="${POS_WEIGHT:-13.18999647945325}"` + - The script always passes `--pos-weight "$POS_WEIGHT"`. +- `scripts/hpc/trainoptuna.sh` + - Same hard-coded default positive weight. +- `src/pepseqpred/apps/train_cli.py` + - If `--pos-weight` is absent, `pos_weight_from_label_shards(label_shards)` uses all provided label shard totals. +- `docs/pv1_cwp_bkp_merge_split_and_pos_weight.md` + - Already notes that automatic train-time `pos_weight` uses shard-level totals, not train-only IDs. + +Why this can hurt: + +- Multi-pathogen data can have very different positive rates than the dataset used to derive the hard-coded value. +- K-fold training may have materially different positive rates per fold. +- Using validation labels to compute the training class weight is also a mild leakage/selection mismatch, even if it is not target leakage in the usual model-fitting sense. + +Planning direction: + +- Compute `pos_weight` from the current run's training IDs only. +- Log per-run and per-fold positive/negative residue counts before training. +- Stop hard-coding old positive weights in HPC defaults unless explicitly requested. + +### 6. Threshold policy is hard-coded and can be overly conservative + +Training chooses thresholds by maximizing recall subject to minimum precision 0.25. If that precision is unreachable, it falls back to the best precision threshold. Ensemble prediction then uses majority vote on member-thresholded binary masks. + +Evidence: + +- `src/pepseqpred/core/train/trainer.py` + - Calls `find_threshold_max_recall_min_precision(y_true, y_prob, min_precision=0.25)`. +- `src/pepseqpred/core/train/threshold.py` + - Fallback favors best precision when minimum precision is unreachable. +- `src/pepseqpred/core/predict/inference.py` + - Ensemble prediction thresholds each member independently and uses majority vote. + - Ties are effectively negative because `votes_needed = n_members // 2 + 1`. + +Why this can hurt: + +- A hard minimum precision can push thresholds high and crush recall on difficult or shifted pathogen groups. +- Majority vote over conservative masks further reduces positive calls. +- The reported validation threshold may not transfer to external multi-pathogen evaluation. + +Planning direction: + +- Make threshold selection configurable. +- Evaluate fixed thresholds, best-F1 thresholds, max-MCC thresholds, and recall-target thresholds. +- For ensembles, compare majority vote with mean-probability thresholding. +- Report threshold status, selected threshold, predicted positive fraction, and recall by pathogen/family. + +### 7. Grouped splits prevent family leakage but are not label-stratified + +The default split type keeps families/groups intact, which is good for leakage control. However, folds are assigned primarily by group size, not by positive/negative support or source balance. + +Evidence: + +- `src/pepseqpred/core/train/split.py` + - `split_ids_grouped` and `build_grouped_kfold_splits` keep groups intact. + - Grouped k-fold assigns larger groups first to balance fold size. + +Why this can hurt: + +- Folds can have very different positive rates. +- Some validation folds may contain source/pathogen groups not meaningfully represented in training. +- Metrics can look much worse than expected if the split is actually a difficult cross-family generalization test. + +Planning direction: + +- Add split reports with per-fold: + - protein count + - valid residue count + - positive residue count + - negative residue count + - positive rate + - source/pathogen/family counts +- Consider grouped stratified splitting where possible. +- Preserve leakage safety, but balance label support more deliberately. + +### 8. Model has no local sequence modeling beyond ESM embeddings + +The FFNN flattens residues and predicts each residue independently after the ESM embedding. It does not model neighboring residue interactions during classification. + +Evidence: + +- `src/pepseqpred/core/models/ffnn.py` + - `PepSeqFFNN.forward` flattens `(B, L, D)` into `(B * L, D)`. + - Output is reshaped back to `(B, L)`. + +Why this can hurt: + +- Residue-level epitope signals are often regional. +- The ESM embedding contains context, but the classifier cannot enforce local smoothness or peptide-level consistency. +- Weak peptide-expanded labels may need a model/objective that understands windows rather than independent residues. + +Planning direction: + +- First fix diagnostics, objective, and weighting before changing architecture. +- Then consider light local context heads, such as 1D convolution, CRF-like smoothing, or pooling over peptide windows. + +Implementation status: + +- A generic training CLI now supports `--model-head ffnn` and `--model-head conv1d`. +- `ffnn` remains the default to preserve previous dense-head behavior. +- New checkpoints record explicit model-head metadata so prediction and evaluation can reload conv checkpoints without architecture flags. + +### 9. Raw protein sequence length is appended as an unnormalized feature + +Resolved: sequence length is now explicit and opt-in. Embedding generation defaults to no appended length feature, and training records only `"raw"` or `"inverse"` in `model_config.seq_len_feature`; a missing key means no appended length feature. + +Original evidence: + +- `src/pepseqpred/core/embeddings/esm2.py` + - `append_seq_len` appends `float(seq_len)` as a column. + - Embedding generation and prediction embedding paths used this behavior by default. + +Why this can hurt: + +- Raw length can be orders of magnitude larger than normalized embedding features. +- It can become a source/pathogen shortcut if sequence length correlates with dataset or label source. +- Multi-pathogen training is particularly susceptible to spurious source-specific features. + +Planning direction: + +- Use the default no-length behavior for new runs. +- Use `--seq-len-feature raw` only for explicit compatibility or ablation runs. +- Use `--seq-len-feature inverse` for the supported bounded transform (`1.0 / seq_len`). + +## Concrete Bug Found + +### Optuna best checkpoint copy path mismatch + +Resolved: `train_optuna_cli.py` now accepts both the legacy copied filename and the trainer's `fully_connected_by_score.pt` when writing the root-level `best_model_by_score.pt`. + +Optuna trial training saves trial checkpoints using one filename, but the CLI later tries to copy a different filename. + +Evidence: + +- `src/pepseqpred/core/train/trainer.py` + - `fit_optuna` saves `fully_connected_by_score.pt`. +- `src/pepseqpred/apps/train_optuna_cli.py` + - Copies `fully_connected_by_score.pt` from the best trial directory, with fallback support for a preexisting `best_model_by_score.pt`. + +Impact: + +- The root-level `best_model_by_score.pt` may never be written. +- Users may accidentally evaluate stale or missing artifacts after Optuna. + +Test coverage: + +- `tests/integration/test_train_clis_inprocess.py` asserts that the root-level best Optuna checkpoint is copied and carries model-head metadata. + +## Configuration Surprises + +### `--use-pos-weight` in Optuna appeared misleading + +Resolved: `--use-pos-weight` has been removed from `train_optuna_cli.py`. Positive weighting is automatic when `--pos-weight` is omitted, and `--pos-weight` remains the manual override. + +Impact: + +- Users may have thought positive weighting was optional in Optuna when it was effectively always active unless code was changed. + +Planning direction: + +- No further action needed for this flag. + +### Threshold minimum precision is not configurable + +The value `0.25` is embedded in trainer evaluation and not exposed as a CLI option. + +Impact: + +- Experiments cannot easily compare different operating points without code edits. + +Planning direction: + +- Expose threshold policy and minimum precision as CLI options. +- Save threshold policy metadata in checkpoints and manifests. + +## Diagnostics To Add Before Major Changes + +Recommended first-pass diagnostics: + +1. Split/fold label-balance report. + - Per split, fold, rank, source, and family. + - Counts for proteins, windows, valid residues, positive residues, negative residues, positive rate. + +2. Window-validity report. + - Number and fraction of zero-valid windows. + - Number and fraction of windows with no positives. + - Same metrics by source/family. + +3. Training-loss denominator report. + - Per batch and per rank valid residue counts. + - Detect ranks/batches contributing gradients from very different valid counts. + +4. Threshold diagnostics. + - Selected threshold. + - Threshold status. + - Predicted positive residue fraction. + - Precision, recall, F1, MCC, PR AUC at multiple fixed thresholds. + +5. Per-protein and per-family evaluation. + - Avoid relying only on flattened residue-level metrics. + - Include peptide-level "any positive" metrics when evaluating reactive/nonreactive peptide tasks. + +6. Positive-weight provenance. + - Record whether `pos_weight` came from CLI, all label shards, or train IDs. + - Record numerator and denominator counts. + +## Validation Notes + +The original findings were produced from static code inspection. Subsequent resolved sections have targeted unit and integration coverage. The real-process Gloo regression is Linux-only, matching CI and HPC; no multi-GPU training job has yet been run for this change. diff --git a/pyproject.toml b/pyproject.toml index b4b3730..7750eac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pepseqpred" -version = "1.1.1" +version = "1.2.1" description = "Residue-level epitope prediction pipeline for peptide/protein workflows." readme = "README.pypi.md" requires-python = ">=3.12" @@ -74,8 +74,8 @@ pepseqpred-predict = "pepseqpred.apps.prediction_cli:main" pepseqpred-preprocess = "pepseqpred.apps.preprocess_cli:main" pepseqpred-prepare-dataset = "pepseqpred.apps.prepare_dataset_cli:main" pepseqpred-eval-ffnn = "pepseqpred.apps.evaluate_ffnn_cli:main" -pepseqpred-train-ffnn = "pepseqpred.apps.train_ffnn_cli:main" -pepseqpred-train-ffnn-optuna = "pepseqpred.apps.train_ffnn_optuna_cli:main" +pepseqpred-train = "pepseqpred.apps.train_cli:main" +pepseqpred-train-optuna = "pepseqpred.apps.train_optuna_cli:main" [tool.setuptools] package-dir = {"" = "src"} diff --git a/scripts/hpc/evaluateffnn.sh b/scripts/hpc/evaluateffnn.sh index d7f0096..9ec75ab 100644 --- a/scripts/hpc/evaluateffnn.sh +++ b/scripts/hpc/evaluateffnn.sh @@ -45,6 +45,8 @@ usage() { echo " MAX_TOKENS default: 1022" echo " EMBED_BATCH_SIZE default: 24" echo " THRESHOLD default: unset" + echo " ENSEMBLE_AGGREGATION default: majority" + echo " ENSEMBLE_THRESHOLD default: unset" echo " ENSEMBLE_SET_INDEX default: 1" echo " EXPECTED_SET_INDEX default: unset (optional guard; fail if resolved set differs)" echo " K_FOLDS default: unset" @@ -104,6 +106,8 @@ MODEL_NAME="${MODEL_NAME:-esm2_t33_650M_UR50D}" MAX_TOKENS="${MAX_TOKENS:-1022}" EMBED_BATCH_SIZE="${EMBED_BATCH_SIZE:-24}" THRESHOLD="${THRESHOLD:-}" +ENSEMBLE_AGGREGATION="${ENSEMBLE_AGGREGATION:-majority}" +ENSEMBLE_THRESHOLD="${ENSEMBLE_THRESHOLD:-}" ENSEMBLE_SET_INDEX="${ENSEMBLE_SET_INDEX:-1}" EXPECTED_SET_INDEX="${EXPECTED_SET_INDEX:-}" K_FOLDS="${K_FOLDS:-}" @@ -253,6 +257,8 @@ PREDICT_ARGS=( --log-json ) [ -n "${THRESHOLD}" ] && PREDICT_ARGS+=(--threshold "${THRESHOLD}") +[ -n "${ENSEMBLE_AGGREGATION}" ] && PREDICT_ARGS+=(--ensemble-aggregation "${ENSEMBLE_AGGREGATION}") +[ -n "${ENSEMBLE_THRESHOLD}" ] && PREDICT_ARGS+=(--ensemble-threshold "${ENSEMBLE_THRESHOLD}") [ -n "${ENSEMBLE_SET_INDEX}" ] && PREDICT_ARGS+=(--ensemble-set-index "${ENSEMBLE_SET_INDEX}") [ -n "${K_FOLDS}" ] && PREDICT_ARGS+=(--k-folds "${K_FOLDS}") [ -n "${EMB_DIM}" ] && PREDICT_ARGS+=(--emb-dim "${EMB_DIM}") @@ -292,6 +298,8 @@ EVAL_ARGS=( --log-json ) [ -n "${THRESHOLD}" ] && EVAL_ARGS+=(--threshold "${THRESHOLD}") +[ -n "${ENSEMBLE_AGGREGATION}" ] && EVAL_ARGS+=(--ensemble-aggregation "${ENSEMBLE_AGGREGATION}") +[ -n "${ENSEMBLE_THRESHOLD}" ] && EVAL_ARGS+=(--ensemble-threshold "${ENSEMBLE_THRESHOLD}") [ -n "${ENSEMBLE_SET_INDEX}" ] && EVAL_ARGS+=(--ensemble-set-index "${ENSEMBLE_SET_INDEX}") [ -n "${K_FOLDS}" ] && EVAL_ARGS+=(--k-folds "${K_FOLDS}") [ -n "${SELECT_BEST_SET_RUNS_CSV}" ] && EVAL_ARGS+=(--select-best-set-runs-csv "${SELECT_BEST_SET_RUNS_CSV}") diff --git a/scripts/hpc/generateembeddings.sh b/scripts/hpc/generateembeddings.sh index 560e0a9..4d29223 100644 --- a/scripts/hpc/generateembeddings.sh +++ b/scripts/hpc/generateembeddings.sh @@ -22,11 +22,14 @@ usage() { echo " model_name default: esm2_t33_650M_UR50D" echo " max_tokens default: 1022" echo " batch_size default: 24" + echo " EMBEDDING_KEY_MODE env default: id-family (id or id-family)" + echo " METADATA_FILE env required when EMBEDDING_KEY_MODE=id-family" + echo " SEQ_LEN_FEATURE env default: none (none, raw, inverse)" echo "" echo "Examples:" - echo " sbatch $0 /scratch/\$USER/data/targets.fasta" - echo " sbatch $0 /scratch/\$USER/data/targets.fasta /scratch/\$USER/esm2 esm2_t33_650M_UR50D 1022 24" - echo " sbatch --export=ALL,IN_FASTA=/scratch/\$USER/data/targets.fasta $0" + echo " sbatch --export=ALL,METADATA_FILE=/scratch/\$USER/data/targets.metadata $0 /scratch/\$USER/data/targets.fasta" + echo " sbatch --export=ALL,METADATA_FILE=/scratch/\$USER/data/targets.metadata $0 /scratch/\$USER/data/targets.fasta /scratch/\$USER/esm2 esm2_t33_650M_UR50D 1022 24" + echo " sbatch --export=ALL,IN_FASTA=/scratch/\$USER/data/targets.fasta,METADATA_FILE=/scratch/\$USER/data/targets.metadata $0" } if [ "${1:-}" = "-h" ] || [ "${1:-}" = "--help" ]; then @@ -49,6 +52,8 @@ OUT_DIR="${2:-${OUT_DIR:-${SCRATCH_DIR}/esm2}}" MODEL_NAME="${3:-${MODEL_NAME:-esm2_t33_650M_UR50D}}" # see ESM-2 documentation for other models MAX_TOKENS="${4:-${MAX_TOKENS:-1022}}" # max number of tokens in model's context window BATCH_SIZE="${5:-${BATCH_SIZE:-24}}" # can probably get away with 16 or 24 on V100, double on A100 +METADATA_FILE="${METADATA_FILE:-}" +SEQ_LEN_FEATURE="${SEQ_LEN_FEATURE:-none}" if [ -z "${IN_FASTA}" ]; then echo "Missing required input FASTA path." @@ -60,6 +65,19 @@ fi EMBEDDING_KEY_MODE="${EMBEDDING_KEY_MODE:-id-family}" KEY_DELIMITER="${KEY_DELIMITER:--}" +METADATA_ARGS=() +if [ "${EMBEDDING_KEY_MODE}" = "id-family" ]; then + if [ -z "${METADATA_FILE}" ]; then + echo "METADATA_FILE is required when EMBEDDING_KEY_MODE=id-family." + exit 1 + fi + if [ ! -f "${METADATA_FILE}" ]; then + echo "Metadata file not found: ${METADATA_FILE}" + exit 1 + fi + METADATA_ARGS+=(--metadata-file "${METADATA_FILE}") +fi + # paths and directories LOG_DIR="logs" EMBEDDING_DIR="artifacts" @@ -85,9 +103,11 @@ ${LAUNCHER} python -u esm.pyz \ --fasta-file "${IN_FASTA}" \ --model-name "${MODEL_NAME}" \ --embedding-key-mode "${EMBEDDING_KEY_MODE}" \ + "${METADATA_ARGS[@]}" \ --key-delimiter "${KEY_DELIMITER}" \ --max-tokens "${MAX_TOKENS}" \ --batch-size "${BATCH_SIZE}" \ + --seq-len-feature "${SEQ_LEN_FEATURE}" \ --num-shards "${NUM_SHARDS}" \ --shard-id "${SHARD_ID}" \ --log-dir "${LOG_DIR}" \ diff --git a/scripts/hpc/predictepitope.sh b/scripts/hpc/predictepitope.sh index ca7ee8d..6d2d316 100644 --- a/scripts/hpc/predictepitope.sh +++ b/scripts/hpc/predictepitope.sh @@ -23,7 +23,10 @@ usage() { echo " USE_SRUN default: 1 (set 0 to run without srun)" echo " MODEL_NAME default: esm2_t33_650M_UR50D" echo " MAX_TOKENS default: 1022" + echo " SEQ_LEN_FEATURE default: auto (auto, none, raw, inverse)" echo " THRESHOLD default: unset (use checkpoint threshold)" + echo " ENSEMBLE_AGGREGATION default: majority" + echo " ENSEMBLE_THRESHOLD default: unset" echo " ENSEMBLE_SET_INDEX default: 1 (schema v2 manifest only)" echo " K_FOLDS default: unset (use all valid members)" echo " LOG_DIR default: logs" @@ -58,7 +61,10 @@ OUTPUT_FASTA="$3" MODEL_NAME="${MODEL_NAME:-esm2_t33_650M_UR50D}" MAX_TOKENS="${MAX_TOKENS:-1022}" +SEQ_LEN_FEATURE="${SEQ_LEN_FEATURE:-auto}" THRESHOLD="${THRESHOLD:-}" +ENSEMBLE_AGGREGATION="${ENSEMBLE_AGGREGATION:-majority}" +ENSEMBLE_THRESHOLD="${ENSEMBLE_THRESHOLD:-}" ENSEMBLE_SET_INDEX="${ENSEMBLE_SET_INDEX:-1}" K_FOLDS="${K_FOLDS:-}" LOG_DIR="${LOG_DIR:-logs}" @@ -76,6 +82,7 @@ CLI_ARGS=( --output-fasta "${OUTPUT_FASTA}" --model-name "${MODEL_NAME}" --max-tokens "${MAX_TOKENS}" + --seq-len-feature "${SEQ_LEN_FEATURE}" --log-dir "${LOG_DIR}" --log-level "${LOG_LEVEL}" --log-json @@ -83,6 +90,8 @@ CLI_ARGS=( # optional decision threshold passes [ -n "${THRESHOLD}" ] && CLI_ARGS+=(--threshold "${THRESHOLD}") +[ -n "${ENSEMBLE_AGGREGATION}" ] && CLI_ARGS+=(--ensemble-aggregation "${ENSEMBLE_AGGREGATION}") +[ -n "${ENSEMBLE_THRESHOLD}" ] && CLI_ARGS+=(--ensemble-threshold "${ENSEMBLE_THRESHOLD}") [ -n "${ENSEMBLE_SET_INDEX}" ] && CLI_ARGS+=(--ensemble-set-index "${ENSEMBLE_SET_INDEX}") [ -n "${K_FOLDS}" ] && CLI_ARGS+=(--k-folds "${K_FOLDS}") diff --git a/scripts/hpc/trainffnn.sh b/scripts/hpc/train.sh similarity index 56% rename from scripts/hpc/trainffnn.sh rename to scripts/hpc/train.sh index 6517700..0df69aa 100644 --- a/scripts/hpc/trainffnn.sh +++ b/scripts/hpc/train.sh @@ -1,5 +1,5 @@ #!/bin/bash -#SBATCH --job-name=ffnn_v1.0 +#SBATCH --job-name=pepseqpred_train #SBATCH --nodes=1 #SBATCH --ntasks=1 #SBATCH --cpus-per-task=20 @@ -7,8 +7,8 @@ #SBATCH --gpus-per-node=a100:4 #SBATCH --mem=256G #SBATCH --time=12:00:00 -#SBATCH --output=/scratch/%u/train_ffnn_slurm/output/%x_%j.out -#SBATCH --error=/scratch/%u/train_ffnn_slurm/error/%x_%j.err +#SBATCH --output=/scratch/%u/train_slurm/output/%x_%j.out +#SBATCH --error=/scratch/%u/train_slurm/error/%x_%j.err # for testing USE_SRUN="${USE_SRUN:-1}" @@ -26,6 +26,15 @@ usage() { echo " N_FOLDS default: 1 (1=single holdout, >1=K-fold ensemble)" echo " SPLIT_SEEDS default: 11,22,33,44,55" echo " TRAIN_SEEDS default: 101,202,303,404,505" + echo " SUBSET default: 0 (all proteins)" + echo " SPLIT_STRATEGY default: size-balanced" + echo " SPLIT_REPORT_JSON default: unset (/split_report.json)" + echo " THRESHOLD_POLICY default: max-recall-min-precision" + echo " THRESHOLD_MIN_PRECISION default: 0.25" + echo " THRESHOLD_MIN_RECALL default: 0.80" + echo " THRESHOLD_FIXED_VALUE default: 0.50" + echo " MODEL_HEAD default: ffnn (ffnn or conv1d)" + echo " SEQ_LEN_FEATURE default: none (none, raw, inverse)" } # require at least one embedding dir, separator (--), one label shard @@ -57,17 +66,30 @@ fi HIDDEN_SIZES="${HIDDEN_SIZES:-150,120,45}" DROPOUTS="${DROPOUTS:-0.1,0.1,0.1}" +MODEL_HEAD="${MODEL_HEAD:-ffnn}" +SEQ_LEN_FEATURE="${SEQ_LEN_FEATURE:-none}" +CONV_CHANNELS="${CONV_CHANNELS:-64}" +CONV_LAYERS="${CONV_LAYERS:-2}" +CONV_KERNEL_SIZE="${CONV_KERNEL_SIZE:-9}" +CONV_DROPOUT="${CONV_DROPOUT:-0.1}" EPOCHS="${EPOCHS:-10}" BEST_MODEL_METRIC="${BEST_MODEL_METRIC:-pr_auc}" +THRESHOLD_POLICY="${THRESHOLD_POLICY:-max-recall-min-precision}" +THRESHOLD_MIN_PRECISION="${THRESHOLD_MIN_PRECISION:-0.25}" +THRESHOLD_MIN_RECALL="${THRESHOLD_MIN_RECALL:-0.80}" +THRESHOLD_FIXED_VALUE="${THRESHOLD_FIXED_VALUE:-0.50}" SPLIT_SEEDS="${SPLIT_SEEDS:-11,22,33,44,55}" TRAIN_SEEDS="${TRAIN_SEEDS:-101,202,303,404,505}" +SPLIT_STRATEGY="${SPLIT_STRATEGY:-size-balanced}" +SPLIT_REPORT_JSON="${SPLIT_REPORT_JSON:-}" +SUBSET="${SUBSET:-0}" N_FOLDS="${N_FOLDS:-1}" -BATCH_SIZE="${BATCH_SIZE:-256}" # ensure batch size is 4 times what you would do for one GPU (for example. 256 = 64 * 4) +BATCH_SIZE="${BATCH_SIZE:-64}" # ensure you account for number of GPUs (ex. 64 * 4 GPUs = 256 total batch size) LR="${LR:-0.001}" WD="${WD:-0.0}" VAL_FRAC="${VAL_FRAC:-0.2}" -POS_WEIGHT="${POS_WEIGHT:-13.18999647945325}" # calculated from previous script -SAVE_PATH="/scratch/$USER/models/${SLURM_JOB_NAME:-ffnn_job}" +POS_WEIGHT="${POS_WEIGHT:-}" # optional manual override; empty uses train-only auto-compute +SAVE_PATH="/scratch/$USER/models/${SLURM_JOB_NAME:-train_job}" RESULTS_CSV="${SAVE_PATH}/runs.csv" ENSEMBLE_MANIFEST="${ENSEMBLE_MANIFEST:-${SAVE_PATH}/ensemble_manifest.json}" NUM_WORKERS="${NUM_WORKERS:-1}" @@ -111,21 +133,44 @@ if [ "${SAVE_VAL_CURVES}" -eq 1 ]; then VAL_CURVE_ARGS+=(--val-plot-formats "$VAL_PLOT_FORMATS") fi -${LAUNCHER} torchrun --nproc_per_node=4 train_ffnn.pyz \ +POS_WEIGHT_ARGS=() +if [ -n "$POS_WEIGHT" ]; then + POS_WEIGHT_ARGS+=(--pos-weight "$POS_WEIGHT") +fi + +SPLIT_REPORT_ARGS=() +if [ -n "$SPLIT_REPORT_JSON" ]; then + SPLIT_REPORT_ARGS+=(--split-report-json "$SPLIT_REPORT_JSON") +fi + +${LAUNCHER} torchrun --nproc_per_node=4 train.pyz \ --embedding-dirs "${EMBEDDING_DIRS[@]}" \ --label-shards "${LABEL_SHARDS[@]}" \ --label-cache-mode "$LABEL_CACHE_MODE" \ --hidden-sizes "$HIDDEN_SIZES" \ --dropouts "$DROPOUTS" \ + --model-head "$MODEL_HEAD" \ + --seq-len-feature "$SEQ_LEN_FEATURE" \ + --conv-channels "$CONV_CHANNELS" \ + --conv-layers "$CONV_LAYERS" \ + --conv-kernel-size "$CONV_KERNEL_SIZE" \ + --conv-dropout "$CONV_DROPOUT" \ --epochs "$EPOCHS" \ "${TRAIN_ARGS[@]}" \ --batch-size "$BATCH_SIZE" \ --lr "$LR" \ --wd "$WD" \ - --pos-weight "$POS_WEIGHT" \ + "${POS_WEIGHT_ARGS[@]}" \ --best-model-metric "$BEST_MODEL_METRIC" \ + --threshold-policy "$THRESHOLD_POLICY" \ + --threshold-min-precision "$THRESHOLD_MIN_PRECISION" \ + --threshold-min-recall "$THRESHOLD_MIN_RECALL" \ + --threshold-fixed-value "$THRESHOLD_FIXED_VALUE" \ --val-frac "$VAL_FRAC" \ + --subset "$SUBSET" \ --split-type "$SPLIT_TYPE" \ + --split-strategy "$SPLIT_STRATEGY" \ + "${SPLIT_REPORT_ARGS[@]}" \ --save-path "$SAVE_PATH" \ --results-csv "$RESULTS_CSV" \ --num-workers "$NUM_WORKERS" \ @@ -133,4 +178,4 @@ ${LAUNCHER} torchrun --nproc_per_node=4 train_ffnn.pyz \ --stride "$STRIDE" \ "${VAL_CURVE_ARGS[@]}" -# USAGE: sbatch trainffnn.sh /scratch/$USER/esm2/artifacts/pts/shard_000 /scratch/$USER/esm2/artifacts/pts/shard_001 /scratch/$USER/esm2/artifacts/pts/shard_002 /scratch/$USER/esm2/artifacts/pts/shard_003 -- /scratch/$USER/labels/labels_shard_000.pt /scratch/$USER/labels/labels_shard_001.pt /scratch/$USER/labels/labels_shard_002.pt /scratch/$USER/labels/labels_shard_003.pt +# USAGE: sbatch train.sh /scratch/$USER/esm2/artifacts/pts/shard_000 /scratch/$USER/esm2/artifacts/pts/shard_001 /scratch/$USER/esm2/artifacts/pts/shard_002 /scratch/$USER/esm2/artifacts/pts/shard_003 -- /scratch/$USER/labels/labels_shard_000.pt /scratch/$USER/labels/labels_shard_001.pt /scratch/$USER/labels/labels_shard_002.pt /scratch/$USER/labels/labels_shard_003.pt diff --git a/scripts/hpc/trainffnnoptuna.sh b/scripts/hpc/trainoptuna.sh similarity index 56% rename from scripts/hpc/trainffnnoptuna.sh rename to scripts/hpc/trainoptuna.sh index be9f29c..f3d5d15 100644 --- a/scripts/hpc/trainffnnoptuna.sh +++ b/scripts/hpc/trainoptuna.sh @@ -1,5 +1,5 @@ #!/bin/bash -#SBATCH --job-name=ffnn_optuna +#SBATCH --job-name=pepseqpred_optuna #SBATCH --nodes=1 #SBATCH --ntasks=1 #SBATCH --cpus-per-task=20 @@ -7,8 +7,8 @@ #SBATCH --gpus-per-node=a100:4 #SBATCH --mem=448G #SBATCH --time=48:00:00 -#SBATCH --output=/scratch/%u/optuna_ffnn/%j/%x.out -#SBATCH --error=/scratch/%u/optuna_ffnn/%j/%x.err +#SBATCH --output=/scratch/%u/optuna_train/%j/%x.out +#SBATCH --error=/scratch/%u/optuna_train/%j/%x.err # for testing USE_SRUN="${USE_SRUN:-1}" @@ -18,6 +18,13 @@ usage() { echo "Usage: $0 -- " echo " embedding_dirs: one or more directories containing per-protein embeddings (.pt)" echo " label_shards: one or more label shard .pt files" + echo " SPLIT_STRATEGY default: size-balanced" + echo " SPLIT_REPORT_JSON default: unset (/split_report.json)" + echo " THRESHOLD_POLICY default: max-recall-min-precision" + echo " THRESHOLD_MIN_PRECISION default: 0.25" + echo " THRESHOLD_MIN_RECALL default: 0.80" + echo " THRESHOLD_FIXED_VALUE default: 0.50" + echo " SEQ_LEN_FEATURE default: none (none, raw, inverse)" echo "" echo "Example:" echo " $0 /scratch/$USER/embeddings/shard1 /scratch/$USER/embeddings/shard2 -- /scratch/$USER/labels/labels_00.pt /scratch/$USER/labels/labels_01.pt" @@ -51,17 +58,23 @@ if [ "${#EMBEDDING_DIRS[@]}" -eq 0 ] || [ "${#LABEL_SHARDS[@]}" -eq 0 ]; then fi # tuning controls -STUDY_NAME="${STUDY_NAME:-ffnn_optuna_v1}" +STUDY_NAME="${STUDY_NAME:-train_optuna_v1}" N_TRIALS="${N_TRIALS:-20}" EPOCHS="${EPOCHS:-15}" SEED="${SEED:-42}" METRIC="${METRIC:-pr_auc}" +THRESHOLD_POLICY="${THRESHOLD_POLICY:-max-recall-min-precision}" +THRESHOLD_MIN_PRECISION="${THRESHOLD_MIN_PRECISION:-0.25}" +THRESHOLD_MIN_RECALL="${THRESHOLD_MIN_RECALL:-0.80}" +THRESHOLD_FIXED_VALUE="${THRESHOLD_FIXED_VALUE:-0.50}" VAL_FRAC="${VAL_FRAC:-0.2}" +SPLIT_STRATEGY="${SPLIT_STRATEGY:-size-balanced}" +SPLIT_REPORT_JSON="${SPLIT_REPORT_JSON:-}" SUBSET="${SUBSET:-0}" NUM_WORKERS="${NUM_WORKERS:-1}" WINDOW_SIZE="${WINDOW_SIZE:-1000}" STRIDE="${STRIDE:-900}" -POS_WEIGHT="${POS_WEIGHT:-13.18999647945325}" # calculated from other script +POS_WEIGHT="${POS_WEIGHT:-}" # optional manual override; empty uses train-only auto-compute SPLIT_TYPE="${SPLIT_TYPE:-id-family}" # id-family or id # output paths @@ -72,11 +85,19 @@ CSV_PATH="${CSV_PATH:-/scratch/$USER/optuna/${STUDY_NAME}_trials.csv}" STORAGE="${STORAGE:-sqlite:////scratch/$USER/optuna/${STUDY_NAME}.db}" # architecture search space +MODEL_HEAD="${MODEL_HEAD:-ffnn}" # ffnn or conv1d +SEQ_LEN_FEATURE="${SEQ_LEN_FEATURE:-none}" # none, raw, inverse ARCH_MODE="${ARCH_MODE:-flat}" # flat, bottleneck, pyramid DEPTH_MIN="${DEPTH_MIN:-2}" DEPTH_MAX="${DEPTH_MAX:-6}" WIDTH_MIN="${WIDTH_MIN:-64}" WIDTH_MAX="${WIDTH_MAX:-512}" +CONV_CHANNEL_CHOICES="${CONV_CHANNEL_CHOICES:-32,64,128}" +CONV_LAYERS_MIN="${CONV_LAYERS_MIN:-1}" +CONV_LAYERS_MAX="${CONV_LAYERS_MAX:-3}" +CONV_KERNEL_SIZE_CHOICES="${CONV_KERNEL_SIZE_CHOICES:-3,5,9,15}" +CONV_DROPOUT_MIN="${CONV_DROPOUT_MIN:-0.0}" +CONV_DROPOUT_MAX="${CONV_DROPOUT_MAX:-0.25}" # optimizer search space LR_MIN="${LR_MIN:-1e-4}" @@ -93,7 +114,7 @@ TIMEOUT_S="${TIMEOUT_S:-0}" # timeout in seconds mkdir -p "$SAVE_PATH" mkdir -p "$(dirname "$CSV_PATH")" mkdir -p "/scratch/$USER/optuna" -mkdir -p "/scratch/$USER/optuna_ffnn/$SLURM_JOB_ID" +mkdir -p "/scratch/$USER/optuna_train/$SLURM_JOB_ID" # load Python Conda environment module purge @@ -113,7 +134,17 @@ fi DDP_TIMEOUT_MIN="${DDP_TIMEOUT_MIN:-60}" export PEPSEQPRED_DDP_TIMEOUT_MIN="$DDP_TIMEOUT_MIN" -${LAUNCHER} torchrun --nproc_per_node=4 train_ffnn_optuna.pyz \ +POS_WEIGHT_ARGS=() +if [ -n "$POS_WEIGHT" ]; then + POS_WEIGHT_ARGS+=(--pos-weight "$POS_WEIGHT") +fi + +SPLIT_REPORT_ARGS=() +if [ -n "$SPLIT_REPORT_JSON" ]; then + SPLIT_REPORT_ARGS+=(--split-report-json "$SPLIT_REPORT_JSON") +fi + +${LAUNCHER} torchrun --nproc_per_node=4 train_optuna.pyz \ --embedding-dirs "${EMBEDDING_DIRS[@]}" \ --label-shards "${LABEL_SHARDS[@]}" \ --study-name "$STUDY_NAME" \ @@ -122,8 +153,16 @@ ${LAUNCHER} torchrun --nproc_per_node=4 train_ffnn_optuna.pyz \ --epochs "$EPOCHS" \ --seed "$SEED" \ --metric "$METRIC" \ + --model-head "$MODEL_HEAD" \ + --seq-len-feature "$SEQ_LEN_FEATURE" \ + --threshold-policy "$THRESHOLD_POLICY" \ + --threshold-min-precision "$THRESHOLD_MIN_PRECISION" \ + --threshold-min-recall "$THRESHOLD_MIN_RECALL" \ + --threshold-fixed-value "$THRESHOLD_FIXED_VALUE" \ --val-frac "$VAL_FRAC" \ --split-type "$SPLIT_TYPE" \ + --split-strategy "$SPLIT_STRATEGY" \ + "${SPLIT_REPORT_ARGS[@]}" \ --subset "$SUBSET" \ --num-workers "$NUM_WORKERS" \ --save-path "$SAVE_PATH" \ @@ -133,15 +172,21 @@ ${LAUNCHER} torchrun --nproc_per_node=4 train_ffnn_optuna.pyz \ --depth-max "$DEPTH_MAX" \ --width-min "$WIDTH_MIN" \ --width-max "$WIDTH_MAX" \ + --conv-channel-choices "$CONV_CHANNEL_CHOICES" \ + --conv-layers-min "$CONV_LAYERS_MIN" \ + --conv-layers-max "$CONV_LAYERS_MAX" \ + --conv-kernel-size-choices "$CONV_KERNEL_SIZE_CHOICES" \ + --conv-dropout-min "$CONV_DROPOUT_MIN" \ + --conv-dropout-max "$CONV_DROPOUT_MAX" \ --batch-sizes "$BATCH_SIZES" \ --lr-min "$LR_MIN" \ --lr-max "$LR_MAX" \ --wd-min "$WD_MIN" \ --wd-max "$WD_MAX" \ - --pos-weight "$POS_WEIGHT" \ + "${POS_WEIGHT_ARGS[@]}" \ --pruner-warmup "$PRUNER_WARMUP" \ --timeout-s "$TIMEOUT_S" \ --window-size "$WINDOW_SIZE" \ --stride "$STRIDE" -# USAGE: sbatch trainffnnoptuna.sh /scratch/$USER/esm2/artifacts/pts/shard_000 /scratch/$USER/esm2/artifacts/pts/shard_001 /scratch/$USER/esm2/artifacts/pts/shard_002 /scratch/$USER/esm2/artifacts/pts/shard_003 -- /scratch/$USER/labels/labels_shard_000.pt /scratch/$USER/labels/labels_shard_001.pt /scratch/$USER/labels/labels_shard_002.pt /scratch/$USER/labels/labels_shard_003.pt +# USAGE: sbatch trainoptuna.sh /scratch/$USER/esm2/artifacts/pts/shard_000 /scratch/$USER/esm2/artifacts/pts/shard_001 /scratch/$USER/esm2/artifacts/pts/shard_002 /scratch/$USER/esm2/artifacts/pts/shard_003 -- /scratch/$USER/labels/labels_shard_000.pt /scratch/$USER/labels/labels_shard_001.pt /scratch/$USER/labels/labels_shard_002.pt /scratch/$USER/labels/labels_shard_003.pt diff --git a/scripts/tools/buildpyz.py b/scripts/tools/buildpyz.py index e0f6210..8b8b764 100644 --- a/scripts/tools/buildpyz.py +++ b/scripts/tools/buildpyz.py @@ -9,7 +9,6 @@ import argparse import shutil import subprocess -import sys import zipapp from pathlib import Path from pyzapps import APPS diff --git a/scripts/tools/pyzapps.py b/scripts/tools/pyzapps.py index 20b7248..952b335 100644 --- a/scripts/tools/pyzapps.py +++ b/scripts/tools/pyzapps.py @@ -13,8 +13,8 @@ "esm": "pepseqpred.apps.esm_cli:main", "labels": "pepseqpred.apps.labels_cli:main", "preprocess": "pepseqpred.apps.preprocess_cli:main", - "train_ffnn": "pepseqpred.apps.train_ffnn_cli:main", - "train_ffnn_optuna": "pepseqpred.apps.train_ffnn_optuna_cli:main", + "train": "pepseqpred.apps.train_cli:main", + "train_optuna": "pepseqpred.apps.train_optuna_cli:main", "predict": "pepseqpred.apps.prediction_cli:main", "evaluate_ffnn": "pepseqpred.apps.evaluate_ffnn_cli:main" } diff --git a/src/pepseqpred/api/predictor.py b/src/pepseqpred/api/predictor.py index 771efac..3b6f3d9 100644 --- a/src/pepseqpred/api/predictor.py +++ b/src/pepseqpred/api/predictor.py @@ -21,7 +21,8 @@ embed_protein_seq, infer_decision_threshold, predict_ensemble_from_embedding, - predict_from_embedding + predict_from_embedding, + resolve_prediction_seq_len_feature, ) _DEFAULT_ESM_MODEL = "esm2_t33_650M_UR50D" @@ -126,6 +127,8 @@ class PepSeqPredictor: Name of the ESM backbone used by this predictor instance. max_tokens : int Maximum residue token budget per ESM pass (excluding CLS/EOS). + seq_len_feature : str + Resolved sequence-length feature mode used for generated embeddings. artifact_mode : str Artifact resolution mode (`single-checkpoint` or `ensemble-manifest`). artifact_meta : Mapping[str, Any] @@ -157,6 +160,7 @@ def __init__( max_tokens: int, artifact_mode: str, artifact_meta: Mapping[str, Any], + seq_len_feature: str = "none", pretrained_meta: Optional[Mapping[str, Any]] = None ) -> None: """Initialize predictor state from already loaded models and metadata.""" @@ -170,6 +174,7 @@ def __init__( self._max_tokens = int(max_tokens) self._artifact_mode = artifact_mode self._artifact_meta = dict(artifact_meta) + self._seq_len_feature = seq_len_feature self._pretrained_meta = ( dict(pretrained_meta) if pretrained_meta is not None else {} ) @@ -185,7 +190,8 @@ def from_artifact( model_name: str = _DEFAULT_ESM_MODEL, max_tokens: int = 1022, device: str = "auto", - model_config: Optional[FFNNModelConfig] = None + model_config: Optional[FFNNModelConfig] = None, + seq_len_feature: str = "auto", ) -> "PepSeqPredictor": """Build a predictor from a checkpoint `.pt` or manifest `.json` artifact. @@ -207,6 +213,9 @@ def from_artifact( Device selector string or `"auto"`. model_config : FFNNModelConfig | None Optional explicit FFNN architecture config for checkpoint loading. + seq_len_feature : str + Sequence-length feature override. `"auto"` resolves from checkpoint + model_config, with a missing key treated as `"none"`. Returns ------- @@ -286,6 +295,16 @@ def from_artifact( f"All ensemble members must share emb_dim for shared-embedding inference, " f"got {sorted(emb_dims)}" ) + seq_len_features = { + resolve_prediction_seq_len_feature(seq_len_feature, cfg) + for cfg in member_model_cfgs + } + if len(seq_len_features) != 1: + raise ValueError( + "All ensemble members must share seq_len_feature for shared-embedding inference, " + f"got {sorted(seq_len_features)}" + ) + resolved_seq_len_feature = next(iter(seq_len_features)) return cls( psp_models=psp_models, @@ -297,7 +316,8 @@ def from_artifact( model_name=model_name, max_tokens=max_tokens, artifact_mode=artifact_mode, - artifact_meta=artifact_meta + artifact_meta=artifact_meta, + seq_len_feature=resolved_seq_len_feature, ) @classmethod @@ -309,7 +329,8 @@ def from_pretrained( k_folds: Optional[int] = None, max_tokens: int = 1022, device: str = "auto", - model_config: Optional[FFNNModelConfig] = None + model_config: Optional[FFNNModelConfig] = None, + seq_len_feature: str = "auto", ) -> "PepSeqPredictor": """Build a predictor from a bundled pretrained model id or alias. @@ -327,6 +348,9 @@ def from_pretrained( Device selector string or `"auto"`. model_config : FFNNModelConfig | None Optional explicit FFNN architecture config for checkpoint loading. + seq_len_feature : str + Sequence-length feature override. `"auto"` resolves from checkpoint + model_config, with a missing key treated as `"none"`. Returns ------- @@ -349,7 +373,8 @@ def from_pretrained( model_name=info.expected_esm_model, max_tokens=max_tokens, device=device, - model_config=model_config + model_config=model_config, + seq_len_feature=seq_len_feature, ) predictor._pretrained_meta = { @@ -458,7 +483,8 @@ def _payload_to_result( "artifact_meta": dict(self._artifact_meta), "model_name": self._model_name, "device": self._device, - "max_tokens": self._max_tokens + "max_tokens": self._max_tokens, + "seq_len_feature": self._seq_len_feature, } if self._pretrained_meta: meta["pretrained"] = dict(self._pretrained_meta) @@ -518,7 +544,8 @@ def predict_sequence( layer=self._layer, batch_converter=self._batch_converter, device=self._device, - max_tokens=self._max_tokens + max_tokens=self._max_tokens, + seq_len_feature=self._seq_len_feature, ) payload, used_thresholds = self._predict_from_embedding( protein_emb=protein_emb, @@ -619,7 +646,8 @@ def load_predictor( model_name: str = _DEFAULT_ESM_MODEL, max_tokens: int = 1022, device: str = "auto", - model_config: Optional[FFNNModelConfig] = None + model_config: Optional[FFNNModelConfig] = None, + seq_len_feature: str = "auto", ) -> PepSeqPredictor: """Create a predictor from an artifact path. @@ -641,6 +669,9 @@ def load_predictor( Device selector string or `"auto"`. model_config : FFNNModelConfig | None Optional explicit FFNN architecture config. + seq_len_feature : str + Sequence-length feature override. `"auto"` resolves from checkpoint + model_config, with a missing key treated as `"none"`. Returns ------- @@ -663,6 +694,7 @@ def load_predictor( max_tokens=max_tokens, device=device, model_config=model_config, + seq_len_feature=seq_len_feature, ) @@ -689,7 +721,8 @@ def load_pretrained_predictor( k_folds: Optional[int] = None, max_tokens: int = 1022, device: str = "auto", - model_config: Optional[FFNNModelConfig] = None + model_config: Optional[FFNNModelConfig] = None, + seq_len_feature: str = "auto", ) -> PepSeqPredictor: """Create a predictor from a bundled pretrained model id or alias. @@ -707,6 +740,9 @@ def load_pretrained_predictor( Device selector string or `"auto"`. model_config : FFNNModelConfig | None Optional explicit FFNN architecture config. + seq_len_feature : str + Sequence-length feature override. `"auto"` resolves from checkpoint + model_config, with a missing key treated as `"none"`. Returns ------- @@ -726,7 +762,8 @@ def load_pretrained_predictor( k_folds=k_folds, max_tokens=max_tokens, device=device, - model_config=model_config + model_config=model_config, + seq_len_feature=seq_len_feature, ) @@ -741,7 +778,8 @@ def predict_sequence( model_name: str = _DEFAULT_ESM_MODEL, max_tokens: int = 1022, device: str = "auto", - model_config: Optional[FFNNModelConfig] = None + model_config: Optional[FFNNModelConfig] = None, + seq_len_feature: str = "auto", ) -> PredictionResult: """Convenience wrapper that loads a predictor and predicts one sequence. @@ -767,6 +805,9 @@ def predict_sequence( Device selector string or `"auto"`. model_config : FFNNModelConfig | None Optional explicit FFNN architecture config. + seq_len_feature : str + Sequence-length feature override. `"auto"` resolves from checkpoint + model_config, with a missing key treated as `"none"`. Returns ------- @@ -788,7 +829,8 @@ def predict_sequence( model_name=model_name, max_tokens=max_tokens, device=device, - model_config=model_config + model_config=model_config, + seq_len_feature=seq_len_feature, ) return predictor.predict_sequence( protein_seq=protein_seq, @@ -808,7 +850,8 @@ def predict_fasta( model_name: str = _DEFAULT_ESM_MODEL, max_tokens: int = 1022, device: str = "auto", - model_config: Optional[FFNNModelConfig] = None + model_config: Optional[FFNNModelConfig] = None, + seq_len_feature: str = "auto", ) -> List[PredictionResult]: """Convenience wrapper to predict FASTA input with one artifact load. @@ -834,6 +877,9 @@ def predict_fasta( Device selector string or `"auto"`. model_config : FFNNModelConfig | None Optional explicit FFNN architecture config. + seq_len_feature : str + Sequence-length feature override. `"auto"` resolves from checkpoint + model_config, with a missing key treated as `"none"`. Returns ------- @@ -855,7 +901,8 @@ def predict_fasta( model_name=model_name, max_tokens=max_tokens, device=device, - model_config=model_config + model_config=model_config, + seq_len_feature=seq_len_feature, ) if output_fasta is None: return predictor.predict_fasta( diff --git a/src/pepseqpred/apps/esm_cli.py b/src/pepseqpred/apps/esm_cli.py index b28aec4..5f45035 100644 --- a/src/pepseqpred/apps/esm_cli.py +++ b/src/pepseqpred/apps/esm_cli.py @@ -33,6 +33,7 @@ normalize_family_value ) from pepseqpred.core.embeddings.esm2 import esm_embeddings_from_fasta +from pepseqpred.core.data.seq_len_feature import EMBEDDING_SEQ_LEN_FEATURES def main() -> None: @@ -139,6 +140,13 @@ def main() -> None: type=int, default=1022, help="Size of context window for ESM model.") + parser.add_argument("--seq-len-feature", + action="store", + dest="seq_len_feature", + type=str, + choices=list(EMBEDDING_SEQ_LEN_FEATURES), + default="none", + help="Optional sequence-length feature to append to residue embeddings.") parser.add_argument("-b", "--batch-size", action="store", dest="batch_size", @@ -283,6 +291,7 @@ def main() -> None: "metadata_file": str(metadata_file) if metadata_file is not None else None, "embedding_key_mode": args.embedding_key_mode, "key_delimiter": args.key_delimiter, + "seq_len_feature": args.seq_len_feature, "output_path": str(os.path.abspath(out_dir)), "device": "cuda" if torch.cuda.is_available() else "cpu", "torch_version": torch.__version__, @@ -305,6 +314,7 @@ def main() -> None: index_csv_path=out_dir/idx_csv_path, key_mode=args.embedding_key_mode, key_delimiter=args.key_delimiter, + seq_len_feature=args.seq_len_feature, logger=logger ) logger.info("run_end", extra={"extra": { diff --git a/src/pepseqpred/apps/evaluate_ffnn_cli.py b/src/pepseqpred/apps/evaluate_ffnn_cli.py index b4baf18..e2ffc50 100644 --- a/src/pepseqpred/apps/evaluate_ffnn_cli.py +++ b/src/pepseqpred/apps/evaluate_ffnn_cli.py @@ -18,6 +18,7 @@ from pepseqpred.core.io.logger import setup_logger from pepseqpred.core.io.read import parse_float_csv, parse_int_csv from pepseqpred.core.data.proteindataset import ProteinDataset, pad_collate +from pepseqpred.core.models.factory import MODEL_HEADS from pepseqpred.core.predict.inference import ( FFNNModelConfig, build_model_from_checkpoint, @@ -25,6 +26,7 @@ predict_member_probabilities_from_embedding ) from pepseqpred.core.train.metrics import compute_eval_metrics +from pepseqpred.core.train.threshold import threshold_diagnostic_grid @dataclass(frozen=True) @@ -159,6 +161,11 @@ def _resolve_best_set_index( def _build_cli_model_config(args: argparse.Namespace) -> FFNNModelConfig | None: """Builds the model configuration from explicit CLI architecture flags.""" + model_head_arg = getattr(args, "model_head", None) + conv_channels_arg = getattr(args, "conv_channels", None) + conv_layers_arg = getattr(args, "conv_layers", None) + conv_kernel_size_arg = getattr(args, "conv_kernel_size", None) + conv_dropout_arg = getattr(args, "conv_dropout", None) any_explicit = any( arg is not None for arg in ( @@ -167,7 +174,12 @@ def _build_cli_model_config(args: argparse.Namespace) -> FFNNModelConfig | None: args.dropouts, args.use_layer_norm, args.use_residual, - args.num_classes + args.num_classes, + model_head_arg, + conv_channels_arg, + conv_layers_arg, + conv_kernel_size_arg, + conv_dropout_arg ) ) if not any_explicit: @@ -184,6 +196,16 @@ def _build_cli_model_config(args: argparse.Namespace) -> FFNNModelConfig | None: missing.append("--use-layer-norm/--no-use-layer-norm") if args.use_residual is None: missing.append("--use-residual/--no-use-residual") + model_head = str(model_head_arg or "ffnn") + if model_head == "conv1d": + if conv_channels_arg is None: + missing.append("--conv-channels") + if conv_layers_arg is None: + missing.append("--conv-layers") + if conv_kernel_size_arg is None: + missing.append("--conv-kernel-size") + if conv_dropout_arg is None: + missing.append("--conv-dropout") if missing: raise ValueError( "When using explicit architecture flags, provide all required values: " @@ -203,7 +225,12 @@ def _build_cli_model_config(args: argparse.Namespace) -> FFNNModelConfig | None: dropouts=tuple(dropouts), num_classes=num_classes, use_layer_norm=bool(args.use_layer_norm), - use_residual=bool(args.use_residual) + use_residual=bool(args.use_residual), + model_head=model_head, + conv_channels=int(conv_channels_arg) if conv_channels_arg is not None else 64, + conv_layers=int(conv_layers_arg) if conv_layers_arg is not None else 2, + conv_kernel_size=int(conv_kernel_size_arg) if conv_kernel_size_arg is not None else 9, + conv_dropout=float(conv_dropout_arg) if conv_dropout_arg is not None else 0.1 ) @@ -538,7 +565,12 @@ def _compute_pr_auc_trapezoid(y_true: np.ndarray, y_prob: np.ndarray) -> float: return float(np.trapezoid(precision, recall)) -def _empty_eval_output(votes_needed: int | None, include_curves: bool) -> Dict[str, Any]: +def _empty_eval_output( + votes_needed: int | None, + include_curves: bool, + ensemble_aggregation: str = "majority", + ensemble_threshold: float | None = None, +) -> Dict[str, Any]: """Returns a canonical empty evaluation payload.""" out = { "processed_proteins": 0, @@ -552,6 +584,11 @@ def _empty_eval_output(votes_needed: int | None, include_curves: bool) -> Dict[s "accuracy": float("nan"), "confusion_matrix": [[0, 0], [0, 0]], "votes_needed": votes_needed, + "ensemble_aggregation": str(ensemble_aggregation), + "ensemble_threshold": ( + float(ensemble_threshold) if ensemble_threshold is not None else None + ), + "threshold_grid": [], "has_both_classes": False, "metrics": { "precision": float("nan"), @@ -598,12 +635,18 @@ def _build_eval_output_from_arrays( proteins_zero_valid_residues: int, votes_needed: int | None, include_curves: bool, - curve_max_points: int + curve_max_points: int, + ensemble_aggregation: str = "majority", + ensemble_threshold: float | None = None, ) -> Dict[str, Any]: """Builds a full evaluation summary from flattened residue-level arrays.""" if y_true.size < 1: out = _empty_eval_output( - votes_needed=votes_needed, include_curves=include_curves) + votes_needed=votes_needed, + include_curves=include_curves, + ensemble_aggregation=ensemble_aggregation, + ensemble_threshold=ensemble_threshold, + ) out["processed_proteins"] = int(processed_proteins) out["proteins_zero_valid_residues"] = int(proteins_zero_valid_residues) out["proteins_with_valid_residues"] = int( @@ -633,6 +676,11 @@ def _build_eval_output_from_arrays( "accuracy": float((y_true == y_pred).mean()) if y_true.size > 0 else float("nan"), "confusion_matrix": [[int(cm[0, 0]), int(cm[0, 1])], [int(cm[1, 0]), int(cm[1, 1])]], "votes_needed": votes_needed, + "ensemble_aggregation": str(ensemble_aggregation), + "ensemble_threshold": ( + float(ensemble_threshold) if ensemble_threshold is not None else None + ), + "threshold_grid": list(threshold_diagnostic_grid(y_true, y_prob)), "has_both_classes": bool(metrics.get("has_both_classes", False)), "metrics": metrics } @@ -977,9 +1025,14 @@ def _evaluate_dataset( best_fold_by: str, best_fold_direction: str, plot_dir: Path | None, - plot_formats: Sequence[str] + plot_formats: Sequence[str], + ensemble_aggregation: str = "majority", + ensemble_threshold: float | None = None, ) -> Dict[str, Any]: """Evaluates one model or an ensemble on the given eval dataset loader.""" + ensemble_aggregation = str(ensemble_aggregation).strip().lower() + if ensemble_aggregation not in {"majority", "mean-prob"}: + raise ValueError("--ensemble-aggregation must be one of: majority, mean-prob") n_members = int(len(psp_models)) if n_members < 1: raise ValueError("At least one model is required for evaluation") @@ -994,7 +1047,18 @@ def _evaluate_dataset( if curve_max_points < 2: raise ValueError("--curve-max-points must be >= 2") - votes_needed = int((n_members // 2) + 1) + resolved_ensemble_threshold: float | None = None + votes_needed: int | None = None + if n_members > 1 and ensemble_aggregation == "majority": + votes_needed = int((n_members // 2) + 1) + elif n_members > 1: + resolved_ensemble_threshold = ( + float(ensemble_threshold) + if ensemble_threshold is not None + else float(sum(float(x) for x in thresholds) / len(thresholds)) + ) + if resolved_ensemble_threshold <= 0.0 or resolved_ensemble_threshold >= 1.0: + raise ValueError("--ensemble-threshold must be between (0.0, 1.0)") all_true: List[torch.Tensor] = [] ens_pred: List[torch.Tensor] = [] @@ -1036,11 +1100,17 @@ def _evaluate_dataset( if n_members == 1: pred_ensemble = member_pred_full[0] prob_ensemble = member_probs_full[0] - else: + elif ensemble_aggregation == "majority": vote_sum = torch.stack(member_pred_full, dim=0).sum(dim=0) pred_ensemble = (vote_sum >= votes_needed).to(torch.int64) prob_ensemble = torch.stack( member_probs_full, dim=0).mean(dim=0) + else: + prob_ensemble = torch.stack( + member_probs_full, dim=0).mean(dim=0) + pred_ensemble = ( + prob_ensemble >= float(resolved_ensemble_threshold) + ).to(torch.int64) y_true_valid = y_row[valid_mask].to(torch.int64).cpu() ens_pred_valid = pred_ensemble[valid_mask].to(torch.int64).cpu() @@ -1097,9 +1167,13 @@ def _evaluate_dataset( y_prob=ens_prob_np, processed_proteins=int(processed), proteins_zero_valid_residues=int(zero_valid), - votes_needed=(votes_needed if n_members > 1 else None), + votes_needed=votes_needed, include_curves=include_curves, - curve_max_points=curve_max_points + curve_max_points=curve_max_points, + ensemble_aggregation=( + ensemble_aggregation if n_members > 1 else "single-model" + ), + ensemble_threshold=resolved_ensemble_threshold, ) if not emit_fold_metrics: @@ -1125,7 +1199,9 @@ def _evaluate_dataset( proteins_zero_valid_residues=int(zero_valid), votes_needed=None, include_curves=include_curves, - curve_max_points=curve_max_points + curve_max_points=curve_max_points, + ensemble_aggregation="member", + ensemble_threshold=float(thresholds[member_idx]), ) fold_eval["fold_index"] = member.fold_index fold_eval["member_index"] = member.member_index @@ -1277,6 +1353,23 @@ def main() -> None: default=None, help="Optional global threshold override in (0.0, 1.0)." ) + parser.add_argument( + "--ensemble-aggregation", + action="store", + dest="ensemble_aggregation", + type=str, + choices=["majority", "mean-prob"], + default="majority", + help="Ensemble aggregation rule for manifest evaluation." + ) + parser.add_argument( + "--ensemble-threshold", + action="store", + dest="ensemble_threshold", + type=float, + default=None, + help="Optional threshold for mean-prob ensemble aggregation." + ) parser.add_argument( "--ensemble-set-index", action="store", @@ -1474,6 +1567,47 @@ def main() -> None: default=None, help="Explicit output classes (binary=1)." ) + parser.add_argument( + "--model-head", + action="store", + dest="model_head", + type=str, + choices=list(MODEL_HEADS), + default=None, + help="Explicit model head used in training." + ) + parser.add_argument( + "--conv-channels", + action="store", + dest="conv_channels", + type=int, + default=None, + help="Explicit Conv1d channels used in training." + ) + parser.add_argument( + "--conv-layers", + action="store", + dest="conv_layers", + type=int, + default=None, + help="Explicit Conv1d layer count used in training." + ) + parser.add_argument( + "--conv-kernel-size", + action="store", + dest="conv_kernel_size", + type=int, + default=None, + help="Explicit Conv1d kernel size used in training." + ) + parser.add_argument( + "--conv-dropout", + action="store", + dest="conv_dropout", + type=float, + default=None, + help="Explicit Conv1d dropout used in training." + ) parser.add_argument( "--use-layer-norm", action="store_true", @@ -1507,6 +1641,10 @@ def main() -> None: raise ValueError("--k-folds must be >= 1") if args.threshold is not None and (args.threshold <= 0.0 or args.threshold >= 1.0): raise ValueError("--threshold must be between (0.0, 1.0)") + if args.ensemble_threshold is not None and ( + args.ensemble_threshold <= 0.0 or args.ensemble_threshold >= 1.0 + ): + raise ValueError("--ensemble-threshold must be between (0.0, 1.0)") if args.subset < 0: raise ValueError("--subset must be >= 0") if args.batch_size < 1: @@ -1595,6 +1733,18 @@ def main() -> None: f"got {sorted(emb_dims)}" ) + resolved_ensemble_threshold = None + if len(psp_models) > 1 and args.ensemble_aggregation == "mean-prob": + resolved_ensemble_threshold = ( + float(args.ensemble_threshold) + if args.ensemble_threshold is not None + else ( + float(args.threshold) + if args.threshold is not None + else float(sum(member_thresholds) / len(member_thresholds)) + ) + ) + base_dataset = ProteinDataset( embedding_dirs=args.embedding_dirs, label_shards=args.label_shards, @@ -1673,16 +1823,23 @@ def main() -> None: ), "threshold": float(member_thresholds[0]) if len(member_thresholds) == 1 else None, "member_thresholds": [float(x) for x in member_thresholds], + "ensemble_aggregation": str(args.ensemble_aggregation), + "ensemble_threshold": resolved_ensemble_threshold, "device": device, "model_cfg_src": ( str(member_model_cfg_srcs[0]) if len(set(member_model_cfg_srcs)) == 1 else "mixed" ), + "model_head": str(first_cfg.model_head), "emb_dim": int(first_cfg.emb_dim), "hidden_sizes": [int(x) for x in first_cfg.hidden_sizes], "use_layer_norm": bool(first_cfg.use_layer_norm), "use_residual": bool(first_cfg.use_residual), + "conv_channels": int(first_cfg.conv_channels), + "conv_layers": int(first_cfg.conv_layers), + "conv_kernel_size": int(first_cfg.conv_kernel_size), + "conv_dropout": float(first_cfg.conv_dropout), "num_classes": int(first_cfg.num_classes), "emit_fold_metrics": bool(args.emit_fold_metrics), "include_curves": bool(args.include_curves), @@ -1710,7 +1867,9 @@ def main() -> None: best_fold_by=str(args.best_fold_by), best_fold_direction=str(args.best_fold_direction), plot_dir=args.plot_dir, - plot_formats=plot_formats + plot_formats=plot_formats, + ensemble_aggregation=str(args.ensemble_aggregation), + ensemble_threshold=resolved_ensemble_threshold ) if not bool(eval_out.get("has_both_classes", False)): logger.warning( @@ -1761,6 +1920,8 @@ def main() -> None: ), "threshold": float(member_thresholds[0]) if len(member_thresholds) == 1 else None, "member_thresholds": [float(x) for x in member_thresholds], + "ensemble_aggregation": str(args.ensemble_aggregation), + "ensemble_threshold": resolved_ensemble_threshold, "emit_fold_metrics": bool(args.emit_fold_metrics), "include_curves": bool(args.include_curves), "curve_max_points": int(args.curve_max_points), @@ -1779,10 +1940,15 @@ def main() -> None: if len(set(member_model_cfg_srcs)) == 1 else "mixed" ), + "model_head": str(first_cfg.model_head), "emb_dim": int(first_cfg.emb_dim), "hidden_sizes": [int(x) for x in first_cfg.hidden_sizes], "use_layer_norm": bool(first_cfg.use_layer_norm), "use_residual": bool(first_cfg.use_residual), + "conv_channels": int(first_cfg.conv_channels), + "conv_layers": int(first_cfg.conv_layers), + "conv_kernel_size": int(first_cfg.conv_kernel_size), + "conv_dropout": float(first_cfg.conv_dropout), "num_classes": int(first_cfg.num_classes), "evaluation": eval_out } diff --git a/src/pepseqpred/apps/prediction_cli.py b/src/pepseqpred/apps/prediction_cli.py index 403e7d7..e9d007d 100644 --- a/src/pepseqpred/apps/prediction_cli.py +++ b/src/pepseqpred/apps/prediction_cli.py @@ -24,14 +24,17 @@ from pepseqpred.core.io.logger import setup_logger from pepseqpred.core.io.read import parse_int_csv, parse_float_csv from pepseqpred.core.embeddings.esm2 import clean_seq +from pepseqpred.core.models.factory import MODEL_HEADS from pepseqpred.core.predict.inference import ( FFNNModelConfig, build_model_from_checkpoint, embed_protein_seq, infer_decision_threshold, predict_ensemble_from_embedding, - predict_from_embedding + predict_from_embedding, + resolve_prediction_seq_len_feature ) +from pepseqpred.core.data.seq_len_feature import PREDICTION_SEQ_LEN_FEATURES @dataclass(frozen=True) @@ -68,6 +71,11 @@ def read_fasta_records(fasta_path: Path | str) -> Iterator[Tuple[str, str]]: def _build_cli_model_config(args: argparse.Namespace) -> FFNNModelConfig | None: """Builds the model configuration using CLI arguments.""" + model_head_arg = getattr(args, "model_head", None) + conv_channels_arg = getattr(args, "conv_channels", None) + conv_layers_arg = getattr(args, "conv_layers", None) + conv_kernel_size_arg = getattr(args, "conv_kernel_size", None) + conv_dropout_arg = getattr(args, "conv_dropout", None) any_explicit = any( arg is not None for arg in ( @@ -76,7 +84,12 @@ def _build_cli_model_config(args: argparse.Namespace) -> FFNNModelConfig | None: args.dropouts, args.use_layer_norm, args.use_residual, - args.num_classes + args.num_classes, + model_head_arg, + conv_channels_arg, + conv_layers_arg, + conv_kernel_size_arg, + conv_dropout_arg ) ) if not any_explicit: @@ -93,6 +106,16 @@ def _build_cli_model_config(args: argparse.Namespace) -> FFNNModelConfig | None: missing.append("--use-layer-norm/--no-use-layer-norm") if args.use_residual is None: missing.append("--use-residual/--no-use-residual") + model_head = str(model_head_arg or "ffnn") + if model_head == "conv1d": + if conv_channels_arg is None: + missing.append("--conv-channels") + if conv_layers_arg is None: + missing.append("--conv-layers") + if conv_kernel_size_arg is None: + missing.append("--conv-kernel-size") + if conv_dropout_arg is None: + missing.append("--conv-dropout") if missing: raise ValueError( "When using explicit architecture flags, provide all required values: " @@ -113,7 +136,12 @@ def _build_cli_model_config(args: argparse.Namespace) -> FFNNModelConfig | None: dropouts=tuple(dropouts), num_classes=num_classes, use_layer_norm=bool(args.use_layer_norm), - use_residual=bool(args.use_residual) + use_residual=bool(args.use_residual), + model_head=model_head, + conv_channels=int(conv_channels_arg) if conv_channels_arg is not None else 64, + conv_layers=int(conv_layers_arg) if conv_layers_arg is not None else 2, + conv_kernel_size=int(conv_kernel_size_arg) if conv_kernel_size_arg is not None else 9, + conv_dropout=float(conv_dropout_arg) if conv_dropout_arg is not None else 0.1 ) @@ -301,6 +329,19 @@ def main() -> None: type=float, default=None, help="Optional global threshold override within (0.0, 1.0).") + parser.add_argument("--ensemble-aggregation", + action="store", + dest="ensemble_aggregation", + type=str, + choices=["majority", "mean-prob"], + default="majority", + help="Ensemble aggregation rule for manifest predictions.") + parser.add_argument("--ensemble-threshold", + action="store", + dest="ensemble_threshold", + type=float, + default=None, + help="Optional threshold for mean-prob ensemble aggregation.") parser.add_argument("--ensemble-set-index", action="store", dest="ensemble_set_index", @@ -325,6 +366,13 @@ def main() -> None: type=int, default=1022, help="ESM residue token budget excluding CLS and EOS.") + parser.add_argument("--seq-len-feature", + action="store", + dest="seq_len_feature", + type=str, + choices=list(PREDICTION_SEQ_LEN_FEATURES), + default="auto", + help="Sequence-length feature mode for generated embeddings; auto uses checkpoint metadata.") parser.add_argument("--log-dir", action="store", dest="log_dir", @@ -367,6 +415,37 @@ def main() -> None: type=int, default=None, help="Explicit output classes (binary=1).") + parser.add_argument("--model-head", + action="store", + dest="model_head", + type=str, + choices=list(MODEL_HEADS), + default=None, + help="Explicit model head used in training.") + parser.add_argument("--conv-channels", + action="store", + dest="conv_channels", + type=int, + default=None, + help="Explicit Conv1d channels used in training.") + parser.add_argument("--conv-layers", + action="store", + dest="conv_layers", + type=int, + default=None, + help="Explicit Conv1d layer count used in training.") + parser.add_argument("--conv-kernel-size", + action="store", + dest="conv_kernel_size", + type=int, + default=None, + help="Explicit Conv1d kernel size used in training.") + parser.add_argument("--conv-dropout", + action="store", + dest="conv_dropout", + type=float, + default=None, + help="Explicit Conv1d dropout used in training.") parser.add_argument("--use-layer-norm", action="store_true", dest="use_layer_norm", @@ -392,6 +471,10 @@ def main() -> None: raise ValueError("--k-folds must be >= 1") if args.threshold is not None and (args.threshold <= 0.0 or args.threshold >= 1.0): raise ValueError("--threshold must be between (0.0, 1.0)") + if args.ensemble_threshold is not None and ( + args.ensemble_threshold <= 0.0 or args.ensemble_threshold >= 1.0 + ): + raise ValueError("--ensemble-threshold must be between (0.0, 1.0)") json_indent = 2 if args.log_json else None logger = setup_logger(log_dir=args.log_dir, @@ -453,6 +536,28 @@ def main() -> None: raise ValueError( f"All ensemble members must share emb_dim for shared-embedding inference, got {sorted(emb_dims)}" ) + seq_len_features = { + resolve_prediction_seq_len_feature(args.seq_len_feature, cfg) + for cfg in member_model_cfgs + } + if len(seq_len_features) != 1: + raise ValueError( + "All ensemble members must share seq_len_feature for shared-embedding inference, " + f"got {sorted(seq_len_features)}" + ) + resolved_seq_len_feature = next(iter(seq_len_features)) + + resolved_ensemble_threshold = None + if len(psp_models) > 1 and args.ensemble_aggregation == "mean-prob": + resolved_ensemble_threshold = ( + float(args.ensemble_threshold) + if args.ensemble_threshold is not None + else ( + float(args.threshold) + if args.threshold is not None + else float(sum(member_thresholds) / len(member_thresholds)) + ) + ) first_cfg = member_model_cfgs[0] logger.info("prediction_init", @@ -469,12 +574,20 @@ def main() -> None: ), "threshold": float(member_thresholds[0]) if len(member_thresholds) == 1 else None, "member_thresholds": [float(x) for x in member_thresholds], + "ensemble_aggregation": str(args.ensemble_aggregation), + "ensemble_threshold": resolved_ensemble_threshold, "device": device, "model_cfg_src": str(member_model_cfg_srcs[0]) if len(set(member_model_cfg_srcs)) == 1 else "mixed", + "model_head": str(first_cfg.model_head), + "seq_len_feature": resolved_seq_len_feature, "emb_dim": int(first_cfg.emb_dim), "hidden_sizes": [int(x) for x in first_cfg.hidden_sizes], "use_layer_norm": bool(first_cfg.use_layer_norm), "use_residual": bool(first_cfg.use_residual), + "conv_channels": int(first_cfg.conv_channels), + "conv_layers": int(first_cfg.conv_layers), + "conv_kernel_size": int(first_cfg.conv_kernel_size), + "conv_dropout": float(first_cfg.conv_dropout), "num_classes": int(first_cfg.num_classes) }}) @@ -493,7 +606,8 @@ def main() -> None: layer=layer, batch_converter=batch_converter, device=device, - max_tokens=args.max_tokens + max_tokens=args.max_tokens, + seq_len_feature=resolved_seq_len_feature, ) if len(psp_models) == 1: pred = predict_from_embedding( @@ -507,7 +621,9 @@ def main() -> None: psp_models=psp_models, protein_emb=protein_emb, device=device, - thresholds=member_thresholds + thresholds=member_thresholds, + aggregation=args.ensemble_aggregation, + ensemble_threshold=resolved_ensemble_threshold ) out_f.write(f">{header}\n{pred['binary_mask']}\n") @@ -538,7 +654,10 @@ def main() -> None: "processed": processed, "failed": failed, "total_residues": total_residues, - "total_epitopes": total_epitopes + "total_epitopes": total_epitopes, + "ensemble_aggregation": str(args.ensemble_aggregation), + "seq_len_feature": resolved_seq_len_feature, + "ensemble_threshold": resolved_ensemble_threshold }}) diff --git a/src/pepseqpred/apps/train_ffnn_cli.py b/src/pepseqpred/apps/train_cli.py similarity index 78% rename from src/pepseqpred/apps/train_ffnn_cli.py rename to src/pepseqpred/apps/train_cli.py index a43ddfa..d8c944b 100644 --- a/src/pepseqpred/apps/train_ffnn_cli.py +++ b/src/pepseqpred/apps/train_cli.py @@ -1,6 +1,6 @@ -"""train_ffnn_cli.py +"""train_cli.py -Handles end-to-end training and evaluation of a PepSeqPredFFNN with the goal to predict the locations +Handles end-to-end training and evaluation of PepSeqPred model heads with the goal to predict the locations of antibody epitopes within a protein sequence downstream. The resulting model will make binary predictions: definite epitope or not epitope, it handles residues labeled uncertain through a masking process. @@ -11,8 +11,8 @@ Usage ----- ->>> # from scripts/hpc/trainffnn.sh (see shell script for CLI config) ->>> sbatch trainffnn.sh /path/to/emb_shard_dir0 ... /path/to/emb_shard_dirN -- \\ +>>> # from scripts/hpc/train.sh (see shell script for CLI config) +>>> sbatch train.sh /path/to/emb_shard_dir0 ... /path/to/emb_shard_dirN -- \\ /path/to/label_shard0.pt ... /path/to/label_shardN.pt """ @@ -32,7 +32,13 @@ from pepseqpred.core.io.logger import setup_logger from pepseqpred.core.io.write import append_csv_row from pepseqpred.core.data.proteindataset import ProteinDataset, pad_collate -from pepseqpred.core.models.ffnn import PepSeqFFNN +from pepseqpred.core.models.factory import ( + MODEL_HEADS, + PepSeqModelConfig, + build_pepseq_model, + model_config_to_dict, + validate_model_config +) from pepseqpred.core.train.trainer import ( Trainer, TrainerConfig, @@ -40,18 +46,31 @@ ) from pepseqpred.core.train.ddp import init_ddp from pepseqpred.core.train.split import ( + SPLIT_STRATEGIES, split_ids, split_ids_grouped, + split_ids_label_stratified, build_kfold_splits, build_grouped_kfold_splits, + build_label_stratified_kfold_splits, + build_label_support_by_id, + build_split_report, partition_ids_weighted, sort_ids_for_locality, shuffle_ids_by_group ) -from pepseqpred.core.train.weights import pos_weight_from_label_shards +from pepseqpred.core.train.weights import ( + compute_pos_neg_counts, + global_pos_neg_counts +) +from pepseqpred.core.train.threshold import THRESHOLD_POLICIES from pepseqpred.core.train.embedding import infer_emb_dim from pepseqpred.core.train.seed import set_all_seeds from pepseqpred.core.io.read import parse_int_csv, parse_float_csv +from pepseqpred.core.data.seq_len_feature import ( + EMBEDDING_SEQ_LEN_FEATURES, + cli_seq_len_feature_to_model +) def summarize_numeric(series: pd.Series) -> Dict[str, Any]: @@ -105,6 +124,21 @@ def _finite_or_none(value: Any) -> float | None: return num if math.isfinite(num) else None +def _split_summary_csv_fields( + prefix: str, + summary: Mapping[str, Any] | None +) -> Dict[str, Any]: + """Flattens split report summary fields for run CSV artifacts.""" + summary = summary or {} + return { + f"{prefix}Proteins": _finite_or_none(summary.get("protein_count")), + f"{prefix}ValidResidues": _finite_or_none(summary.get("valid_residues")), + f"{prefix}PosResidues": _finite_or_none(summary.get("positive_residues")), + f"{prefix}NegResidues": _finite_or_none(summary.get("negative_residues")), + f"{prefix}PositiveRate": _finite_or_none(summary.get("positive_rate")), + } + + def _parse_plot_formats(raw: str) -> Tuple[str, ...]: """Parses comma-separated plot file formats.""" tokens = [t.strip().lower() for t in str(raw).split(",")] @@ -670,12 +704,28 @@ def _check_family_leakage(train_ids: List[str], val_ids: List[str], family_group def _build_run_plans( args: argparse.Namespace, protein_ids: List[str], - family_groups: Dict[str, str] + split_groups: Dict[str, str] | None = None, + family_groups: Dict[str, str] | None = None, + label_support_by_id: Mapping[str, Mapping[str, Any]] | None = None ) -> Tuple[List[RunPlan], Dict[str, Any]]: """Builds run plans from unified seed lists and n-folds configuration.""" if len(protein_ids) == 0: raise ValueError("No proteins found to train on") + split_strategy = str(getattr(args, "split_strategy", "size-balanced")) + # Backward-compatibility: legacy callers pass family_groups as the third + # positional argument, which now maps to split_groups. + if family_groups is None and str(args.split_type) == "id-family" and split_groups is not None: + family_groups = dict(split_groups) + family_groups = family_groups or {} + if split_groups is None: + split_groups = ( + family_groups + if str(args.split_type) == "id-family" + else {protein_id: protein_id for protein_id in protein_ids} + ) + label_support_by_id = label_support_by_id or {} + n_folds = int(args.n_folds) if n_folds < 1: raise ValueError("--n-folds must be >= 1") @@ -702,7 +752,18 @@ def _build_run_plans( if n_folds == 1: run_plans: List[RunPlan] = [] for run_index, (split_seed, train_seed) in enumerate(zip(split_seeds, train_seeds), start=1): - if args.split_type == "id-family": + if split_strategy == "label-stratified": + train_ids_all, val_ids_all = split_ids_label_stratified( + protein_ids, + args.val_frac, + split_seed, + split_groups, + label_support_by_id, + ) + if args.split_type == "id-family": + _check_family_leakage( + train_ids_all, val_ids_all, family_groups) + elif args.split_type == "id-family": train_ids_all, val_ids_all = split_ids_grouped( protein_ids, args.val_frac, split_seed, family_groups ) @@ -733,6 +794,7 @@ def _build_run_plans( "train_seeds": [int(x) for x in train_seeds], "n_folds": 1, "n_sets": int(n_sets), + "split_strategy": split_strategy, "train_mode": train_mode } @@ -742,7 +804,15 @@ def _build_run_plans( zip(split_seeds, train_seeds), start=1 ): - if args.split_type == "id-family": + if split_strategy == "label-stratified": + fold_splits = build_label_stratified_kfold_splits( + protein_ids, + n_folds=n_folds, + seed=int(set_split_seed), + groups=split_groups, + support_by_id=label_support_by_id, + ) + elif args.split_type == "id-family": fold_splits = build_grouped_kfold_splits( protein_ids, n_folds=n_folds, seed=int(set_split_seed), groups=family_groups ) @@ -794,14 +864,15 @@ def _build_run_plans( "n_folds": int(n_folds), "n_sets": int(n_sets), "ensemble_seed_mode": "set-paired", + "split_strategy": split_strategy, "train_mode": train_mode } def main() -> None: - """Handles command-line argument parsing and high-level execution of the Train FFNN program.""" + """Handles command-line argument parsing and high-level execution of training.""" parser = argparse.ArgumentParser( - description="Train PepSeqPred FFNN on protein ESM-2 embeddings for binary residue-level epitope prediction.") + description="Train a PepSeqPred model head on protein ESM-2 embeddings for binary residue-level epitope prediction.") parser.add_argument("--embedding-dirs", nargs="+", required=True, @@ -832,6 +903,44 @@ def main() -> None: action="store_true", dest="use_residual", help="If set, residuals are used in feed-forward calculation") + parser.add_argument("--model-head", + action="store", + dest="model_head", + type=str, + choices=list(MODEL_HEADS), + default="ffnn", + help="Classifier head to train.") + parser.add_argument("--conv-channels", + action="store", + dest="conv_channels", + type=int, + default=64, + help="Number of local Conv1d channels when --model-head=conv1d.") + parser.add_argument("--conv-layers", + action="store", + dest="conv_layers", + type=int, + default=2, + help="Number of Conv1d layers when --model-head=conv1d.") + parser.add_argument("--conv-kernel-size", + action="store", + dest="conv_kernel_size", + type=int, + default=9, + help="Odd Conv1d kernel size when --model-head=conv1d.") + parser.add_argument("--conv-dropout", + action="store", + dest="conv_dropout", + type=float, + default=0.1, + help="Dropout applied after each Conv1d activation when --model-head=conv1d.") + parser.add_argument("--seq-len-feature", + action="store", + dest="seq_len_feature", + type=str, + choices=list(EMBEDDING_SEQ_LEN_FEATURES), + default="none", + help="Sequence-length feature mode used in embedding tensors.") parser.add_argument("--epochs", action="store", dest="epochs", @@ -867,7 +976,7 @@ def main() -> None: action="store", type=float, default=None, - help="Optionally include a pre-calculated postive class weight") + help="Optional manual positive class weight; omitted means compute from the current training split") parser.add_argument("--save-path", action="store", dest="save_path", @@ -891,6 +1000,15 @@ def main() -> None: default="id-family", choices=["id", "id-family"], help="Data partition type, use ID only or ID and taxonomic family.") + parser.add_argument("--split-strategy", + type=str, + default="size-balanced", + choices=list(SPLIT_STRATEGIES), + help="Split assignment strategy. Default preserves existing size-balanced behavior.") + parser.add_argument("--split-report-json", + type=Path, + default=None, + help="Optional split report JSON path. Defaults to /split_report.json.") parser.add_argument("--num-workers", action="store", dest="num_workers", @@ -902,13 +1020,13 @@ def main() -> None: dest="window_size", type=int, default=1000, - help="Window size for long protein sequences (<= 0 to disable)") + help="Training window size for long protein sequences (<= 0 to disable; validation uses full proteins)") parser.add_argument("--stride", action="store", dest="stride", type=int, default=900, - help="Stride between windows for long proteins") + help="Stride between training windows for long proteins") parser.add_argument("--no-collapse-labels", dest="collapse_labels", action="store_false", @@ -916,7 +1034,7 @@ def main() -> None: parser.add_argument("--no-pad-last-window", dest="pad_last_window", action="store_false", - help="Disable padding of final short window") + help="Disable padding of final short training window") parser.add_argument("--no-cache-label-shard", dest="cache_current_label_shard", action="store_false", @@ -950,6 +1068,23 @@ def main() -> None: choices=["loss", "precision", "recall", "f1", "mcc", "auc", "auc10", "pr_auc", "res_balanced_acc"], help="Metric used to choose the best model checkpoint per run") + parser.add_argument("--threshold-policy", + type=str, + default="max-recall-min-precision", + choices=list(THRESHOLD_POLICIES), + help="Validation threshold selection policy.") + parser.add_argument("--threshold-min-precision", + type=float, + default=0.25, + help="Minimum precision for max-recall-min-precision threshold selection.") + parser.add_argument("--threshold-min-recall", + type=float, + default=0.80, + help="Minimum recall for min-recall-max-precision threshold selection.") + parser.add_argument("--threshold-fixed-value", + type=float, + default=0.50, + help="Fixed threshold used when --threshold-policy=fixed.") parser.add_argument("--results-csv", type=Path, default=None, @@ -993,9 +1128,15 @@ def main() -> None: ) if len(unknown) > 0: parser.error(f"unrecognized arguments: {' '.join(unknown)}") + if args.threshold_min_precision < 0.0 or args.threshold_min_precision > 1.0: + raise ValueError("--threshold-min-precision must be in [0.0, 1.0]") + if args.threshold_min_recall < 0.0 or args.threshold_min_recall > 1.0: + raise ValueError("--threshold-min-recall must be in [0.0, 1.0]") + if args.threshold_fixed_value <= 0.0 or args.threshold_fixed_value >= 1.0: + raise ValueError("--threshold-fixed-value must be between (0.0, 1.0)") logger = setup_logger(json_lines=True, json_indent=2, - name="train_ffnn_cli") + name="train_cli") ddp = init_ddp() rank = ddp["rank"] if ddp is not None else 0 @@ -1049,18 +1190,37 @@ def main() -> None: if args.subset > 0: protein_ids = protein_ids[:args.subset] - # parition data by ID + family or just ID + split_report_json = args.split_report_json or ( + args.save_path / "split_report.json") + + # partition data by ID + family or just ID family_groups: Dict[str, str] = {} missing_family_ids = 0 - if args.split_type == "id-family": - for protein_id in protein_ids: - family = base_dataset.embedding_family_by_id.get(protein_id) - if family is None or str(family).strip() == "": - # singleton group when family missing fallback - family_groups[protein_id] = f"__missing_family__:{protein_id}" - missing_family_ids += 1 - else: - family_groups[protein_id] = str(family) + for protein_id in protein_ids: + family = base_dataset.embedding_family_by_id.get(protein_id) + if family is None or str(family).strip() == "": + # singleton group when family missing fallback + family_groups[protein_id] = f"__missing_family__:{protein_id}" + missing_family_ids += 1 + else: + family_groups[protein_id] = str(family) + split_groups = ( + family_groups + if args.split_type == "id-family" + else {protein_id: protein_id for protein_id in protein_ids} + ) + + if ddp is None or rank == 0: + label_support_by_id = build_label_support_by_id( + protein_ids, + base_dataset.label_index, + ) + else: + label_support_by_id = {} + if ddp is not None: + obj = [label_support_by_id] + dist.broadcast_object_list(obj, src=0) + label_support_by_id = obj[0] # estimate relative workload without tensor I/O by using embedding file size. id_weights: Dict[str, float] = {} @@ -1079,12 +1239,52 @@ def main() -> None: for protein_id in protein_ids } - run_plans, split_meta = _build_run_plans(args, protein_ids, family_groups) + run_plans, split_meta = _build_run_plans( + args, + protein_ids, + split_groups, + family_groups, + label_support_by_id, + ) + split_summary_by_run: Dict[int, Mapping[str, Any]] = {} if rank == 0: + split_report_payload = build_split_report( + run_splits=[ + { + "run_index": plan.run_index, + "train_mode": plan.train_mode, + "split_seed": plan.split_seed, + "train_seed": plan.train_seed, + "fold_index": plan.fold_index, + "n_folds": plan.n_folds, + "ensemble_set_index": plan.ensemble_set_index, + "train_ids": plan.train_ids_all, + "val_ids": plan.val_ids_all, + } + for plan in run_plans + ], + support_by_id=label_support_by_id, + families_by_id=family_groups, + split_type=str(args.split_type), + split_strategy=str(args.split_strategy), + ) + split_report_json.parent.mkdir(parents=True, exist_ok=True) + split_report_json.write_text( + json.dumps(_sanitize_for_json(split_report_payload), + indent=2, allow_nan=False), + encoding="utf-8", + ) + split_summary_by_run = { + int(entry["run_index"]): entry + for entry in split_report_payload["runs"] + if entry.get("run_index") is not None + } logger.info("run_plan_init", extra={"extra": { "train_mode": str(split_meta["train_mode"]), "n_runs": len(run_plans), "split_type": args.split_type, + "split_strategy": args.split_strategy, + "split_report_json": str(split_report_json), "missing_family_ids": missing_family_ids, "n_folds": int(split_meta["n_folds"]) if "n_folds" in split_meta else None, "n_sets": int(split_meta["n_sets"]) if "n_sets" in split_meta else None, @@ -1105,6 +1305,33 @@ def main() -> None: raise ValueError( "--hidden-sizes and --dropouts must be the same length" ) + emb_dim = infer_emb_dim(base_dataset.embedding_index) + model_config = validate_model_config( + PepSeqModelConfig( + emb_dim=emb_dim, + hidden_sizes=hidden_sizes, + dropouts=dropouts, + num_classes=1, + use_layer_norm=bool(args.use_layer_norm), + use_residual=bool(args.use_residual), + model_head=str(args.model_head), + conv_channels=int(args.conv_channels), + conv_layers=int(args.conv_layers), + conv_kernel_size=int(args.conv_kernel_size), + conv_dropout=float(args.conv_dropout), + seq_len_feature=cli_seq_len_feature_to_model( + args.seq_len_feature), + ) + ) + model_config_payload = model_config_to_dict(model_config) + if rank == 0: + logger.info( + "model_config_resolved", + extra={"extra": { + **model_config_payload, + "seq_len_feature": str(args.seq_len_feature), + }} + ) # per-run loop for run_plan in run_plans: @@ -1196,10 +1423,10 @@ def main() -> None: protein_ids=val_ids, label_index=base_dataset.label_index, embedding_index=base_dataset.embedding_index, - window_size=args.window_size if args.window_size > 0 else None, - stride=args.stride, + window_size=None, + stride=1, collapse_labels=args.collapse_labels, - pad_last_window=args.pad_last_window, + pad_last_window=False, return_meta=False, cache_current_label_shard=args.cache_current_label_shard, drop_label_after_use=args.drop_label_after_use, @@ -1225,20 +1452,36 @@ def main() -> None: val_data, **loader_kwargs) if val_data is not None else None # compute or store positive weight - pos_weight = None + pos_weight_source = "cli" + train_pos_residues = None + train_neg_residues = None if args.pos_weight is not None: pos_weight = float(args.pos_weight) else: - pos_weight = pos_weight_from_label_shards(label_shards) + local_pos, local_neg = compute_pos_neg_counts(train_loader) + train_pos_residues, train_neg_residues = global_pos_neg_counts( + local_pos, + local_neg, + ddp + ) + pos_weight = float(train_neg_residues / max(train_pos_residues, 1)) + pos_weight_source = "train_loader" + + if rank == 0: + logger.info( + "pos_weight_resolved", + extra={ + "extra": { + "run_index": int(run_index), + "source": pos_weight_source, + "train_pos_residues": train_pos_residues, + "train_neg_residues": train_neg_residues, + "pos_weight": float(pos_weight) + } + } + ) - # build our FFNN model - emb_dim = infer_emb_dim(base_dataset.embedding_index) - model = PepSeqFFNN(emb_dim=emb_dim, - hidden_sizes=hidden_sizes, - dropouts=dropouts, - use_layer_norm=args.use_layer_norm, - use_residual=args.use_residual, - num_classes=1) + model = build_pepseq_model(model_config) if ddp is not None: device = torch.device(f"cuda:{ddp['local_rank']}") @@ -1257,12 +1500,17 @@ def main() -> None: learning_rate=args.lr, weight_decay=args.weight_decay, device="cuda" if torch.cuda.is_available() else "cpu", - pos_weight=pos_weight) + pos_weight=pos_weight, + threshold_policy=args.threshold_policy, + threshold_min_precision=args.threshold_min_precision, + threshold_min_recall=args.threshold_min_recall, + threshold_fixed_value=args.threshold_fixed_value) trainer = Trainer(model=model, train_loader=train_loader, logger=logger, val_loader=val_loader, - config=config) + config=config, + model_config=model_config) # run training, only save if rank 0 or single rank run if ddp is None or rank == 0: @@ -1318,6 +1566,9 @@ def main() -> None: "best_val_loss": best_val_loss, "best_score_value": best_score_value }}) + split_entry = split_summary_by_run.get(int(run_index), {}) + train_split_summary = split_entry.get("train", {}) + val_split_summary = split_entry.get("validation", {}) row = { "RunIndex": run_index, "TrainMode": run_plan.train_mode, @@ -1333,6 +1584,16 @@ def main() -> None: "NFolds": run_plan.n_folds, "SplitSeed": split_seed, "TrainSeed": train_seed, + "SplitStrategy": str(args.split_strategy), + "SplitReportJson": str(split_report_json), + "ModelHead": str(model_config.model_head), + "SeqLenFeature": str(args.seq_len_feature), + "ConvChannels": int(model_config.conv_channels), + "ConvLayers": int(model_config.conv_layers), + "ConvKernelSize": int(model_config.conv_kernel_size), + "ConvDropout": float(model_config.conv_dropout), + **_split_summary_csv_fields("Train", train_split_summary), + **_split_summary_csv_fields("Val", val_split_summary), "RunSaveDir": str(run_save_dir) if run_save_dir is not None else None, "CheckpointPath": str(checkpoint_path) if checkpoint_path is not None else None, "BestMetricKey": args.best_model_metric, @@ -1340,6 +1601,16 @@ def main() -> None: "BestEpoch": best_epoch, "BestValLoss": best_val_loss, "Threshold": threshold, + "ThresholdPolicy": str(best_metrics.get("threshold_policy", args.threshold_policy)), + "ThresholdStatus": str(best_metrics.get("threshold_status", "")), + "ThresholdMinPrecision": _finite_or_none( + best_metrics.get("threshold_min_precision", float("nan"))), + "ThresholdMinRecall": _finite_or_none( + best_metrics.get("threshold_min_recall", float("nan"))), + "ThresholdFixedValue": _finite_or_none( + best_metrics.get("threshold_fixed_value", float("nan"))), + "ThresholdPredPosFrac": _finite_or_none( + best_metrics.get("threshold_pred_pos_frac", float("nan"))), "PR_AUC": _finite_or_none(best_metrics.get("pr_auc", float("nan"))), "F1": _finite_or_none(best_metrics.get("f1", float("nan"))), "MCC": _finite_or_none(best_metrics.get("mcc", float("nan"))), @@ -1362,11 +1633,23 @@ def main() -> None: "n_runs": int(len(run_rows)), "train_mode": str(split_meta["train_mode"]), "split_type": str(args.split_type), + "split_strategy": str(args.split_strategy), + "split_report_json": str(split_report_json), + "seq_len_feature": str(args.seq_len_feature), + "model_config": model_config_payload, "best_model_metric": str(args.best_model_metric), + "threshold_policy": str(args.threshold_policy), + "threshold_min_precision": float(args.threshold_min_precision), + "threshold_min_recall": float(args.threshold_min_recall), + "threshold_fixed_value": float(args.threshold_fixed_value), "split_seeds": [int(x) for x in split_meta["split_seeds"]], "train_seeds": [int(x) for x in split_meta["train_seeds"]], "metrics": { "BestMetricValue": summarize_numeric(df_runs["BestMetricValue"]), + "TrainPositiveRate": summarize_numeric(df_runs["TrainPositiveRate"]), + "ValPositiveRate": summarize_numeric(df_runs["ValPositiveRate"]), + "Threshold": summarize_numeric(df_runs["Threshold"]), + "ThresholdPredPosFrac": summarize_numeric(df_runs["ThresholdPredPosFrac"]), "PR_AUC": summarize_numeric(df_runs["PR_AUC"]), "F1": summarize_numeric(df_runs["F1"]), "MCC": summarize_numeric(df_runs["MCC"]), @@ -1421,6 +1704,11 @@ def main() -> None: "train_seed": int(row["TrainSeed"]), "checkpoint": row.get("CheckpointPath"), "threshold": row.get("Threshold"), + "threshold_policy": row.get("ThresholdPolicy"), + "threshold_status": row.get("ThresholdStatus"), + "threshold_min_precision": row.get("ThresholdMinPrecision"), + "threshold_min_recall": row.get("ThresholdMinRecall"), + "threshold_fixed_value": row.get("ThresholdFixedValue"), "status": row.get("Status"), "best_metric_value": row.get("BestMetricValue") }) @@ -1481,11 +1769,19 @@ def main() -> None: "ensemble_type": "kfold_majority_vote", "train_mode": str(split_meta["train_mode"]), "split_type": str(args.split_type), + "split_strategy": str(args.split_strategy), + "split_report_json": str(split_report_json), + "seq_len_feature": str(args.seq_len_feature), + "model_config": model_config_payload, "n_folds": int(split_meta["n_folds"]), "set_index": int(entry["set_index"]), "split_seed": int(entry["split_seed"]), "train_seed": int(entry["train_seed"]), "best_model_metric": str(args.best_model_metric), + "threshold_policy": str(args.threshold_policy), + "threshold_min_precision": float(args.threshold_min_precision), + "threshold_min_recall": float(args.threshold_min_recall), + "threshold_fixed_value": float(args.threshold_fixed_value), "n_members": int(len(members)), "n_valid_members": int(len(valid_members)), "voting": { @@ -1530,12 +1826,20 @@ def main() -> None: "ensemble_type": "kfold_majority_vote", "train_mode": str(split_meta["train_mode"]), "split_type": str(args.split_type), + "split_strategy": str(args.split_strategy), + "split_report_json": str(split_report_json), + "seq_len_feature": str(args.seq_len_feature), + "model_config": model_config_payload, "n_folds": int(split_meta["n_folds"]), "n_sets": int(n_sets), "ensemble_seed_mode": str( split_meta.get("ensemble_seed_mode", "set-paired") ), "best_model_metric": str(args.best_model_metric), + "threshold_policy": str(args.threshold_policy), + "threshold_min_precision": float(args.threshold_min_precision), + "threshold_min_recall": float(args.threshold_min_recall), + "threshold_fixed_value": float(args.threshold_fixed_value), "sets": set_payloads } root_manifest_path = args.ensemble_manifest or ( diff --git a/src/pepseqpred/apps/train_ffnn_optuna_cli.py b/src/pepseqpred/apps/train_optuna_cli.py similarity index 60% rename from src/pepseqpred/apps/train_ffnn_optuna_cli.py rename to src/pepseqpred/apps/train_optuna_cli.py index 273b3ad..2d2711c 100644 --- a/src/pepseqpred/apps/train_ffnn_optuna_cli.py +++ b/src/pepseqpred/apps/train_optuna_cli.py @@ -1,9 +1,9 @@ -"""train_ffnn_optuna_cli.py +"""train_optuna_cli.py -This CLI is very similar to `train_ffnn_cli.py`, except it optimizes for the best possible hyperparameters -within the user-defined ranges in the shell script `scripts/hpc/trainffnnoptuna.sh`. +This CLI is very similar to `train_cli.py`, except it optimizes for the best possible hyperparameters +within the user-defined ranges in the shell script `scripts/hpc/trainoptuna.sh`. -Handles end-to-end training, evaluation, and hyperparameter optimization of a PepSeqPredFFNN with the goal +Handles end-to-end training, evaluation, and hyperparameter optimization of a PepSeqPred model head with the goal to predict the locations of antibody epitopes within a protein sequence downstream. The resulting model will make binary predictions: definite epitope or not epitope, it handles residues labeled uncertain through a masking process. @@ -17,14 +17,15 @@ Usage ----- ->>> # from scripts/hpc/trainffnnoptuna.sh (see shell script for CLI config) ->>> sbatch trainffnnoptuna.sh /path/to/emb_shard_dir0 ... /path/to/emb_shard_dirN -- \\ +>>> # from scripts/hpc/trainoptuna.sh (see shell script for CLI config) +>>> sbatch trainoptuna.sh /path/to/emb_shard_dir0 ... /path/to/emb_shard_dirN -- \\ /path/to/label_shard0.pt ... /path/to/label_shardN.pt """ import argparse import json +import math import time import random from pathlib import Path @@ -38,18 +39,36 @@ from pepseqpred.core.io.logger import setup_logger from pepseqpred.core.io.write import append_csv_row from pepseqpred.core.data.proteindataset import ProteinDataset, pad_collate -from pepseqpred.core.models.ffnn import PepSeqFFNN +from pepseqpred.core.models.factory import ( + MODEL_HEADS, + PepSeqModelConfig, + build_pepseq_model, + model_config_to_dict, + validate_model_config +) from pepseqpred.core.train.trainer import Trainer, TrainerConfig from pepseqpred.core.train.ddp import init_ddp from pepseqpred.core.train.split import ( + SPLIT_STRATEGIES, split_ids, split_ids_grouped, + split_ids_label_stratified, + build_label_support_by_id, + build_split_report, partition_ids_weighted, sort_ids_for_locality ) -from pepseqpred.core.train.weights import pos_weight_from_label_shards +from pepseqpred.core.train.weights import ( + compute_pos_neg_counts, + global_pos_neg_counts +) +from pepseqpred.core.train.threshold import THRESHOLD_POLICIES from pepseqpred.core.train.embedding import infer_emb_dim from pepseqpred.core.train.seed import set_all_seeds +from pepseqpred.core.data.seq_len_feature import ( + EMBEDDING_SEQ_LEN_FEATURES, + cli_seq_len_feature_to_model +) def _broadcast_params(params: Dict[str, Any], ddp: Dict[str, Any] | None) -> Dict[str, Any]: @@ -61,6 +80,46 @@ def _broadcast_params(params: Dict[str, Any], ddp: Dict[str, Any] | None) -> Dic return obj[0] +def _finite_or_none(value: Any) -> float | None: + """Tries to convert number to float if finite, otherwise returns None.""" + try: + num = float(value) + except (TypeError, ValueError): + return None + return num if math.isfinite(num) else None + + +def _parse_int_choices(raw: str, flag: str) -> Tuple[int, ...]: + """Parse a comma-separated list of integer choices.""" + values: list[int] = [] + for token in str(raw).split(","): + token = token.strip() + if not token: + continue + try: + values.append(int(token)) + except ValueError as e: + raise ValueError(f"{flag} must contain only integers") from e + if len(values) < 1: + raise ValueError(f"{flag} must include at least one value") + return tuple(values) + + +def _split_summary_csv_fields( + prefix: str, + summary: Dict[str, Any] | None +) -> Dict[str, Any]: + """Flattens split report summary fields for trial CSV artifacts.""" + summary = summary or {} + return { + f"{prefix}Proteins": _finite_or_none(summary.get("protein_count")), + f"{prefix}ValidResidues": _finite_or_none(summary.get("valid_residues")), + f"{prefix}PosResidues": _finite_or_none(summary.get("positive_residues")), + f"{prefix}NegResidues": _finite_or_none(summary.get("negative_residues")), + f"{prefix}PositiveRate": _finite_or_none(summary.get("positive_rate")), + } + + def build_hidden_sizes(trial: optuna.trial.Trial, depth_min: int, depth_max: int, @@ -124,7 +183,7 @@ def build_hidden_sizes(trial: optuna.trial.Trial, def main() -> None: """Parses CLI arguments and runs Optuna study.""" parser = argparse.ArgumentParser( - description="Optuna tuning CLI for PepSeqPredFFNN.") + description="Optuna tuning CLI for PepSeqPred model heads.") parser.add_argument("--embedding-dirs", nargs="+", required=True, @@ -161,6 +220,37 @@ def main() -> None: choices=["precision", "recall", "f1", "mcc", "auc", "pr_auc"], help="Metric to maximize") + parser.add_argument("--model-head", + action="store", + dest="model_head", + type=str, + choices=list(MODEL_HEADS), + default="ffnn", + help="Classifier head to tune.") + parser.add_argument("--seq-len-feature", + action="store", + dest="seq_len_feature", + type=str, + choices=list(EMBEDDING_SEQ_LEN_FEATURES), + default="none", + help="Sequence-length feature mode used in embedding tensors.") + parser.add_argument("--threshold-policy", + type=str, + default="max-recall-min-precision", + choices=list(THRESHOLD_POLICIES), + help="Validation threshold selection policy.") + parser.add_argument("--threshold-min-precision", + type=float, + default=0.25, + help="Minimum precision for max-recall-min-precision threshold selection.") + parser.add_argument("--threshold-min-recall", + type=float, + default=0.80, + help="Minimum recall for min-recall-max-precision threshold selection.") + parser.add_argument("--threshold-fixed-value", + type=float, + default=0.50, + help="Fixed threshold used when --threshold-policy=fixed.") parser.add_argument("--val-frac", type=float, default=0.2, @@ -174,6 +264,15 @@ def main() -> None: default="id-family", choices=["id", "id-family"], help="Data partition type, use ID only or ID and taxonomic family.") + parser.add_argument("--split-strategy", + type=str, + default="size-balanced", + choices=list(SPLIT_STRATEGIES), + help="Split assignment strategy. Default preserves existing size-balanced behavior.") + parser.add_argument("--split-report-json", + type=Path, + default=None, + help="Optional split report JSON path. Defaults to /split_report.json.") parser.add_argument("--num-workers", type=int, default=4, @@ -183,7 +282,7 @@ def main() -> None: action="store", type=float, default=None, - help="Optionally include a pre-calculated postive class weight") + help="Optional manual positive class weight; omitted means compute from the current training split") parser.add_argument("--save-path", type=Path, default=Path("checkpoints/ffnn_optuna"), @@ -217,6 +316,30 @@ def main() -> None: type=str, default="32,64,128", help="Comma separated batch sizes") + parser.add_argument("--conv-channel-choices", + type=str, + default="32,64,128", + help="Comma-separated Conv1d channel choices used when --model-head=conv1d.") + parser.add_argument("--conv-layers-min", + type=int, + default=1, + help="Minimum Conv1d layers used when --model-head=conv1d.") + parser.add_argument("--conv-layers-max", + type=int, + default=3, + help="Maximum Conv1d layers used when --model-head=conv1d.") + parser.add_argument("--conv-kernel-size-choices", + type=str, + default="3,5,9,15", + help="Comma-separated odd Conv1d kernel choices used when --model-head=conv1d.") + parser.add_argument("--conv-dropout-min", + type=float, + default=0.0, + help="Minimum Conv1d dropout used when --model-head=conv1d.") + parser.add_argument("--conv-dropout-max", + type=float, + default=0.25, + help="Maximum Conv1d dropout used when --model-head=conv1d.") parser.add_argument("--lr-min", type=float, default=1e-4, @@ -233,9 +356,6 @@ def main() -> None: type=float, default=1e-2, help="Max weight decay") - parser.add_argument("--use-pos-weight", - action="store_true", - help="Use positive weight to handle positive vs. negative class imbalances.") parser.add_argument("--pruner-warmup", type=int, default=2, @@ -249,13 +369,13 @@ def main() -> None: dest="window_size", type=int, default=1000, - help="Window size for long protein sequences (<= 0 to disable)") + help="Training window size for long protein sequences (<= 0 to disable; validation uses full proteins)") parser.add_argument("--stride", action="store", dest="stride", type=int, default=900, - help="Stride between windows for long proteins") + help="Stride between training windows for long proteins") parser.add_argument("--no-collapse-labels", dest="collapse_labels", action="store_false", @@ -263,7 +383,7 @@ def main() -> None: parser.add_argument("--no-pad-last-window", dest="pad_last_window", action="store_false", - help="Disable padding of final short window") + help="Disable padding of final short training window") parser.add_argument("--no-cache-label-shard", dest="cache_current_label_shard", action="store_false", @@ -273,10 +393,24 @@ def main() -> None: action="store_false", help="Keep labels in memory after each protein is processed") args = parser.parse_args() + if args.threshold_min_precision < 0.0 or args.threshold_min_precision > 1.0: + raise ValueError("--threshold-min-precision must be in [0.0, 1.0]") + if args.threshold_min_recall < 0.0 or args.threshold_min_recall > 1.0: + raise ValueError("--threshold-min-recall must be in [0.0, 1.0]") + if args.threshold_fixed_value <= 0.0 or args.threshold_fixed_value >= 1.0: + raise ValueError("--threshold-fixed-value must be between (0.0, 1.0)") + if args.conv_layers_min < 1: + raise ValueError("--conv-layers-min must be >= 1") + if args.conv_layers_max < args.conv_layers_min: + raise ValueError("--conv-layers-max must be >= --conv-layers-min") + if args.conv_dropout_min < 0.0 or args.conv_dropout_min > 1.0: + raise ValueError("--conv-dropout-min must be in [0.0, 1.0]") + if args.conv_dropout_max < args.conv_dropout_min or args.conv_dropout_max > 1.0: + raise ValueError("--conv-dropout-max must be in [--conv-dropout-min, 1.0]") args.save_path.mkdir(parents=True, exist_ok=True) logger = setup_logger(json_lines=True, json_indent=2, - name="optuna_train_ffnn") + name="train_optuna_cli") ddp = init_ddp() rank = ddp["rank"] if ddp is not None else 0 @@ -318,24 +452,55 @@ def main() -> None: if args.subset > 0: protein_ids = protein_ids[:args.subset] - # parition data by ID + family or just ID - if args.split_type == "id-family": - family_groups: Dict[str, str] = {} - missing_family_ids = 0 - - for protein_id in protein_ids: - family = base_dataset.embedding_family_by_id.get(protein_id) - if family is None or str(family).strip() == "": - # singleton group when family missing fallback - family_groups[protein_id] = f"__missing_family__:{protein_id}" - missing_family_ids += 1 - else: - family_groups[protein_id] = str(family) + split_report_json = args.split_report_json or ( + args.save_path / "split_report.json") + # partition data by ID + family or just ID + family_groups: Dict[str, str] = {} + missing_family_ids = 0 + for protein_id in protein_ids: + family = base_dataset.embedding_family_by_id.get(protein_id) + if family is None or str(family).strip() == "": + family_groups[protein_id] = f"__missing_family__:{protein_id}" + missing_family_ids += 1 + else: + family_groups[protein_id] = str(family) + split_groups = ( + family_groups + if args.split_type == "id-family" + else {protein_id: protein_id for protein_id in protein_ids} + ) + + if ddp is None or rank == 0: + label_support_by_id = build_label_support_by_id( + protein_ids, + base_dataset.label_index, + ) + else: + label_support_by_id = {} + if ddp is not None: + obj = [label_support_by_id] + dist.broadcast_object_list(obj, src=0) + label_support_by_id = obj[0] + + if args.split_strategy == "label-stratified": + train_ids_all, val_ids_all = split_ids_label_stratified( + protein_ids, + args.val_frac, + seed, + split_groups, + label_support_by_id, + ) + elif args.split_type == "id-family": train_ids_all, val_ids_all = split_ids_grouped( protein_ids, args.val_frac, seed, family_groups ) + else: + train_ids_all, val_ids_all = split_ids( + protein_ids, args.val_frac, seed + ) + if args.split_type == "id-family": train_families = {family_groups[pid] for pid in train_ids_all} val_families = {family_groups[pid] for pid in val_ids_all} overlap = train_families & val_families @@ -343,24 +508,54 @@ def main() -> None: raise RuntimeError( f"Family leakage detected for split_type='id-family': n_overlap={len(overlap)}" ) - - if rank == 0: - logger.info("family_split_summary", - extra={"extra": { - "split_type": args.split_type, - "train_ids": len(train_ids_all), - "val_ids": len(val_ids_all), - "val_families": len(val_families), - "missing_family_ids": missing_family_ids - }}) else: - train_ids_all, val_ids_all = split_ids( - protein_ids, args.val_frac, seed - ) + val_families = set() if len(train_ids_all) == 0: raise ValueError("Global split produced 0 train IDs") + split_report_payload = None + train_split_summary: Dict[str, Any] = {} + val_split_summary: Dict[str, Any] = {} + if rank == 0: + split_report_payload = build_split_report( + run_splits=[ + { + "run_index": 1, + "train_mode": "optuna-holdout", + "split_seed": int(seed), + "train_seed": int(seed), + "fold_index": None, + "n_folds": 1, + "ensemble_set_index": None, + "train_ids": train_ids_all, + "val_ids": val_ids_all, + } + ], + support_by_id=label_support_by_id, + families_by_id=family_groups, + split_type=str(args.split_type), + split_strategy=str(args.split_strategy), + ) + split_report_json.parent.mkdir(parents=True, exist_ok=True) + split_report_json.write_text( + json.dumps(split_report_payload, indent=2, allow_nan=False), + encoding="utf-8", + ) + split_entry = split_report_payload["runs"][0] + train_split_summary = dict(split_entry["train"]) + val_split_summary = dict(split_entry["validation"]) + logger.info("family_split_summary", + extra={"extra": { + "split_type": args.split_type, + "split_strategy": args.split_strategy, + "split_report_json": str(split_report_json), + "train_ids": len(train_ids_all), + "val_ids": len(val_ids_all), + "val_families": len(val_families), + "missing_family_ids": missing_family_ids + }}) + # weight by embedding file size to reduce I/O id_weights: Dict[str, float] = {} for protein_id in protein_ids: @@ -458,10 +653,10 @@ def main() -> None: protein_ids=val_ids, label_index=base_dataset.label_index, embedding_index=base_dataset.embedding_index, - window_size=args.window_size if args.window_size > 0 else None, - stride=args.stride, + window_size=None, + stride=1, collapse_labels=args.collapse_labels, - pad_last_window=args.pad_last_window, + pad_last_window=False, return_meta=False, cache_current_label_shard=args.cache_current_label_shard, drop_label_after_use=args.drop_label_after_use @@ -473,6 +668,60 @@ def main() -> None: for x in args.batch_sizes.split(",") if x.strip()] if len(batch_sizes) == 0: raise ValueError("No batch sizes provided") + conv_channel_choices = _parse_int_choices( + args.conv_channel_choices, + "--conv-channel-choices", + ) + conv_kernel_size_choices = _parse_int_choices( + args.conv_kernel_size_choices, + "--conv-kernel-size-choices", + ) + bad_kernels = [x for x in conv_kernel_size_choices if x < 1 or x % 2 != 1] + if bad_kernels: + raise ValueError( + "--conv-kernel-size-choices must contain only positive odd integers" + ) + + # resolve positive class weight once from the current training split + pos_weight_source = "cli" + train_pos_residues = None + train_neg_residues = None + if args.pos_weight is not None: + resolved_pos_weight = float(args.pos_weight) + else: + count_loader_kwargs = { + "batch_size": batch_sizes[0], + "shuffle": False, + "num_workers": args.num_workers, + "pin_memory": pin, + "collate_fn": pad_collate + } + if args.num_workers > 0: + count_loader_kwargs["multiprocessing_context"] = "spawn" + count_loader_kwargs["prefetch_factor"] = 4 + + count_loader = DataLoader(train_data, **count_loader_kwargs) + local_pos, local_neg = compute_pos_neg_counts(count_loader) + train_pos_residues, train_neg_residues = global_pos_neg_counts( + local_pos, + local_neg, + ddp + ) + resolved_pos_weight = float(train_neg_residues / max(train_pos_residues, 1)) + pos_weight_source = "train_loader" + + if rank == 0: + logger.info( + "pos_weight_resolved", + extra={ + "extra": { + "source": pos_weight_source, + "train_pos_residues": train_pos_residues, + "train_neg_residues": train_neg_residues, + "pos_weight": float(resolved_pos_weight) + } + } + ) # setup Optuna study emb_dim = infer_emb_dim(base_dataset.embedding_index) @@ -516,13 +765,42 @@ def _sample_params(trial: optuna.trial.Trial) -> Dict[str, Any]: wd = trial.suggest_float( "weight_decay", args.wd_min, args.wd_max, log=True) batch_size = trial.suggest_categorical("batch_size", batch_sizes) + if args.model_head == "conv1d": + conv_channels = trial.suggest_categorical( + "conv_channels", + list(conv_channel_choices), + ) + conv_layers = trial.suggest_int( + "conv_layers", + args.conv_layers_min, + args.conv_layers_max, + ) + conv_kernel_size = trial.suggest_categorical( + "conv_kernel_size", + list(conv_kernel_size_choices), + ) + conv_dropout = trial.suggest_float( + "conv_dropout", + args.conv_dropout_min, + args.conv_dropout_max, + ) + else: + conv_channels = int(conv_channel_choices[0]) + conv_layers = int(args.conv_layers_min) + conv_kernel_size = int(conv_kernel_size_choices[0]) + conv_dropout = float(args.conv_dropout_min) return { + "model_head": str(args.model_head), "hidden_sizes": hidden_sizes, "dropouts": dropouts, "depth": depth, "use_layer_norm": use_layer_norm, "use_residual": use_residual, + "conv_channels": int(conv_channels), + "conv_layers": int(conv_layers), + "conv_kernel_size": int(conv_kernel_size), + "conv_dropout": float(conv_dropout), "learning_rate": lr, "weight_decay": wd, "batch_size": batch_size @@ -537,6 +815,10 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa depth = int(params["depth"]) use_layer_norm = bool(params["use_layer_norm"]) use_residual = bool(params["use_residual"]) + conv_channels = int(params["conv_channels"]) + conv_layers = int(params["conv_layers"]) + conv_kernel_size = int(params["conv_kernel_size"]) + conv_dropout = float(params["conv_dropout"]) lr = float(params["learning_rate"]) wd = float(params["weight_decay"]) batch_size = int(params["batch_size"]) @@ -557,12 +839,24 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa val_loader = DataLoader( val_data, **loader_kwargs) if val_data is not None else None - model = PepSeqFFNN(emb_dim=emb_dim, - hidden_sizes=hidden_sizes, - dropouts=dropouts, - use_layer_norm=use_layer_norm, - use_residual=use_residual, - num_classes=1) + model_config = validate_model_config( + PepSeqModelConfig( + emb_dim=emb_dim, + hidden_sizes=hidden_sizes, + dropouts=dropouts, + num_classes=1, + use_layer_norm=use_layer_norm, + use_residual=use_residual, + model_head=str(params["model_head"]), + conv_channels=conv_channels, + conv_layers=conv_layers, + conv_kernel_size=conv_kernel_size, + conv_dropout=conv_dropout, + seq_len_feature=cli_seq_len_feature_to_model( + args.seq_len_feature), + ) + ) + model = build_pepseq_model(model_config) device = torch.device(f"cuda:{ddp['local_rank']}") if ddp is not None else torch.device( "cuda" if torch.cuda.is_available() else "cpu") @@ -571,12 +865,8 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa model = DDP(model, device_ids=[ ddp["local_rank"]], output_device=ddp["local_rank"]) - # compute or store positive weight (like class weight) - pos_weight = None - if args.pos_weight is not None: - pos_weight = float(args.pos_weight) - else: - pos_weight = pos_weight_from_label_shards(label_shards) + # use the train-split positive weight resolved before trial search + pos_weight = resolved_pos_weight # setup config and trainer class config = TrainerConfig(epochs=args.epochs, @@ -584,7 +874,11 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa learning_rate=lr, weight_decay=wd, device="cuda" if torch.cuda.is_available() else "cpu", - pos_weight=pos_weight) + pos_weight=pos_weight, + threshold_policy=args.threshold_policy, + threshold_min_precision=args.threshold_min_precision, + threshold_min_recall=args.threshold_min_recall, + threshold_fixed_value=args.threshold_fixed_value) trial_dir = None if rank == 0 and trial is not None: @@ -596,7 +890,8 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa train_loader=train_loader, logger=logger, val_loader=val_loader, - config=config) + config=config, + model_config=model_config) # start and time trial start = time.time() @@ -612,7 +907,9 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa row: Dict[str, Any] = { "RunID": f"{args.study_name}_trial_{trial.number:04d}", "Timestamp": pd.Timestamp.utcnow().isoformat(), - "ModelVersion": "ffnn_optuna", + "ModelVersion": "train_optuna", + "ModelHead": str(model_config.model_head), + "SeqLenFeature": str(args.seq_len_feature), "NumParameters": int(sum(p.numel() for p in model.parameters())), "HiddenLayers": int(depth), "HiddenSizes": ",".join(str(x) for x in hidden_sizes), @@ -620,11 +917,19 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa "Dropout": float(dropouts[0]) if len(dropouts) else 0.0, "UseLayerNorm": bool(use_layer_norm), "UseResidual": bool(use_residual), + "ConvChannels": int(model_config.conv_channels), + "ConvLayers": int(model_config.conv_layers), + "ConvKernelSize": int(model_config.conv_kernel_size), + "ConvDropout": float(model_config.conv_dropout), "BatchSize": int(batch_size), "LearningRate": float(lr), "WeightDecay": float(wd), "PosWeight": float(pos_weight) if pos_weight else 1.0, "Epochs": int(args.epochs), + "SplitStrategy": str(args.split_strategy), + "SplitReportJson": str(split_report_json), + **_split_summary_csv_fields("Train", train_split_summary), + **_split_summary_csv_fields("Val", val_split_summary), "BestValLossAtScore": float(best_val_loss_at_score), "BestEpoch": int(best_epoch), "Precision": float(best_metrics.get("precision", float("nan"))), @@ -636,8 +941,12 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa "AUC10": float(best_metrics.get("auc10", float("nan"))), "PR_AUC": float(best_metrics.get("pr_auc", float("nan"))), "Threshold": float(best_metrics.get("threshold", float("nan"))), + "ThresholdPolicy": str(best_metrics.get("threshold_policy", args.threshold_policy)), "ThresholdStatus": str(best_metrics.get("threshold_status", "")), "ThresholdMinPrecision": float(best_metrics.get("threshold_min_precision", float("nan"))), + "ThresholdMinRecall": float(best_metrics.get("threshold_min_recall", float("nan"))), + "ThresholdFixedValue": float(best_metrics.get("threshold_fixed_value", float("nan"))), + "ThresholdPredPosFrac": float(best_metrics.get("threshold_pred_pos_frac", float("nan"))), "ScoreKey": args.metric, "ScoreValue": float(best_score), "ElapsedSec": float(elapsed), @@ -650,6 +959,8 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa trial.set_user_attr("best_val_loss_at_score", float(best_val_loss_at_score)) trial.set_user_attr("best_metrics", best_metrics) + trial.set_user_attr("seq_len_feature", str(args.seq_len_feature)) + trial.set_user_attr("model_config", model_config_to_dict(model_config)) return float(best_score) @@ -698,13 +1009,35 @@ def _run_trial(params: Dict[str, Any], trial: optuna.trial.Trial | None) -> floa "best_value": float(best.value), "best_params": dict(best.params), "best_user_attrs": dict(best.user_attrs), - "metric": args.metric} + "metric": args.metric, + "model_head": str(args.model_head), + "seq_len_feature": str(args.seq_len_feature), + "split_strategy": str(args.split_strategy), + "split_report_json": str(split_report_json), + "threshold_policy": str(args.threshold_policy), + "threshold_min_precision": float(args.threshold_min_precision), + "threshold_min_recall": float(args.threshold_min_recall), + "threshold_fixed_value": float(args.threshold_fixed_value)} best_trial_json.write_text(json.dumps(best_payload, indent=2)) best_trial_dir = Path(best.user_attrs["trial_dir"]) - src = best_trial_dir / "best_model_by_score.pt" - if src.exists(): + src = next( + ( + path for path in ( + best_trial_dir / "best_model_by_score.pt", + best_trial_dir / "fully_connected_by_score.pt", + ) + if path.exists() + ), + None, + ) + if src is not None: best_ckpt_path.write_bytes(src.read_bytes()) + else: + logger.warning( + "best_checkpoint_missing", + extra={"extra": {"trial_dir": str(best_trial_dir)}} + ) logger.info("best_results", extra={"extra": best_payload}) diff --git a/src/pepseqpred/core/data/seq_len_feature.py b/src/pepseqpred/core/data/seq_len_feature.py new file mode 100644 index 0000000..da0a77d --- /dev/null +++ b/src/pepseqpred/core/data/seq_len_feature.py @@ -0,0 +1,79 @@ +"""Sequence-length feature mode helpers.""" + +from typing import Any + +SEQ_LEN_FEATURE_NONE = "none" +SEQ_LEN_FEATURE_RAW = "raw" +SEQ_LEN_FEATURE_INVERSE = "inverse" +SEQ_LEN_FEATURE_AUTO = "auto" + +EMBEDDING_SEQ_LEN_FEATURES = ( + SEQ_LEN_FEATURE_NONE, + SEQ_LEN_FEATURE_RAW, + SEQ_LEN_FEATURE_INVERSE, +) +MODEL_SEQ_LEN_FEATURES = ( + SEQ_LEN_FEATURE_RAW, + SEQ_LEN_FEATURE_INVERSE, +) +PREDICTION_SEQ_LEN_FEATURES = ( + SEQ_LEN_FEATURE_AUTO, + SEQ_LEN_FEATURE_NONE, + SEQ_LEN_FEATURE_RAW, + SEQ_LEN_FEATURE_INVERSE, +) + + +def normalize_embedding_seq_len_feature(value: Any) -> str: + """Normalize an embedding-generation sequence-length feature mode.""" + if value is None: + return SEQ_LEN_FEATURE_NONE + mode = str(value).strip().lower() + if mode not in EMBEDDING_SEQ_LEN_FEATURES: + raise ValueError( + "seq_len_feature must be one of: " + + ", ".join(EMBEDDING_SEQ_LEN_FEATURES) + ) + return mode + + +def normalize_model_seq_len_feature(value: Any) -> str | None: + """Normalize model-config sequence-length feature metadata.""" + if value is None: + return None + mode = str(value).strip().lower() + if mode not in MODEL_SEQ_LEN_FEATURES: + raise ValueError( + "model_config.seq_len_feature must be either 'raw' or 'inverse' " + "when present" + ) + return mode + + +def cli_seq_len_feature_to_model(value: Any) -> str | None: + """Convert a CLI mode into model-config metadata.""" + mode = normalize_embedding_seq_len_feature(value) + if mode == SEQ_LEN_FEATURE_NONE: + return None + return mode + + +def model_seq_len_feature_to_embedding(value: Any) -> str: + """Convert model-config metadata into an embedding feature mode.""" + mode = normalize_model_seq_len_feature(value) + if mode is None: + return SEQ_LEN_FEATURE_NONE + return mode + + +def normalize_prediction_seq_len_feature(value: Any) -> str: + """Normalize a prediction-time sequence-length feature mode.""" + if value is None: + return SEQ_LEN_FEATURE_AUTO + mode = str(value).strip().lower() + if mode not in PREDICTION_SEQ_LEN_FEATURES: + raise ValueError( + "seq_len_feature must be one of: " + + ", ".join(PREDICTION_SEQ_LEN_FEATURES) + ) + return mode diff --git a/src/pepseqpred/core/embeddings/esm2.py b/src/pepseqpred/core/embeddings/esm2.py index 98f03a6..0066cfe 100644 --- a/src/pepseqpred/core/embeddings/esm2.py +++ b/src/pepseqpred/core/embeddings/esm2.py @@ -4,9 +4,9 @@ This module normalizes protein sequences, batches inputs by a token budget to reduce padding, and computes per-residue embeddings for both short and long -proteins (using sliding windows for long sequences). It also appends sequence -length as a feature and writes per-sequence `.pt` embeddings plus a CSV index -describing stored artifacts. +proteins (using sliding windows for long sequences). It can optionally append a +sequence-length feature and writes per-sequence `.pt` embeddings plus a CSV +index describing stored artifacts. """ import os @@ -19,6 +19,12 @@ import numpy as np import pandas as pd from pepseqpred.core.io.keys import build_emb_stem, normalize_family_value +from pepseqpred.core.data.seq_len_feature import ( + SEQ_LEN_FEATURE_INVERSE, + SEQ_LEN_FEATURE_NONE, + SEQ_LEN_FEATURE_RAW, + normalize_embedding_seq_len_feature +) def clean_seq(seq: str) -> str: @@ -117,8 +123,9 @@ def compute_window_embedding(token: torch.Tensor, Returns ------- torch.Tensor - An array of shape (L, D+1) where L is the number of residues (excluding CLS and EOS), and D+1 is - the embedding dimension for the chosen layer plus an additional column for the protein sequence length. + An array of shape (L, D) where L is the number of residues + (excluding CLS and EOS), and D is the embedding dimension for the + chosen layer. """ seq_len = token.size(1) - 2 # remove CLS and EOS tokens @@ -238,7 +245,50 @@ def append_seq_len(res_vec: np.ndarray, seq_len: int) -> np.ndarray: np.ndarray The updated embedding array with the protein sequence length appended. """ - col = np.full((res_vec.shape[0], 1), float(seq_len), dtype=res_vec.dtype) + return apply_seq_len_feature( + res_vec=res_vec, + seq_len=seq_len, + seq_len_feature=SEQ_LEN_FEATURE_RAW, + ) + + +def apply_seq_len_feature( + res_vec: np.ndarray, + seq_len: int, + seq_len_feature: str | None = SEQ_LEN_FEATURE_NONE, +) -> np.ndarray: + """ + Optionally append a sequence-length feature to residue embeddings. + + Parameters + ---------- + res_vec : np.ndarray + Per-residue ESM embeddings. + seq_len : int + Protein sequence length. + seq_len_feature : str or None + `"none"` leaves embeddings unchanged, `"raw"` appends `seq_len`, + and `"inverse"` appends `1.0 / seq_len`. + + Returns + ------- + np.ndarray + Embeddings with shape `(L, D)` for `"none"` or `(L, D+1)` for + appended feature modes. + """ + mode = normalize_embedding_seq_len_feature(seq_len_feature) + if mode == SEQ_LEN_FEATURE_NONE: + return res_vec + if int(seq_len) <= 0: + raise ValueError("seq_len must be > 0 when appending length feature") + + value = float(seq_len) + if mode == SEQ_LEN_FEATURE_INVERSE: + value = 1.0 / value + elif mode != SEQ_LEN_FEATURE_RAW: + raise ValueError(f"Unsupported seq_len_feature='{seq_len_feature}'") + + col = np.full((res_vec.shape[0], 1), value, dtype=res_vec.dtype) return np.concatenate([res_vec, col], axis=1) @@ -253,6 +303,7 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, index_csv_path: Path | str = "esm2_seq_index.csv", key_mode: str = "id-family", key_delimiter: str = "-", + seq_len_feature: str | None = SEQ_LEN_FEATURE_NONE, logger: Optional[logging.Logger] = None) -> Tuple[pd.DataFrame, List[str]]: """ Generate per residue ESM embeddings for sequences in a DataFrame and write outputs. @@ -284,6 +335,8 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, - "id-family": use `{id_col}{key_delimiter}{viral_family}`. key_delimiter : str Delimiter between ID and viral family when `key_mode` is "id-family". + seq_len_feature : str or None + Sequence-length feature mode: `"none"`, `"raw"`, or `"inverse"`. logger : logging.Logger or None Logger to use. If None, uses esm_cli logger. @@ -295,6 +348,7 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, """ # set up logger and start timer logger = logger or logging.getLogger("esm_cli") + seq_len_feature = normalize_embedding_seq_len_feature(seq_len_feature) index_records = [] t0 = time.perf_counter() @@ -383,7 +437,8 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, "extra": { "total_sequences": len(df), "key_mode": key_mode, - "key_delimiter": key_delimiter + "key_delimiter": key_delimiter, + "seq_len_feature": seq_len_feature }}) # load ESM embedding model @@ -456,8 +511,9 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, seq_key = str(batch[batch_idx][0]) seq_meta = metadata_by_key.get(seq_key, {}) - # append sequence length column - res_vec = append_seq_len(res_vec, len_) # (L, D+1) + # optionally append sequence length column + res_vec = apply_seq_len_feature( + res_vec, len_, seq_len_feature) # save as .pt per sequence torch.save(torch.from_numpy(res_vec), @@ -475,6 +531,7 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, "embed_dim": int(res_vec.shape[1]), # entire sequence length "original_seq_len": int(len_), + "seq_len_feature": seq_len_feature, "handle": "short", "model": model_name, "storage": ".pt"}) @@ -489,9 +546,9 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, # (L, D) single_token, model, layer, device).cpu().numpy().astype(np.float32) - # append sequence length column - res_vec = append_seq_len( - res_vec, len(protein_seq)) # (L, D+1) + # optionally append sequence length column + res_vec = apply_seq_len_feature( + res_vec, len(protein_seq), seq_len_feature) seq_meta = metadata_by_key.get(str(protein_key), {}) # save logs per-residue embeddings for sliding window long sequences @@ -503,6 +560,7 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, "length": int(res_vec.shape[0]), "embed_dim": int(res_vec.shape[1]), "original_seq_len": len(protein_seq), + "seq_len_feature": seq_len_feature, "handle": "long", "model": model_name, "storage": ".pt"}) @@ -534,6 +592,7 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, "length", "embed_dim", "original_seq_len", + "seq_len_feature", "handle", "model", "storage", @@ -565,6 +624,7 @@ def esm_embeddings_from_fasta(fasta_df: pd.DataFrame, "handled_short": int(by_handle.get("short", 0)), "handled_long": int(by_handle.get("long", 0)), "total_duration_s": round(elapsed, 3), + "seq_len_feature": seq_len_feature, "artifacts_path": str(os.path.abspath(per_seq_dir)), "index_csv_path": str(os.path.abspath(index_csv_path)) }}) diff --git a/src/pepseqpred/core/models/__init__.py b/src/pepseqpred/core/models/__init__.py index e69de29..da6aecd 100644 --- a/src/pepseqpred/core/models/__init__.py +++ b/src/pepseqpred/core/models/__init__.py @@ -0,0 +1,23 @@ +from pepseqpred.core.models.factory import ( + MODEL_HEADS, + PepSeqModelConfig, + build_pepseq_model, + model_config_from_mapping, + model_config_to_dict, + normalize_model_head, + validate_model_config, +) +from pepseqpred.core.models.ffnn import FFBlock, PepSeqConvFFNN, PepSeqFFNN + +__all__ = [ + "FFBlock", + "PepSeqFFNN", + "PepSeqConvFFNN", + "PepSeqModelConfig", + "MODEL_HEADS", + "build_pepseq_model", + "model_config_from_mapping", + "model_config_to_dict", + "normalize_model_head", + "validate_model_config", +] diff --git a/src/pepseqpred/core/models/factory.py b/src/pepseqpred/core/models/factory.py new file mode 100644 index 0000000..166191c --- /dev/null +++ b/src/pepseqpred/core/models/factory.py @@ -0,0 +1,164 @@ +"""Model configuration and construction helpers for PepSeqPred classifiers.""" + +from dataclasses import asdict, dataclass +from typing import Any, Mapping, Tuple +import torch.nn as nn +from pepseqpred.core.models.ffnn import PepSeqConvFFNN, PepSeqFFNN +from pepseqpred.core.data.seq_len_feature import normalize_model_seq_len_feature + +MODEL_HEAD_FFNN = "ffnn" +MODEL_HEAD_CONV1D = "conv1d" +MODEL_HEADS = (MODEL_HEAD_FFNN, MODEL_HEAD_CONV1D) + + +@dataclass(frozen=True) +class PepSeqModelConfig: + """Architecture configuration for a PepSeqPred residue classifier.""" + + emb_dim: int + hidden_sizes: Tuple[int, ...] + dropouts: Tuple[float, ...] + num_classes: int + use_layer_norm: bool + use_residual: bool + model_head: str = MODEL_HEAD_FFNN + conv_channels: int = 64 + conv_layers: int = 2 + conv_kernel_size: int = 9 + conv_dropout: float = 0.1 + seq_len_feature: str | None = None + + +def normalize_model_head(model_head: str | None) -> str: + """Normalize and validate a model-head token.""" + head = str(model_head or MODEL_HEAD_FFNN).strip().lower() + if head not in MODEL_HEADS: + raise ValueError( + f"model_head must be one of: {', '.join(MODEL_HEADS)}") + return head + + +def validate_model_config(config: PepSeqModelConfig) -> PepSeqModelConfig: + """Validate architecture config values and return a normalized config.""" + model_head = normalize_model_head(config.model_head) + hidden_sizes = tuple(int(x) for x in config.hidden_sizes) + dropouts = tuple(float(x) for x in config.dropouts) + if len(hidden_sizes) != len(dropouts): + raise ValueError("hidden_sizes and dropouts must have the same length") + if int(config.emb_dim) <= 0: + raise ValueError("emb_dim must be > 0") + if int(config.num_classes) != 1: + raise ValueError( + f"Inference expects binary residue model (num_classes=1), got {config.num_classes}" + ) + if int(config.conv_channels) <= 0: + raise ValueError("conv_channels must be > 0") + if int(config.conv_layers) < 1: + raise ValueError("conv_layers must be >= 1") + if int(config.conv_kernel_size) < 1 or int(config.conv_kernel_size) % 2 != 1: + raise ValueError("conv_kernel_size must be a positive odd integer") + if float(config.conv_dropout) < 0.0 or float(config.conv_dropout) > 1.0: + raise ValueError("conv_dropout must be in [0.0, 1.0]") + seq_len_feature = normalize_model_seq_len_feature( + config.seq_len_feature) + return PepSeqModelConfig( + emb_dim=int(config.emb_dim), + hidden_sizes=hidden_sizes, + dropouts=dropouts, + num_classes=int(config.num_classes), + use_layer_norm=bool(config.use_layer_norm), + use_residual=bool(config.use_residual), + model_head=model_head, + conv_channels=int(config.conv_channels), + conv_layers=int(config.conv_layers), + conv_kernel_size=int(config.conv_kernel_size), + conv_dropout=float(config.conv_dropout), + seq_len_feature=seq_len_feature, + ) + + +def model_config_from_mapping(raw: Mapping[str, Any]) -> PepSeqModelConfig: + """Build a validated model config from checkpoint or CLI metadata.""" + if not isinstance(raw, Mapping): + raise ValueError("model_config must be a mapping") + + missing = [ + key for key in ( + "emb_dim", + "hidden_sizes", + "dropouts", + "num_classes", + "use_layer_norm", + "use_residual", + ) + if key not in raw + ] + if missing: + raise ValueError( + "model_config missing required fields: " + ", ".join(missing) + ) + if "seq_len_feature" in raw and raw["seq_len_feature"] is None: + raise ValueError( + "model_config.seq_len_feature must be either 'raw' or 'inverse' " + "when present" + ) + + return validate_model_config( + PepSeqModelConfig( + emb_dim=int(raw["emb_dim"]), + hidden_sizes=tuple(int(x) for x in raw["hidden_sizes"]), + dropouts=tuple(float(x) for x in raw["dropouts"]), + num_classes=int(raw["num_classes"]), + use_layer_norm=bool(raw["use_layer_norm"]), + use_residual=bool(raw["use_residual"]), + model_head=normalize_model_head(raw.get("model_head", MODEL_HEAD_FFNN)), + conv_channels=int(raw.get("conv_channels", 64)), + conv_layers=int(raw.get("conv_layers", 2)), + conv_kernel_size=int(raw.get("conv_kernel_size", 9)), + conv_dropout=float(raw.get("conv_dropout", 0.1)), + seq_len_feature=( + normalize_model_seq_len_feature(raw["seq_len_feature"]) + if "seq_len_feature" in raw + else None + ), + ) + ) + + +def model_config_to_dict(config: PepSeqModelConfig) -> dict[str, Any]: + """Return JSON-serializable model config metadata.""" + normalized = validate_model_config(config) + out = asdict(normalized) + out["hidden_sizes"] = [int(x) for x in normalized.hidden_sizes] + out["dropouts"] = [float(x) for x in normalized.dropouts] + if normalized.seq_len_feature is None: + out.pop("seq_len_feature", None) + return out + + +def build_pepseq_model(config: PepSeqModelConfig) -> nn.Module: + """Construct a PepSeqPred model from a validated config.""" + config = validate_model_config(config) + if config.model_head == MODEL_HEAD_FFNN: + return PepSeqFFNN( + emb_dim=config.emb_dim, + hidden_sizes=config.hidden_sizes, + dropouts=config.dropouts, + num_classes=config.num_classes, + use_layer_norm=config.use_layer_norm, + use_residual=config.use_residual, + ) + if config.model_head == MODEL_HEAD_CONV1D: + return PepSeqConvFFNN( + emb_dim=config.emb_dim, + hidden_sizes=config.hidden_sizes, + dropouts=config.dropouts, + num_classes=config.num_classes, + use_layer_norm=config.use_layer_norm, + use_residual=config.use_residual, + conv_channels=config.conv_channels, + conv_layers=config.conv_layers, + conv_kernel_size=config.conv_kernel_size, + conv_dropout=config.conv_dropout, + ) + raise ValueError(f"Unsupported model_head={config.model_head}") diff --git a/src/pepseqpred/core/models/ffnn.py b/src/pepseqpred/core/models/ffnn.py index ad0f393..693d51b 100644 --- a/src/pepseqpred/core/models/ffnn.py +++ b/src/pepseqpred/core/models/ffnn.py @@ -6,7 +6,7 @@ residual connections, plus a reusable feed-forward block. """ -from typing import Iterable, Sequence +from typing import Sequence import torch import torch.nn as nn from .base import PepSeqClassifierBase @@ -111,7 +111,7 @@ def __init__(self, emb_dim: int = 1281, if len(hidden_sizes) != len(dropouts): raise ValueError("hidden_sizes and dropouts must be the same size") - layers: Iterable[nn.Module] = [] + layers: list[nn.Module] = [] in_features = emb_dim # arbitrary depth FFNN (default is 3 layers) @@ -140,7 +140,7 @@ def forward(self, X: torch.Tensor) -> torch.Tensor: Returns ------- Tensor - Logits of shape (B, C), where C is the number of output classes. + Logits of shape (B, L). """ if X.dim() != 3: raise ValueError(f"Expected shape (B, L, E), got {X.shape}") @@ -155,3 +155,114 @@ def forward(self, X: torch.Tensor) -> torch.Tensor: logits = self.ff_model(X) # (B * L, 1) logits = logits.view(B, L) # (B, L) return logits + + +class PepSeqConvFFNN(PepSeqClassifierBase): + """ + Residue-level classifier with a lightweight local Conv1d context head. + + The model preserves the same public forward contract as :class:`PepSeqFFNN`: + input tensors are shaped `(B, L, E)` and output logits are shaped `(B, L)`. + Local residue context is computed by applying Conv1d over the sequence + dimension, then concatenating those local features with the original ESM + embeddings before the dense FFNN classifier. + """ + + def __init__( + self, + emb_dim: int = 1281, + hidden_sizes: Sequence[int] = (150, 120, 45), + dropouts: Sequence[float] = (0.2, 0.2, 0.2), + num_classes: int = 1, + use_layer_norm: bool = False, + use_residual: bool = False, + conv_channels: int = 64, + conv_layers: int = 2, + conv_kernel_size: int = 9, + conv_dropout: float = 0.1, + ): + super().__init__(emb_dim=emb_dim, num_classes=num_classes) + + if len(hidden_sizes) != len(dropouts): + raise ValueError("hidden_sizes and dropouts must be the same size") + if int(conv_channels) <= 0: + raise ValueError("conv_channels must be > 0") + if int(conv_layers) < 1: + raise ValueError("conv_layers must be >= 1") + if int(conv_kernel_size) < 1 or int(conv_kernel_size) % 2 != 1: + raise ValueError("conv_kernel_size must be a positive odd integer") + if float(conv_dropout) < 0.0 or float(conv_dropout) > 1.0: + raise ValueError("conv_dropout must be in [0.0, 1.0]") + + self.conv_channels = int(conv_channels) + self.conv_layers = int(conv_layers) + self.conv_kernel_size = int(conv_kernel_size) + self.conv_dropout = float(conv_dropout) + + conv_modules: list[nn.Module] = [] + in_channels = emb_dim + padding = self.conv_kernel_size // 2 + for _ in range(self.conv_layers): + conv_modules.append( + nn.Conv1d( + in_channels=in_channels, + out_channels=self.conv_channels, + kernel_size=self.conv_kernel_size, + padding=padding, + ) + ) + conv_modules.append(nn.ReLU()) + conv_modules.append(nn.Dropout(self.conv_dropout)) + in_channels = self.conv_channels + self.conv_model = nn.Sequential(*conv_modules) + + layers: list[nn.Module] = [] + in_features = emb_dim + self.conv_channels + for hidden_size, p in zip(hidden_sizes, dropouts): + layers.append( + FFBlock( + in_dim=in_features, + out_dim=hidden_size, + dropout_p=p, + use_layer_norm=use_layer_norm, + use_residual=use_residual, + ) + ) + in_features = hidden_size + + layers.append(nn.Linear(in_features, num_classes)) + self.ff_model = nn.Sequential(*layers) + + def forward(self, X: torch.Tensor) -> torch.Tensor: + """ + Forward pass for residue-level epitope classification. + + Parameters + ---------- + X : Tensor + Input tensor of shape `(B, L, E)`, where B is the batch size, + L is the protein length, and E is the embedding dimension. + + Returns + ------- + Tensor + Logits of shape `(B, L)`. + """ + if X.dim() != 3: + raise ValueError(f"Expected shape (B, L, E), got {X.shape}") + + if X.size(-1) != self.emb_dim: + raise ValueError( + f"Expected emb_dim {self.emb_dim}, got {X.size(-1)}") + + B, L, D = X.shape + local = self.conv_model(X.transpose(1, 2)).transpose(1, 2) + if local.shape[:2] != X.shape[:2]: + raise ValueError( + f"Conv head changed sequence shape from {tuple(X.shape[:2])} " + f"to {tuple(local.shape[:2])}" + ) + features = torch.cat((X, local), dim=-1) + logits = self.ff_model(features.reshape(B * L, D + self.conv_channels)) + logits = logits.view(B, L) + return logits diff --git a/src/pepseqpred/core/predict/__init__.py b/src/pepseqpred/core/predict/__init__.py index ab34d4e..943e442 100644 --- a/src/pepseqpred/core/predict/__init__.py +++ b/src/pepseqpred/core/predict/__init__.py @@ -3,6 +3,7 @@ from pepseqpred.core.predict.artifacts import PredictionMember, resolve_prediction_members from pepseqpred.core.predict.inference import ( FFNNModelConfig, + PepSeqModelConfig, build_model_from_checkpoint, embed_protein_seq, infer_decision_threshold, @@ -18,6 +19,7 @@ "PredictionMember", "resolve_prediction_members", "FFNNModelConfig", + "PepSeqModelConfig", "build_model_from_checkpoint", "embed_protein_seq", "infer_decision_threshold", diff --git a/src/pepseqpred/core/predict/inference.py b/src/pepseqpred/core/predict/inference.py index 9c6c878..48c59b5 100644 --- a/src/pepseqpred/core/predict/inference.py +++ b/src/pepseqpred/core/predict/inference.py @@ -1,12 +1,27 @@ -from dataclasses import dataclass import math import re from typing import Tuple, Dict, Any, Mapping, List, Optional, Sequence import esm import torch import numpy as np -from pepseqpred.core.models.ffnn import PepSeqFFNN -from pepseqpred.core.embeddings.esm2 import clean_seq, compute_window_embedding, append_seq_len +from pepseqpred.core.models.factory import ( + MODEL_HEAD_FFNN, + PepSeqModelConfig, + build_pepseq_model, + model_config_from_mapping, + validate_model_config +) +from pepseqpred.core.embeddings.esm2 import ( + apply_seq_len_feature, + clean_seq, + compute_window_embedding +) +from pepseqpred.core.data.seq_len_feature import ( + SEQ_LEN_FEATURE_AUTO, + SEQ_LEN_FEATURE_NONE, + model_seq_len_feature_to_embedding, + normalize_prediction_seq_len_feature +) _HIDDEN_LAYER_RE = re.compile(r"^ff_model\.(\d+)\.linear\.weight$") _OUTPUT_LAYER_RE = re.compile(r"^ff_model\.(\d+)\.weight$") @@ -14,14 +29,19 @@ _SKIP_LINEAR_RE = re.compile(r"^ff_model\.\d+\.skip\.weight$") -@dataclass -class FFNNModelConfig: - emb_dim: int - hidden_sizes: Tuple[int, ...] - dropouts: Tuple[float, ...] - num_classes: int - use_layer_norm: bool - use_residual: bool +FFNNModelConfig = PepSeqModelConfig + + +def resolve_prediction_seq_len_feature( + seq_len_feature: str | None, + model_config: PepSeqModelConfig, +) -> str: + """Resolve prediction embedding length-feature mode from override/config.""" + mode = normalize_prediction_seq_len_feature(seq_len_feature) + if mode == SEQ_LEN_FEATURE_AUTO: + return model_seq_len_feature_to_embedding( + getattr(model_config, "seq_len_feature", None)) + return mode def normalize_state_dict_keys(state: Mapping[str, Any]) -> Dict[str, Any]: @@ -47,7 +67,7 @@ def normalize_state_dict_keys(state: Mapping[str, Any]) -> Dict[str, Any]: return out_dict -def infer_model_config_from_state(state: Mapping[str, Any]) -> FFNNModelConfig: +def infer_model_config_from_state(state: Mapping[str, Any]) -> PepSeqModelConfig: """ Infer PepSeqFFNN architecture directly from checkpoint `model_state_dict`. @@ -58,7 +78,7 @@ def infer_model_config_from_state(state: Mapping[str, Any]) -> FFNNModelConfig: Returns ------- - FFNNModelConfig + PepSeqModelConfig A fully populated model configuration dataclass to be used in model building. Raises @@ -137,21 +157,22 @@ def infer_model_config_from_state(state: Mapping[str, Any]) -> FFNNModelConfig: # dropout is not stored in state_dict dropouts = tuple(0.0 for _ in hidden_sizes) - return FFNNModelConfig( + return PepSeqModelConfig( emb_dim=emb_dim, hidden_sizes=hidden_sizes, dropouts=dropouts, num_classes=num_classes, use_layer_norm=use_layer_norm, - use_residual=use_residual + use_residual=use_residual, + model_head=MODEL_HEAD_FFNN, ) def build_model_from_checkpoint( checkpoint: Mapping[str, Any], device: str = "cpu", - model_config: Optional[FFNNModelConfig] = None -) -> Tuple[PepSeqFFNN, FFNNModelConfig, str]: + model_config: Optional[PepSeqModelConfig] = None +) -> Tuple[torch.nn.Module, PepSeqModelConfig, str]: """ Builds and loads model exactly from training checkpoint. @@ -161,22 +182,22 @@ def build_model_from_checkpoint( State dictionary obtained from saved trained model checkpoint. device : str Device type to use for building model checkpoint, default is `"cpu"`, `"cuda"` if accepted if GPUs are available. - model_config : FFNNModelConfig or None + model_config : PepSeqModelConfig or None A fully populated model configuration dataclass to be used in model building. Returns ------- - PepSeqFFNN + torch.nn.Module Fully populated PepSeqPred model ready to use in inference. - FFNNModelConfig + PepSeqModelConfig Model configuration dataclass returned for logging/tracking purposes. str - The model configuration source, either `"cli"` if passed explicitly or `"state_dict"` if inferred from the state dictionary. + The model configuration source: `"cli"`, `"checkpoint"`, or `"state_dict"`. Raises ------ ValueError - If `checkpoint` is not of type `Mapping` or `"model_state_dict"` is not a key in `checkpoint`. If `state` is not of type `Mapping`. If the number of output classes is not 1 (binary), if the embedding dimension is <= 0, or if the number hidden layer sizes is not the same as the number of dropouts. If the checkpoint weights are incompatible with the model architecture provided. + If `checkpoint` is not of type `Mapping` or `"model_state_dict"` is not a key in `checkpoint`. If `state` is not of type `Mapping`. If the model config is invalid or if the checkpoint weights are incompatible with the model architecture provided. """ if not isinstance(checkpoint, Mapping) or "model_state_dict" not in checkpoint: raise ValueError( @@ -189,35 +210,22 @@ def build_model_from_checkpoint( # can either pass config directly or infer from the state dict if model_config is not None: - config = model_config + config = validate_model_config(model_config) config_src = "cli" + elif isinstance(checkpoint.get("model_config"), Mapping): + config = model_config_from_mapping(checkpoint["model_config"]) + config_src = "checkpoint" else: config = infer_model_config_from_state(state) config_src = "state_dict" - # validate config - if config.num_classes != 1: - raise ValueError( - f"Inference expects binary residue model (num_classes=1), got {config.num_classes}") - if config.emb_dim <= 0: - raise ValueError("emb_dim must be > 0") - if len(config.hidden_sizes) != len(config.dropouts): - raise ValueError("hidden_sizes and dropouts must have the same length") - - model = PepSeqFFNN( - emb_dim=config.emb_dim, - hidden_sizes=config.hidden_sizes, - dropouts=config.dropouts, - num_classes=config.num_classes, - use_layer_norm=config.use_layer_norm, - use_residual=config.use_residual - ) + model = build_pepseq_model(config) try: model.load_state_dict(state, strict=True) except RuntimeError as e: raise ValueError( "Checkpoint weights are incompatible with the provided model architecture. " - "Check --emb-dim/--hidden-sizes/--dropouts/--use-layer-norm/--use-residual" + "Check --model-head and architecture flags." ) from e model.eval().to(device) return model, config, config_src @@ -292,7 +300,8 @@ def embed_protein_seq(protein_seq: str, layer: int, batch_converter: esm.data.BatchConverter, device: str, - max_tokens: int = 1022) -> torch.Tensor: + max_tokens: int = 1022, + seq_len_feature: str | None = SEQ_LEN_FEATURE_NONE) -> torch.Tensor: """ Generates the embedding for an entire protein sequence. @@ -311,6 +320,8 @@ def embed_protein_seq(protein_seq: str, The device to run the embedding model on (`"cpu"` or `"cuda"`). max_tokens : int Maximum number of tokens the ESM model can fit in its context window. Default is 1022. + seq_len_feature : str or None + Sequence-length feature mode: `"none"`, `"raw"`, or `"inverse"`. Returns ------- @@ -346,13 +357,17 @@ def embed_protein_seq(protein_seq: str, batch_tokens, esm_model, layer, device) rep_np = rep.numpy().astype(np.float32) - rep_np = append_seq_len(rep_np, seq_len) + rep_np = apply_seq_len_feature( + rep_np, + seq_len=seq_len, + seq_len_feature=seq_len_feature, + ) return torch.from_numpy(rep_np) def predict_member_probabilities_from_embedding( - psp_model: PepSeqFFNN, + psp_model: torch.nn.Module, protein_emb: torch.Tensor, device: str ) -> torch.Tensor: @@ -361,7 +376,7 @@ def predict_member_probabilities_from_embedding( Parameters ---------- - psp_model : PepSeqFFNN + psp_model : torch.nn.Module The PepSeqPred model to use for predicting member probabilities. protein_emb : torch.Tensor Protein embedding to pass through model. @@ -394,7 +409,7 @@ def predict_member_probabilities_from_embedding( return torch.sigmoid(logits)[0].detach().cpu() -def predict_from_embedding(psp_model: PepSeqFFNN, +def predict_from_embedding(psp_model: torch.nn.Module, protein_emb: torch.Tensor, device: str, threshold: float = 0.5) -> Dict[str, Any]: @@ -403,7 +418,7 @@ def predict_from_embedding(psp_model: PepSeqFFNN, Parameters ---------- - psp_model : PepSeqFFNN + psp_model : torch.nn.Module The PepSeqPred model to use for predicting member probabilities. protein_emb : torch.Tensor Protein embedding to pass through model. @@ -436,17 +451,19 @@ def predict_from_embedding(psp_model: PepSeqFFNN, def predict_ensemble_from_embedding( - psp_models: Sequence[PepSeqFFNN], + psp_models: Sequence[torch.nn.Module], protein_emb: torch.Tensor, device: str, - thresholds: Sequence[float] + thresholds: Sequence[float], + aggregation: str = "majority", + ensemble_threshold: float | None = None ) -> Dict[str, Any]: """ - Predict residue-level mask using strict majority vote across model members. + Predict residue-level mask from an ensemble of model members. Parameters ---------- - psp_models : Sequence[PepSeqFFNN] + psp_models : Sequence[torch.nn.Module] A sequence of PepSeqPred models to use for majority vote prediction of member probabilities. protein_emb : torch.Tensor Protein embedding to pass through models. @@ -454,6 +471,13 @@ def predict_ensemble_from_embedding( Device type: `"cuda"` for GPUs, otherwise `"cpu"`. thresholds : Sequence[float] A sequence of thresholds between `0.0` and `1.0`, that determine the cutoff for non-epitope vs definite epitope. Default is `0.5`. + aggregation : str + Ensemble aggregation rule. `"majority"` thresholds each member and + applies strict majority vote. `"mean-prob"` thresholds the mean + probability across members. + ensemble_threshold : float or None + Threshold used for `"mean-prob"` aggregation. If omitted, the mean + of member thresholds is used. Returns ------- @@ -464,9 +488,13 @@ def predict_ensemble_from_embedding( ------ ValueError If no models are provided, if number of models and thresholds differs, - if any threshold is outside `(0.0, 1.0)`, if member probability shapes - are invalid, or if member output lengths are inconsistent. + if any threshold is outside `(0.0, 1.0)`, if aggregation is invalid, + if member probability shapes are invalid, or if member output lengths + are inconsistent. """ + aggregation = str(aggregation).strip().lower() + if aggregation not in {"majority", "mean-prob"}: + raise ValueError("aggregation must be one of: majority, mean-prob") if len(psp_models) < 1: raise ValueError( "At least one model is required for ensemble prediction") @@ -497,36 +525,58 @@ def predict_ensemble_from_embedding( raise ValueError( "All ensemble members must produce the same sequence length") - vote_sum = torch.stack(member_masks, dim=0).sum(dim=0) n_members = int(len(member_masks)) - votes_needed = int((n_members // 2) + 1) - majority_mask = (vote_sum >= votes_needed).to(torch.int64) mean_probs = torch.stack(member_probs, dim=0).mean(dim=0) + votes_needed: int | None = None + if aggregation == "majority": + vote_sum = torch.stack(member_masks, dim=0).sum(dim=0) + votes_needed = int((n_members // 2) + 1) + mask = (vote_sum >= votes_needed).to(torch.int64) + payload_threshold = float("nan") + resolved_ensemble_threshold = None + else: + resolved_ensemble_threshold = ( + float(ensemble_threshold) + if ensemble_threshold is not None + else float(sum(float(x) for x in thresholds) / len(thresholds)) + ) + if resolved_ensemble_threshold <= 0.0 or resolved_ensemble_threshold >= 1.0: + raise ValueError("ensemble_threshold must be in (0.0, 1.0)") + mask = (mean_probs >= resolved_ensemble_threshold).to(torch.int64) + payload_threshold = float(resolved_ensemble_threshold) + out = _build_prediction_payload( probs=mean_probs, - mask=majority_mask, - threshold=float("nan") + mask=mask, + threshold=payload_threshold ) out["n_members"] = n_members out["votes_needed"] = votes_needed out["member_thresholds"] = [float(x) for x in thresholds] + out["ensemble_aggregation"] = aggregation + out["ensemble_threshold"] = ( + float(resolved_ensemble_threshold) + if resolved_ensemble_threshold is not None + else None + ) return out -def predict_protein(psp_model: PepSeqFFNN, +def predict_protein(psp_model: torch.nn.Module, esm_model: torch.nn.Module, layer: int, batch_converter: esm.data.BatchConverter, protein_seq: str, max_tokens: int, device: str, - threshold: float = 0.5) -> Dict[str, Any]: + threshold: float = 0.5, + seq_len_feature: str | None = SEQ_LEN_FEATURE_NONE) -> Dict[str, Any]: """ Predict residue-level binary epitope mask for one protein sequence Parameters ---------- - psp_model : PepSeqFFNN + psp_model : torch.nn.Module The PepSeqPred model to use for predicting member probabilities. esm_model : torch.nn.Module The ESM-2 model to use for embedding generation. @@ -542,6 +592,8 @@ def predict_protein(psp_model: PepSeqFFNN, Device type: `"cuda"` for GPUs, otherwise `"cpu"`. threshold : float The threshold between `0.0` and `1.0`, that determine the cutoff for non-epitope vs definite epitope. Default is `0.5`. + seq_len_feature : str or None + Sequence-length feature mode used for generated embeddings. Returns ------- @@ -562,7 +614,8 @@ def predict_protein(psp_model: PepSeqFFNN, layer=layer, batch_converter=batch_converter, device=device, - max_tokens=max_tokens + max_tokens=max_tokens, + seq_len_feature=seq_len_feature, ) return predict_from_embedding( psp_model=psp_model, diff --git a/src/pepseqpred/core/train/split.py b/src/pepseqpred/core/train/split.py index fd24e35..c6067be 100644 --- a/src/pepseqpred/core/train/split.py +++ b/src/pepseqpred/core/train/split.py @@ -8,7 +8,467 @@ import math import random -from typing import List, Dict, Tuple, Any, Optional, Set +from pathlib import Path +from typing import List, Dict, Tuple, Any, Optional, Set, Mapping, Sequence + +import torch + + +SPLIT_STRATEGIES: Tuple[str, ...] = ("size-balanced", "label-stratified") + + +def _positive_rate(pos_residues: int, valid_residues: int) -> float | None: + if valid_residues <= 0: + return None + return float(pos_residues / valid_residues) + + +def _zero_stats() -> Dict[str, int]: + return { + "protein_count": 0, + "valid_residues": 0, + "positive_residues": 0, + "negative_residues": 0, + } + + +def _add_stats(left: Mapping[str, int], right: Mapping[str, int]) -> Dict[str, int]: + return { + "protein_count": int(left.get("protein_count", 0)) + int(right.get("protein_count", 0)), + "valid_residues": int(left.get("valid_residues", 0)) + int(right.get("valid_residues", 0)), + "positive_residues": int(left.get("positive_residues", 0)) + int(right.get("positive_residues", 0)), + "negative_residues": int(left.get("negative_residues", 0)) + int(right.get("negative_residues", 0)), + } + + +def count_label_tensor_support(labels: torch.Tensor) -> Dict[str, int]: + """ + Count valid, positive, and negative residues using ProteinDataset masking semantics. + + Binary labels mark all residues valid. Three-column labels use columns + [definite epitope, uncertain, not epitope] and exclude uncertain residues. + """ + if not torch.is_tensor(labels): + raise TypeError(f"labels must be a torch.Tensor, not {type(labels)}") + y = labels.detach().cpu() + if y.dim() == 2 and y.size(1) == 3: + def_col = y[:, 0].float() + unc_col = y[:, 1].float() + valid_mask = unc_col == 0 + valid = int(valid_mask.sum().item()) + pos = int(((def_col == 1) & valid_mask).sum().item()) + neg = int(valid - pos) + else: + y_flat = y.view(-1) + valid = int(y_flat.numel()) + pos = int((y_flat == 1).sum().item()) + neg = int(valid - pos) + return { + "valid_residues": valid, + "positive_residues": pos, + "negative_residues": neg, + } + + +def build_label_support_by_id( + ids: Sequence[str], + label_index: Mapping[str, Path | str], +) -> Dict[str, Dict[str, Any]]: + """ + Build residue label-support summaries for selected protein IDs. + + Label shards are loaded once per shard path. Missing IDs are retained with + zero support and a non-OK status so split reports can account for them. + """ + ids_ordered = [str(protein_id) for protein_id in ids] + ids_by_shard: Dict[Path, List[str]] = {} + support_by_id: Dict[str, Dict[str, Any]] = {} + + for protein_id in ids_ordered: + shard_raw = label_index.get(protein_id) + if shard_raw is None: + support_by_id[protein_id] = { + "protein_id": protein_id, + "valid_residues": 0, + "positive_residues": 0, + "negative_residues": 0, + "label_shard": None, + "status": "missing_label", + } + continue + ids_by_shard.setdefault(Path(shard_raw), []).append(protein_id) + + for shard_path in sorted(ids_by_shard.keys(), key=lambda p: str(p)): + payload = torch.load(shard_path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict) or "labels" not in payload: + raise TypeError( + f"Label shard {shard_path} must be a dict with 'labels' key" + ) + labels_obj = payload["labels"] + if not isinstance(labels_obj, dict): + raise TypeError(f"'labels' in {shard_path} must be a dict") + + for protein_id in ids_by_shard[shard_path]: + labels = labels_obj.get(protein_id) + if labels is None: + support_by_id[protein_id] = { + "protein_id": protein_id, + "valid_residues": 0, + "positive_residues": 0, + "negative_residues": 0, + "label_shard": str(shard_path), + "status": "missing_label_tensor", + } + continue + counts = count_label_tensor_support(labels) + support_by_id[protein_id] = { + "protein_id": protein_id, + "valid_residues": int(counts["valid_residues"]), + "positive_residues": int(counts["positive_residues"]), + "negative_residues": int(counts["negative_residues"]), + "label_shard": str(shard_path), + "status": "ok", + } + del payload + + return support_by_id + + +def _support_stats_for_ids( + ids: Sequence[str], + support_by_id: Mapping[str, Mapping[str, Any]], +) -> Dict[str, int]: + stats = _zero_stats() + stats["protein_count"] = int(len(ids)) + for protein_id in ids: + support = support_by_id.get(str(protein_id), {}) + stats["valid_residues"] += int(support.get("valid_residues", 0)) + stats["positive_residues"] += int(support.get("positive_residues", 0)) + stats["negative_residues"] += int(support.get("negative_residues", 0)) + return stats + + +def _build_group_items( + ids: Sequence[str], + groups: Mapping[str, str], + support_by_id: Mapping[str, Mapping[str, Any]], +) -> List[Dict[str, Any]]: + group_to_ids: Dict[str, List[str]] = {} + for protein_id in ids: + group_id = str(groups.get(str(protein_id), str(protein_id))) + group_to_ids.setdefault(group_id, []).append(str(protein_id)) + + items: List[Dict[str, Any]] = [] + for group_id, members in group_to_ids.items(): + stats = _support_stats_for_ids(members, support_by_id) + items.append({ + "group_id": group_id, + "ids": list(members), + "stats": stats, + }) + return items + + +def _balance_score( + stats: Mapping[str, int], + target: Mapping[str, float], + overall_rate: float | None, + total: Mapping[str, int], +) -> float: + protein_total = max(float(total.get("protein_count", 0)), 1.0) + valid_total = max(float(total.get("valid_residues", 0)), 1.0) + pos_total = max(float(total.get("positive_residues", 0)), 1.0) + neg_total = max(float(total.get("negative_residues", 0)), 1.0) + + protein_err = abs(float(stats.get("protein_count", 0)) - float(target["protein_count"])) / protein_total + valid_err = abs(float(stats.get("valid_residues", 0)) - float(target["valid_residues"])) / valid_total + pos_err = abs(float(stats.get("positive_residues", 0)) - float(target["positive_residues"])) / pos_total + neg_err = abs(float(stats.get("negative_residues", 0)) - float(target["negative_residues"])) / neg_total + + rate_err = 0.0 + valid = int(stats.get("valid_residues", 0)) + if overall_rate is not None and valid > 0: + rate = float(stats.get("positive_residues", 0) / valid) + rate_err = abs(rate - overall_rate) + + empty_valid_penalty = ( + 0.25 if float(target["valid_residues"]) > 0.0 and valid == 0 else 0.0 + ) + return ( + protein_err + + (0.5 * valid_err) + + (2.0 * pos_err) + + (2.0 * neg_err) + + (3.0 * rate_err) + + empty_valid_penalty + ) + + +def _group_sort_key( + item: Mapping[str, Any], + overall_rate: float | None, +) -> Tuple[float, int, int, str]: + stats = item["stats"] + valid = int(stats["valid_residues"]) + rate = ( + float(stats["positive_residues"] / valid) + if valid > 0 + else (overall_rate if overall_rate is not None else 0.0) + ) + rate_delta = abs(rate - overall_rate) if overall_rate is not None else 0.0 + return ( + -rate_delta, + -int(stats["valid_residues"]), + -int(stats["protein_count"]), + str(item["group_id"]), + ) + + +def split_ids_label_stratified( + ids: List[str], + val_frac: float, + seed: int, + groups: Dict[str, str], + support_by_id: Mapping[str, Mapping[str, Any]], +) -> Tuple[List[str], List[str]]: + """ + Split IDs while keeping groups intact and balancing residue label support. + + This is an opt-in alternative to `split_ids_grouped`; it uses per-protein + valid/positive/negative residue counts instead of group size alone. + """ + ids = [str(protein_id) for protein_id in ids] + if val_frac < 0.0 or val_frac > 1.0: + return ids, [] + if val_frac == 0.0: + return ids, [] + if val_frac == 1.0: + return [], ids + if len(ids) == 0: + return [], [] + + group_items = _build_group_items(ids, groups, support_by_id) + if len(group_items) == 0: + return [], [] + + total = _support_stats_for_ids(ids, support_by_id) + overall_rate = _positive_rate( + total["positive_residues"], total["valid_residues"]) + target = { + "protein_count": float(int(len(ids) * val_frac)), + "valid_residues": float(total["valid_residues"] * val_frac), + "positive_residues": float(total["positive_residues"] * val_frac), + "negative_residues": float(total["negative_residues"] * val_frac), + } + if target["protein_count"] <= 0.0: + return ids, [] + + rng = random.Random(seed) + rng.shuffle(group_items) + group_items.sort(key=lambda item: _group_sort_key(item, overall_rate)) + + val_groups: Set[str] = set() + val_stats = _zero_stats() + remaining = list(group_items) + + while remaining: + current_score = _balance_score(val_stats, target, overall_rate, total) + candidates: List[Tuple[float, int, int, str, int, Dict[str, Any]]] = [] + for idx, item in enumerate(remaining): + if len(val_groups) + 1 >= len(group_items): + continue + next_stats = _add_stats(val_stats, item["stats"]) + score = _balance_score(next_stats, target, overall_rate, total) + candidates.append(( + score, + abs(int(next_stats["protein_count"]) - int(target["protein_count"])), + -int(next_stats["valid_residues"]), + str(item["group_id"]), + idx, + item, + )) + + if len(candidates) == 0: + break + + best_score, _protein_err, _neg_valid, _group_id, idx, best_item = min( + candidates, + key=lambda row: row[:5], + ) + should_take = ( + int(val_stats["protein_count"]) < int(target["protein_count"]) + or best_score < current_score + ) + if not should_take: + break + + val_groups.add(str(best_item["group_id"])) + val_stats = _add_stats(val_stats, best_item["stats"]) + remaining.pop(idx) + + val_ids = [ + protein_id + for protein_id in ids + if str(groups.get(protein_id, protein_id)) in val_groups + ] + train_ids = [ + protein_id + for protein_id in ids + if str(groups.get(protein_id, protein_id)) not in val_groups + ] + return train_ids, val_ids + + +def build_label_stratified_kfold_splits( + ids: List[str], + n_folds: int, + seed: int, + groups: Dict[str, str], + support_by_id: Mapping[str, Mapping[str, Any]], +) -> List[Tuple[List[str], List[str]]]: + """Build K-fold splits with group integrity and residue label-support balance.""" + ids = [str(protein_id) for protein_id in ids] + if n_folds < 2: + raise ValueError("n_folds must be >= 2") + if len(ids) == 0: + raise ValueError("ids cannot be empty") + + group_items = _build_group_items(ids, groups, support_by_id) + if len(group_items) < n_folds: + raise ValueError( + f"n_folds={n_folds} cannot exceed number of groups={len(group_items)}" + ) + + total = _support_stats_for_ids(ids, support_by_id) + overall_rate = _positive_rate( + total["positive_residues"], total["valid_residues"]) + target = { + "protein_count": float(total["protein_count"] / n_folds), + "valid_residues": float(total["valid_residues"] / n_folds), + "positive_residues": float(total["positive_residues"] / n_folds), + "negative_residues": float(total["negative_residues"] / n_folds), + } + + rng = random.Random(seed) + rng.shuffle(group_items) + group_items.sort(key=lambda item: _group_sort_key(item, overall_rate)) + + fold_groups: List[Set[str]] = [set() for _ in range(n_folds)] + fold_stats: List[Dict[str, int]] = [_zero_stats() for _ in range(n_folds)] + + for item in group_items: + candidates = [] + for fold_idx in range(n_folds): + next_stats = list(fold_stats) + next_stats[fold_idx] = _add_stats(fold_stats[fold_idx], item["stats"]) + score = sum( + _balance_score(stats, target, overall_rate, total) + for stats in next_stats + ) + candidates.append(( + score, + int(fold_stats[fold_idx]["protein_count"]), + len(fold_groups[fold_idx]), + fold_idx, + )) + fold_idx = min(candidates, key=lambda row: row)[3] + fold_groups[fold_idx].add(str(item["group_id"])) + fold_stats[fold_idx] = _add_stats(fold_stats[fold_idx], item["stats"]) + + if any(len(member_groups) == 0 for member_groups in fold_groups): + raise RuntimeError("Label-stratified grouped K-fold assignment produced an empty fold") + + out: List[Tuple[List[str], List[str]]] = [] + for fold_idx in range(n_folds): + val_group_ids = fold_groups[fold_idx] + val_ids = [ + protein_id + for protein_id in ids + if str(groups.get(protein_id, protein_id)) in val_group_ids + ] + train_ids = [ + protein_id + for protein_id in ids + if str(groups.get(protein_id, protein_id)) not in val_group_ids + ] + if len(val_ids) == 0 or len(train_ids) == 0: + raise RuntimeError( + f"Fold {fold_idx + 1} has empty split (train={len(train_ids)}, val={len(val_ids)})" + ) + out.append((train_ids, val_ids)) + return out + + +def summarize_split_ids( + ids: Sequence[str], + support_by_id: Mapping[str, Mapping[str, Any]], + families_by_id: Optional[Mapping[str, str]] = None, +) -> Dict[str, Any]: + """Summarize protein and residue support for a train/validation split.""" + ids = [str(protein_id) for protein_id in ids] + stats = _support_stats_for_ids(ids, support_by_id) + family_counts: Dict[str, int] = {} + shard_counts: Dict[str, int] = {} + families_by_id = families_by_id or {} + for protein_id in ids: + family = str(families_by_id.get(protein_id, "__unavailable__")) + family_counts[family] = family_counts.get(family, 0) + 1 + support = support_by_id.get(protein_id, {}) + shard = support.get("label_shard") + shard_key = str(shard) if shard is not None else "__missing_label__" + shard_counts[shard_key] = shard_counts.get(shard_key, 0) + 1 + + return { + "protein_count": int(stats["protein_count"]), + "valid_residues": int(stats["valid_residues"]), + "positive_residues": int(stats["positive_residues"]), + "negative_residues": int(stats["negative_residues"]), + "positive_rate": _positive_rate( + stats["positive_residues"], stats["valid_residues"]), + "family_counts": dict(sorted(family_counts.items())), + "label_shard_counts": dict(sorted(shard_counts.items())), + "pathogen_counts": {"__unavailable__": int(stats["protein_count"])}, + } + + +def build_split_report( + run_splits: Sequence[Mapping[str, Any]], + support_by_id: Mapping[str, Mapping[str, Any]], + families_by_id: Mapping[str, str], + split_type: str, + split_strategy: str, +) -> Dict[str, Any]: + """Build split report JSON payload for training and Optuna split plans.""" + all_ids: List[str] = [] + seen: Set[str] = set() + entries: List[Dict[str, Any]] = [] + for split in run_splits: + train_ids = [str(x) for x in split.get("train_ids", [])] + val_ids = [str(x) for x in split.get("val_ids", [])] + for protein_id in train_ids + val_ids: + if protein_id not in seen: + seen.add(protein_id) + all_ids.append(protein_id) + entries.append({ + "run_index": split.get("run_index"), + "train_mode": split.get("train_mode"), + "split_seed": split.get("split_seed"), + "train_seed": split.get("train_seed"), + "fold_index": split.get("fold_index"), + "n_folds": split.get("n_folds"), + "ensemble_set_index": split.get("ensemble_set_index"), + "train": summarize_split_ids(train_ids, support_by_id, families_by_id), + "validation": summarize_split_ids(val_ids, support_by_id, families_by_id), + }) + + return { + "schema_version": 1, + "split_type": str(split_type), + "split_strategy": str(split_strategy), + "pathogen_metadata_status": "unavailable", + "all_ids": summarize_split_ids(all_ids, support_by_id, families_by_id), + "runs": entries, + } def split_ids(ids: List[str], val_frac: float, seed: int) -> Tuple[List[str], List[str]]: diff --git a/src/pepseqpred/core/train/threshold.py b/src/pepseqpred/core/train/threshold.py index 16ccb02..a76250f 100644 --- a/src/pepseqpred/core/train/threshold.py +++ b/src/pepseqpred/core/train/threshold.py @@ -1,21 +1,117 @@ -"""threshold.py +"""Threshold selection utilities for PepSeqPred classification outputs.""" -Threshold selection utilities for PepSeqPred classification outputs. +from typing import Any, Dict, Sequence, Tuple -Provides helpers to compute confusion statistics from probabilities and to -select a threshold that maximizes recall subject to a minimum precision. -""" - -from typing import Dict, Tuple, Optional, Any import numpy as np +THRESHOLD_POLICIES: Tuple[str, ...] = ( + "max-recall-min-precision", + "best-f1", + "best-mcc", + "min-recall-max-precision", + "fixed", +) + +DEFAULT_THRESHOLD_GRID: Tuple[float, ...] = ( + 0.05, + 0.10, + 0.20, + 0.30, + 0.40, + 0.50, + 0.60, + 0.70, + 0.80, + 0.90, + 0.95, +) + + def _safe_divide(n: float, d: float) -> float: """Checks for divide by zero before division operation.""" return float(n / d) if d != 0.0 else 0.0 -def _confusion_from_probs(y_true: np.ndarray, y_prob: np.ndarray, threshold: float) -> Tuple[int, int, int, int]: +def _validate_probability(value: float, name: str) -> float: + try: + out = float(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be numeric") from exc + if not np.isfinite(out) or out <= 0.0 or out >= 1.0: + raise ValueError(f"{name} must be in (0.0, 1.0)") + return out + + +def _validate_constraint(value: float, name: str) -> float: + try: + out = float(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be numeric") from exc + if not np.isfinite(out) or out < 0.0 or out > 1.0: + raise ValueError(f"{name} must be in [0.0, 1.0]") + return out + + +def _as_arrays(y_true: np.ndarray, y_prob: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + y_true_np = np.asarray(y_true).reshape(-1).astype(np.int64, copy=False) + y_prob_np = np.asarray(y_prob).reshape(-1).astype(np.float64, copy=False) + if y_true_np.shape[0] != y_prob_np.shape[0]: + raise ValueError( + f"y_true/y_prob length mismatch: {y_true_np.shape[0]} != {y_prob_np.shape[0]}" + ) + if y_prob_np.size > 0 and not bool(np.isfinite(y_prob_np).all()): + raise ValueError("y_prob must contain only finite values") + return y_true_np, y_prob_np + + +def _mcc_from_counts(tp: int, fp: int, tn: int, fn: int) -> float: + denom = float((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)) + if denom <= 0.0: + return 0.0 + return float(((tp * tn) - (fp * fn)) / np.sqrt(denom)) + + +def _row_from_counts( + threshold: float, + tp: int, + fp: int, + tn: int, + fn: int, +) -> Dict[str, Any]: + precision = _safe_divide(tp, tp + fp) + recall = _safe_divide(tp, tp + fn) + f1 = _safe_divide(2.0 * precision * recall, precision + recall) + mcc = _mcc_from_counts(tp, fp, tn, fn) + total = int(tp + fp + tn + fn) + pred_pos = int(tp + fp) + pred_neg = int(tn + fn) + support_pos = int(tp + fn) + support_neg = int(tn + fp) + return { + "threshold": float(threshold), + "tp": int(tp), + "fp": int(fp), + "tn": int(tn), + "fn": int(fn), + "precision": float(precision), + "recall": float(recall), + "f1": float(f1), + "mcc": float(mcc), + "support_pos": support_pos, + "support_neg": support_neg, + "valid_residues": total, + "pred_pos": pred_pos, + "pred_neg": pred_neg, + "pred_pos_frac": _safe_divide(pred_pos, total), + } + + +def _confusion_from_probs( + y_true: np.ndarray, + y_prob: np.ndarray, + threshold: float, +) -> Tuple[int, int, int, int]: """Builds a confusion matrix given model probabilities for threshold calculation.""" y_pred = (y_prob >= threshold).astype(np.int64) y_true = y_true.astype(np.int64) @@ -29,44 +125,19 @@ def _confusion_from_probs(y_true: np.ndarray, y_prob: np.ndarray, threshold: flo return tp, fp, tn, fn -def find_threshold_max_recall_min_precision( +def _row_from_threshold( y_true: np.ndarray, y_prob: np.ndarray, - min_precision: float = 0.50 + threshold: float, ) -> Dict[str, Any]: - """ - Finds the threshold that maximizes recall subject such that `precision >= min_precision`. - If no precision meets this constraint, the best possible threshold is returned. - - Parameters - ---------- - y_true : ndarray - An array of the true labels for a batch of residues. - y_prob : ndarray - An array of the model's estimated probabilities that a residue is an epitope for a given batch. - min_precision : float - Minimum accepted precision while recall is optimized. Default is `0.50`. - - Returns - ------- - Dict[str, Any] - A dictionary containing the most optimal threshold, confusion matrix - used to calculate the most optimal threshold, precision and recall - at that threshold, the minimum precision accepted, and the status - which is either `"ok"` if precision >= `min_precision` otherwise - `"min_precision_unreachable"`. - """ + tp, fp, tn, fn = _confusion_from_probs(y_true, y_prob, threshold) + return _row_from_counts(threshold, tp, fp, tn, fn) + + +def _candidate_rows(y_true: np.ndarray, y_prob: np.ndarray) -> Sequence[Dict[str, Any]]: if y_true.size == 0: - return { - "threshold": float("nan"), - "tp": 0, "fp": 0, "tn": 0, "fn": 0, - "precision": float("nan"), - "recall": float("nan"), - "status": "no_valid_residues", - "min_precision": min_precision - } - - # sort highest to lowest predicted probs + return [] + order = np.argsort(y_prob, kind="mergesort")[::-1] y_sorted = y_true.astype(np.int64, copy=False)[order] p_sorted = y_prob.astype(np.float64, copy=False)[order] @@ -79,74 +150,255 @@ def find_threshold_max_recall_min_precision( total_neg = int(is_neg.sum()) last_index = np.flatnonzero(np.r_[p_sorted[1:] != p_sorted[:-1], True]) - best: Optional[Dict[str, Any]] = None - best_fallback: Optional[Dict[str, Any]] = None - - def _consider(row: Dict[str, Any]) -> None: - """Best threshold and best fallback computation.""" - nonlocal best, best_fallback - if best_fallback is None or ( - row["precision"] > best_fallback["precision"] or ( - row["precision"] == best_fallback["precision"] and - row["threshold"] > best_fallback["threshold"] - ) - ): - best_fallback = row - if row["precision"] >= min_precision: - if best is None or ( - row["recall"] > best["recall"] or ( - row["recall"] == best["recall"] and - row["precision"] > best["precision"] - ) or ( - row["recall"] == best["recall"] and - row["precision"] == best["precision"] and - row["threshold"] > best["threshold"] - ) - ): - best = row - + rows = [] for index in last_index: - thresh = float(p_sorted[index]) tp = int(tp_cum[index]) fp = int(fp_cum[index]) fn = int(total_pos - tp) tn = int(total_neg - fp) - precision = _safe_divide(tp, tp + fp) - recall = _safe_divide(tp, tp + fn) - - row = {"threshold": float(thresh), - "tp": tp, "fp": fp, "tn": tn, "fn": fn, - "precision": precision, - "recall": recall} + rows.append(_row_from_counts(float(p_sorted[index]), tp, fp, tn, fn)) - # get best fallback and best thresholds - _consider(row) - - # if max prob is below valid threshold (<1), represents "predict no positives" + # If possible, include the valid threshold just below 1.0 that predicts + # no positive residues. This keeps the legacy selector behavior. near_one = float(np.nextafter(1.0, 0.0)) if float(p_sorted[0]) < near_one: - _consider({ - "threshold": near_one, - "tp": 0, "fp": 0, "tn": total_neg, "fn": total_pos, - "precision": 0.0, - "recall": 0.0 - }) - - if best is not None: - best["status"] = "ok" - best["min_precision"] = min_precision - return best - - # minimum precision contraint failed - if best_fallback is not None: - best_fallback["status"] = "min_precision_unreachable" - best_fallback["min_precision"] = min_precision - - return best_fallback if best_fallback is not None else { + rows.append(_row_from_counts(near_one, 0, 0, total_neg, total_pos)) + + return rows + + +def _empty_selection( + policy: str, + min_precision: float, + min_recall: float, + fixed_threshold: float, +) -> Dict[str, Any]: + return { "threshold": float("nan"), - "tp": 0, "fp": 0, "tn": 0, "fn": 0, + "tp": 0, + "fp": 0, + "tn": 0, + "fn": 0, "precision": float("nan"), "recall": float("nan"), + "f1": float("nan"), + "mcc": float("nan"), + "support_pos": 0, + "support_neg": 0, + "valid_residues": 0, + "pred_pos": 0, + "pred_neg": 0, + "pred_pos_frac": float("nan"), "status": "no_valid_residues", - "min_precision": min_precision + "policy": policy, + "min_precision": float(min_precision), + "min_recall": float(min_recall), + "fixed_threshold": float(fixed_threshold), } + + +def _with_policy( + row: Dict[str, Any], + *, + policy: str, + status: str, + min_precision: float, + min_recall: float, + fixed_threshold: float, +) -> Dict[str, Any]: + out = dict(row) + out["status"] = status + out["policy"] = policy + out["min_precision"] = float(min_precision) + out["min_recall"] = float(min_recall) + out["fixed_threshold"] = float(fixed_threshold) + return out + + +def _best_row(rows: Sequence[Dict[str, Any]], keys: Sequence[Tuple[str, bool]]) -> Dict[str, Any]: + def key_fn(row: Dict[str, Any]) -> Tuple[float, ...]: + values = [] + for key, higher_is_better in keys: + value = float(row[key]) + values.append(value if higher_is_better else -value) + return tuple(values) + + return max(rows, key=key_fn) + + +def select_threshold( + y_true: np.ndarray, + y_prob: np.ndarray, + *, + policy: str = "max-recall-min-precision", + min_precision: float = 0.25, + min_recall: float = 0.80, + fixed_threshold: float = 0.50, +) -> Dict[str, Any]: + """Selects a decision threshold from validation labels and probabilities.""" + policy = str(policy).strip().lower() + if policy not in THRESHOLD_POLICIES: + raise ValueError( + f"Unsupported threshold policy '{policy}'. Expected one of: {', '.join(THRESHOLD_POLICIES)}" + ) + min_precision = _validate_constraint(min_precision, "min_precision") + min_recall = _validate_constraint(min_recall, "min_recall") + fixed_threshold = _validate_probability(fixed_threshold, "fixed_threshold") + y_true_np, y_prob_np = _as_arrays(y_true, y_prob) + + if y_true_np.size == 0: + return _empty_selection(policy, min_precision, min_recall, fixed_threshold) + + if policy == "fixed": + row = _row_from_threshold(y_true_np, y_prob_np, fixed_threshold) + return _with_policy( + row, + policy=policy, + status="ok", + min_precision=min_precision, + min_recall=min_recall, + fixed_threshold=fixed_threshold, + ) + + rows = _candidate_rows(y_true_np, y_prob_np) + if len(rows) == 0: + return _empty_selection(policy, min_precision, min_recall, fixed_threshold) + + if policy == "max-recall-min-precision": + eligible = [row for row in rows if float(row["precision"]) >= min_precision] + if eligible: + row = _best_row( + eligible, + ( + ("recall", True), + ("precision", True), + ("threshold", True), + ), + ) + return _with_policy( + row, + policy=policy, + status="ok", + min_precision=min_precision, + min_recall=min_recall, + fixed_threshold=fixed_threshold, + ) + row = _best_row(rows, (("precision", True), ("threshold", True))) + return _with_policy( + row, + policy=policy, + status="min_precision_unreachable", + min_precision=min_precision, + min_recall=min_recall, + fixed_threshold=fixed_threshold, + ) + + if policy == "best-f1": + row = _best_row( + rows, + ( + ("f1", True), + ("precision", True), + ("recall", True), + ("threshold", True), + ), + ) + return _with_policy( + row, + policy=policy, + status="ok", + min_precision=min_precision, + min_recall=min_recall, + fixed_threshold=fixed_threshold, + ) + + if policy == "best-mcc": + row = _best_row( + rows, + ( + ("mcc", True), + ("f1", True), + ("precision", True), + ("recall", True), + ("threshold", True), + ), + ) + return _with_policy( + row, + policy=policy, + status="ok", + min_precision=min_precision, + min_recall=min_recall, + fixed_threshold=fixed_threshold, + ) + + eligible = [row for row in rows if float(row["recall"]) >= min_recall] + if eligible: + row = _best_row( + eligible, + ( + ("precision", True), + ("recall", True), + ("threshold", True), + ), + ) + return _with_policy( + row, + policy=policy, + status="ok", + min_precision=min_precision, + min_recall=min_recall, + fixed_threshold=fixed_threshold, + ) + + row = _best_row( + rows, + ( + ("recall", True), + ("precision", True), + ("threshold", True), + ), + ) + return _with_policy( + row, + policy=policy, + status="min_recall_unreachable", + min_precision=min_precision, + min_recall=min_recall, + fixed_threshold=fixed_threshold, + ) + + +def threshold_diagnostic_grid( + y_true: np.ndarray, + y_prob: np.ndarray, + thresholds: Sequence[float] = DEFAULT_THRESHOLD_GRID, +) -> Sequence[Dict[str, Any]]: + """Computes fixed-threshold diagnostics for validation/evaluation payloads.""" + y_true_np, y_prob_np = _as_arrays(y_true, y_prob) + if y_true_np.size == 0: + return [] + rows = [] + for threshold in thresholds: + threshold_f = _validate_probability(threshold, "threshold") + row = _row_from_threshold(y_true_np, y_prob_np, threshold_f) + rows.append(row) + return rows + + +def find_threshold_max_recall_min_precision( + y_true: np.ndarray, + y_prob: np.ndarray, + min_precision: float = 0.50, +) -> Dict[str, Any]: + """ + Finds the threshold that maximizes recall subject to precision >= min_precision. + + This compatibility wrapper preserves the public helper used by older code. + """ + return select_threshold( + y_true, + y_prob, + policy="max-recall-min-precision", + min_precision=min_precision, + ) diff --git a/src/pepseqpred/core/train/trainer.py b/src/pepseqpred/core/train/trainer.py index 4f317f9..c46037a 100644 --- a/src/pepseqpred/core/train/trainer.py +++ b/src/pepseqpred/core/train/trainer.py @@ -7,10 +7,9 @@ """ import logging -import contextlib from pathlib import Path from dataclasses import dataclass -from typing import Optional, List, Dict, Any, Tuple +from typing import Optional, List, Dict, Any, Tuple, Iterator import torch import torch.nn as nn from torch.nn.parallel import DistributedDataParallel as TorchDDP @@ -19,8 +18,15 @@ import optuna from .ddp import ddp_rank, ddp_all_reduce_sum, ddp_gather_all_1d from .metrics import compute_eval_metrics -from .threshold import find_threshold_max_recall_min_precision +from .threshold import select_threshold, threshold_diagnostic_grid from .curveartifacts import write_validation_curve_artifacts +from pepseqpred.core.models.factory import PepSeqModelConfig, model_config_to_dict + + +Batch = ( + Tuple[torch.Tensor, torch.Tensor] + | Tuple[torch.Tensor, torch.Tensor, torch.Tensor] +) @dataclass @@ -34,6 +40,10 @@ class TrainerConfig: # should only train using GPUs (but can be changed to "cpu") device: str = "cuda" pos_weight: Optional[float] = None + threshold_policy: str = "max-recall-min-precision" + threshold_min_precision: float = 0.25 + threshold_min_recall: float = 0.80 + threshold_fixed_value: float = 0.50 @dataclass(frozen=True) @@ -66,11 +76,13 @@ def __init__(self, model: nn.Module, train_loader: DataLoader, logger: logging.Logger, val_loader: Optional[DataLoader] = None, - config: TrainerConfig = TrainerConfig()): + config: TrainerConfig = TrainerConfig(), + model_config: Optional[PepSeqModelConfig] = None): self.model = model self.train_loader = train_loader self.val_loader = val_loader self.config = config + self.model_config = model_config self.logger = logger self.device = torch.device( @@ -102,10 +114,102 @@ def __init__(self, model: nn.Module, "weight_decay": self.config.weight_decay, "num_params": num_params, "has_val_loader": self.val_loader is not None, - "pos_weight": self.config.pos_weight + "pos_weight": self.config.pos_weight, + "threshold_policy": self.config.threshold_policy, + "threshold_min_precision": self.config.threshold_min_precision, + "threshold_min_recall": self.config.threshold_min_recall, + "threshold_fixed_value": self.config.threshold_fixed_value }}) - def _batch_step(self, batch: torch.Tensor, train: bool = True) -> Dict[str, Any]: + @staticmethod + def _make_zero_valid_dummy_batch(batch: Batch) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build a minimal graph-compatible batch with no valid residues.""" + if len(batch) == 2: + X, y = batch + mask = None + elif len(batch) == 3: + X, y, mask = batch + else: + raise ValueError( + f"Expected a two- or three-tensor batch, got {len(batch)} elements") + + if X.dim() != 3: + raise ValueError( + f"Expected embedding shape (B, L, D), got {tuple(X.shape)}") + if y.dim() != 2: + raise ValueError( + f"Expected target shape (B, L), got {tuple(y.shape)}") + + dummy_x = X.new_zeros((1, 1, X.size(-1))) + dummy_y = y.new_zeros((1, 1)) + if mask is None: + dummy_mask = torch.zeros( + (1, 1), dtype=torch.long, device=y.device) + else: + dummy_mask = mask.new_zeros((1, 1)) + return dummy_x, dummy_y, dummy_mask + + def _synchronized_training_batches( + self, + loader: DataLoader + ) -> Iterator[Tuple[Batch, bool]]: + """Yield real or zero-valid batches until every DDP rank is exhausted. + + Every rank performs the same active-rank collective before a forward + pass. Once a rank exhausts its local loader, it continues yielding a + minimal zero-valid batch so custom loss collectives, DDP backward, and + optimizer state remain synchronized. The current model heads do not use + batch-statistics layers; adding BatchNorm or SyncBatchNorm requires + revisiting this dummy-forward behavior. + """ + iterator = iter(loader) + dummy_batch: Optional[Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor + ]] = None + first_iteration = True + + while True: + try: + batch = next(iterator) + has_local_batch = True + if dummy_batch is None: + dummy_batch = self._make_zero_valid_dummy_batch(batch) + except StopIteration: + batch = None + has_local_batch = False + + sync_state = torch.tensor( + [1 if has_local_batch else 0, 1], + device=self.device, + dtype=torch.int64, + ) + sync_state = ddp_all_reduce_sum(sync_state) + active_ranks = int(sync_state[0].item()) + world_size = int(sync_state[1].item()) + + if active_ranks == 0: + return + + if first_iteration and active_ranks != world_size: + raise RuntimeError( + "At least one DDP rank has no initial training batch; " + "training requires every rank to start with data" + ) + first_iteration = False + + if has_local_batch: + if batch is None: + raise RuntimeError("Training iterator lost a local batch") + yield batch, True + else: + if dummy_batch is None: + raise RuntimeError( + "Cannot construct a synchronized dummy batch before " + "receiving a real training batch" + ) + yield dummy_batch, False + + def _batch_step(self, batch: Batch, train: bool = True) -> Dict[str, Any]: """Steps through a batch to train and optimize the model.""" # ensure we grab mask when applicable if len(batch) == 2: @@ -140,23 +244,49 @@ def _batch_step(self, batch: torch.Tensor, train: bool = True) -> Dict[str, Any] if mask.shape != y.shape: raise ValueError( f"Expected mask shape {tuple(y.shape)}, got {tuple(mask.shape)}") - denom = mask.float().sum() - if denom.item() == 0.0: - loss = loss_raw.sum() * 0.0 - # if train: - # return {"loss": 0.0, "n": 0} - else: - loss = (loss_raw * mask.float()).sum() / denom + mask_float = mask.float() + loss_num = (loss_raw * mask_float).sum() + n_tensor = mask_float.sum() + else: + loss_num = loss_raw.sum() + n_tensor = loss_raw.new_tensor(y.numel()) + + if n_tensor.item() == 0.0: + loss = loss_num * 0.0 else: - loss = loss_raw.mean() + loss = loss_num / n_tensor + + n = int(n_tensor.item()) # optimize model in training batch if train: self.optimizer.zero_grad(set_to_none=True) - loss.backward() + # keep DDP ranks aligned before deciding whether to skip backward + step_stats = ddp_all_reduce_sum(torch.stack(( + n_tensor.detach().to(dtype=torch.float64), + torch.ones((), device=y.device, dtype=torch.float64) + ))) + global_valid_count = int(step_stats[0].item()) + # complete DDP's backward sequence without advancing Adam when the + # entire synchronized step contains no valid residues + if global_valid_count == 0: + (loss_num * 0.0).backward() + return { + "loss": float(loss.item()), + "n": n, + "global_valid_count": 0, + "optimizer_step": False + } + train_loss = loss_num * ( + step_stats[1] / step_stats[0]).to(dtype=loss_num.dtype) + train_loss.backward() self.optimizer.step() - n = int(mask.sum().item() if mask is not None else int(y.numel())) - return {"loss": float(loss.item()), "n": n} + return { + "loss": float(loss.item()), + "n": n, + "global_valid_count": global_valid_count, + "optimizer_step": True + } probs = torch.sigmoid(logits) # (B, L) y_flat = y.to(torch.long).view(-1) @@ -197,18 +327,32 @@ def _run_epoch( total_samples = 0 total_pos = 0 total_neg = 0 + synchronized_steps = 0 + optimizer_steps = 0 + zero_valid_steps = 0 + real_batches = 0 + dummy_batches = 0 all_y: List[torch.Tensor] = [] all_probs: List[torch.Tensor] = [] + if train: + batches = self._synchronized_training_batches(loader) + else: + batches = ((batch, True) for batch in loader) + # use inference mode for eval torch_ctx = torch.enable_grad() if train else torch.inference_mode() - loop_ctx = self.model.join() if train and hasattr( - self.model, "join") else contextlib.nullcontext() - with loop_ctx: - for batch in loader: - with torch_ctx: - out = self._batch_step(batch, train=train) + with torch_ctx: + for batch, is_real_batch in batches: + out = self._batch_step(batch, train=train) + + if train: + synchronized_steps += 1 + real_batches += int(is_real_batch) + dummy_batches += int(not is_real_batch) + optimizer_steps += int(out["optimizer_step"]) + zero_valid_steps += int(out["global_valid_count"] == 0) # collect data for metrics if not train and out["n"] > 0: @@ -259,6 +403,47 @@ def _run_epoch( total_pos = int(pos_sum.item()) total_neg = int(neg_sum.item()) + sync_summary = None + if train: + local_sync_stats = torch.tensor( + [real_batches, dummy_batches], + device=self.device, + dtype=torch.int64 + ) + gathered_sync_stats, sync_sizes = ddp_gather_all_1d( + local_sync_stats, self.device) + if ddp_rank() == 0: + per_rank = [] + for rank_index, size in enumerate(sync_sizes): + if size != 2: + raise RuntimeError( + "Expected two training synchronization counters " + f"from rank {rank_index}, got {size}" + ) + rank_stats = gathered_sync_stats[rank_index][:size] + per_rank.append({ + "rank": int(rank_index), + "real_batches": int(rank_stats[0].item()), + "dummy_batches": int(rank_stats[1].item()) + }) + + total_dummy_batches = sum( + entry["dummy_batches"] for entry in per_rank) + total_rank_steps = synchronized_steps * len(per_rank) + sync_summary = { + "synchronized_steps": int(synchronized_steps), + "optimizer_steps": int(optimizer_steps), + "zero_valid_steps": int(zero_valid_steps), + "dummy_batch_fraction": ( + float(total_dummy_batches / total_rank_steps) + if total_rank_steps > 0 + else 0.0 + ), + "per_rank": per_rank + } + self.logger.info( + "train_sync_summary", extra={"extra": sync_summary}) + # compute eval metrics cm = None eval_metrics = None @@ -308,14 +493,37 @@ def _run_epoch( "auc10": float("nan"), "pr_auc": float("nan"), "threshold": float("nan"), + "threshold_policy": str(self.config.threshold_policy), "threshold_status": "no_valid_residues", - "threshold_min_precision": 0.25 + "threshold_min_precision": float(self.config.threshold_min_precision), + "threshold_min_recall": float(self.config.threshold_min_recall), + "threshold_fixed_value": float(self.config.threshold_fixed_value), + "threshold_precision": float("nan"), + "threshold_recall": float("nan"), + "threshold_f1": float("nan"), + "threshold_mcc": float("nan"), + "threshold_tp": 0, + "threshold_fp": 0, + "threshold_tn": 0, + "threshold_fn": 0, + "threshold_support_pos": 0, + "threshold_support_neg": 0, + "threshold_pred_pos_residues": 0, + "threshold_pred_neg_residues": 0, + "threshold_pred_pos_frac": float("nan"), + "threshold_grid": [] } avg_acc = float("nan") else: # compute predictions at most optimal threshold calculated - thresh_out = find_threshold_max_recall_min_precision( - y_true, y_prob, min_precision=0.25) + thresh_out = select_threshold( + y_true, + y_prob, + policy=self.config.threshold_policy, + min_precision=self.config.threshold_min_precision, + min_recall=self.config.threshold_min_recall, + fixed_threshold=self.config.threshold_fixed_value + ) best_thresh = float(thresh_out["threshold"]) y_pred = (y_prob >= best_thresh).astype(np.int64) @@ -332,8 +540,31 @@ def _run_epoch( eval_metrics = compute_eval_metrics(y_true, y_pred, y_prob) eval_metrics["threshold"] = best_thresh + eval_metrics["threshold_policy"] = thresh_out["policy"] eval_metrics["threshold_status"] = thresh_out["status"] eval_metrics["threshold_min_precision"] = thresh_out["min_precision"] + eval_metrics["threshold_min_recall"] = thresh_out["min_recall"] + eval_metrics["threshold_fixed_value"] = thresh_out["fixed_threshold"] + eval_metrics["threshold_precision"] = thresh_out["precision"] + eval_metrics["threshold_recall"] = thresh_out["recall"] + eval_metrics["threshold_f1"] = thresh_out["f1"] + eval_metrics["threshold_mcc"] = thresh_out["mcc"] + eval_metrics["threshold_tp"] = int(thresh_out["tp"]) + eval_metrics["threshold_fp"] = int(thresh_out["fp"]) + eval_metrics["threshold_tn"] = int(thresh_out["tn"]) + eval_metrics["threshold_fn"] = int(thresh_out["fn"]) + eval_metrics["threshold_support_pos"] = int( + thresh_out["support_pos"]) + eval_metrics["threshold_support_neg"] = int( + thresh_out["support_neg"]) + eval_metrics["threshold_pred_pos_residues"] = int( + thresh_out["pred_pos"]) + eval_metrics["threshold_pred_neg_residues"] = int( + thresh_out["pred_neg"]) + eval_metrics["threshold_pred_pos_frac"] = float( + thresh_out["pred_pos_frac"]) + eval_metrics["threshold_grid"] = list( + threshold_diagnostic_grid(y_true, y_prob)) # log confusion matrix if (not train and cm is not None @@ -351,8 +582,12 @@ def _run_epoch( "balanced_acc": balanced_acc, "per_class_acc": per_class_acc.tolist(), "threshold": eval_metrics["threshold"], + "threshold_policy": eval_metrics["threshold_policy"], "threshold_status": eval_metrics["threshold_status"], "threshold_min_precision": eval_metrics["threshold_min_precision"], + "threshold_min_recall": eval_metrics["threshold_min_recall"], + "threshold_fixed_value": eval_metrics["threshold_fixed_value"], + "threshold_pred_pos_frac": eval_metrics["threshold_pred_pos_frac"], "precision": eval_metrics["precision"], "recall": eval_metrics["recall"], "f1": eval_metrics["f1"], @@ -365,7 +600,16 @@ def _run_epoch( # handle training vs eval output out = {"loss": avg_loss, "n_residues": total_samples, "pos_residues": total_pos, "neg_residues": total_neg} - if not train: + if train: + out.update({ + "synchronized_steps": int(synchronized_steps), + "optimizer_steps": int(optimizer_steps), + "zero_valid_steps": int(zero_valid_steps), + "real_batches": int(real_batches), + "dummy_batches": int(dummy_batches), + "sync_summary": sync_summary + }) + else: out["acc"] = avg_acc out["eval_metrics"] = eval_metrics if capture_eval_arrays and ddp_rank() == 0: @@ -551,6 +795,11 @@ def _save_checkpoint(self, path: Path | str, epoch: int, loss: float, metrics: O "optim_state_dict": self.optimizer.state_dict(), "epoch": epoch, "config": self.config.__dict__, + "model_config": ( + model_config_to_dict(self.model_config) + if self.model_config is not None + else None + ), "best_loss": loss, "metrics": metrics} torch.save(state, path) diff --git a/src/pepseqpred/core/train/weights.py b/src/pepseqpred/core/train/weights.py index 167511c..a6c39b8 100644 --- a/src/pepseqpred/core/train/weights.py +++ b/src/pepseqpred/core/train/weights.py @@ -45,6 +45,37 @@ def compute_pos_neg_counts(loader: DataLoader) -> Tuple[int, int]: return pos, neg +def global_pos_neg_counts( + local_pos: int, + local_neg: int, + ddp: Dict[str, Any] | None +) -> Tuple[int, int]: + """ + Aggregate positive/negative residue counts across DDP ranks. + + Parameters + ---------- + local_pos : int + Local count of positive residues. + local_neg : int + Local count of negative residues. + ddp : Dict[str, Any] | None + DDP metadata dict, or `None` if DDP is disabled. + + Returns + ------- + Tuple[int, int] + Global `(pos_count, neg_count)` across ranks when DDP is enabled. + """ + if ddp is None: + return int(local_pos), int(local_neg) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + t = torch.tensor([local_pos, local_neg], device=device, dtype=torch.long) + dist.all_reduce(t, op=dist.ReduceOp.SUM) + return int(t[0].item()), int(t[1].item()) + + def global_pos_weight(local_pos: int, local_neg: int, ddp: Dict[str, Any] | None) -> float: """ Compute global negative/positive class weight across ranks if DDP is running. @@ -63,12 +94,7 @@ def global_pos_weight(local_pos: int, local_neg: int, ddp: Dict[str, Any] | None float Ratio of negatives to positives, aggregated across ranks when DDP is enabled. """ - if ddp is None: - return float(local_neg / max(local_pos, 1)) - t = torch.tensor([local_pos, local_neg], device=torch.device("cuda")) - dist.all_reduce(t, op=dist.ReduceOp.SUM) - pos = int(t[0].item()) - neg = int(t[1].item()) + pos, neg = global_pos_neg_counts(local_pos, local_neg, ddp) return float(neg / max(pos, 1)) diff --git a/tests/e2e/test_train_to_predict_e2e.py b/tests/e2e/test_train_to_predict_e2e.py index 87806c3..9845792 100644 --- a/tests/e2e/test_train_to_predict_e2e.py +++ b/tests/e2e/test_train_to_predict_e2e.py @@ -38,7 +38,7 @@ def to(self, _device): def __call__(self, batch_tokens, repr_layers, return_contacts=False): _ = return_contacts batch_size, token_len = batch_tokens.shape - # append_seq_len -> final emb dim is 4 (matches training fixture) + # Base fake ESM dim is 3; raw seq_len_feature makes prediction dim 4. rep_dim = 3 reps = torch.ones((batch_size, token_len, rep_dim), dtype=torch.float32) @@ -62,7 +62,7 @@ def test_train_then_predict_e2e(training_artifacts, tmp_path: Path, monkeypatch) train_cmd = [ sys.executable, "-m", - "pepseqpred.apps.train_ffnn_cli", + "pepseqpred.apps.train_cli", "--embedding-dirs", str(training_artifacts["embedding_dir"]), "--label-shards", @@ -77,6 +77,8 @@ def test_train_then_predict_e2e(training_artifacts, tmp_path: Path, monkeypatch) "8", "--dropouts", "0.1", + "--seq-len-feature", + "raw", "--val-frac", "0.5", "--split-seeds", diff --git a/tests/integration/test_eval_ffnn_cli_smoke.py b/tests/integration/test_eval_ffnn_cli_smoke.py index 4755fe4..4e149af 100644 --- a/tests/integration/test_eval_ffnn_cli_smoke.py +++ b/tests/integration/test_eval_ffnn_cli_smoke.py @@ -1,9 +1,16 @@ import json import sys from pathlib import Path + import pytest import torch + import pepseqpred.apps.evaluate_ffnn_cli as evaluate_ffnn_cli +from pepseqpred.core.models.factory import ( + PepSeqModelConfig, + build_pepseq_model, + model_config_to_dict, +) from pepseqpred.core.models.ffnn import PepSeqFFNN pytestmark = pytest.mark.integration @@ -30,7 +37,37 @@ def _write_checkpoint(path: Path, threshold: float = 0.5) -> None: ) -def test_eval_ffnn_cli_single_checkpoint_smoke(training_artifacts, tmp_path: Path, monkeypatch): +def _write_conv_checkpoint(path: Path, threshold: float = 0.5) -> None: + cfg = PepSeqModelConfig( + emb_dim=4, + hidden_sizes=(3,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False, + model_head="conv1d", + conv_channels=3, + conv_layers=1, + conv_kernel_size=3, + conv_dropout=0.0, + ) + model = build_pepseq_model(cfg) + for param in model.parameters(): + torch.nn.init.constant_(param, 0.0) + + torch.save( + { + "model_state_dict": model.state_dict(), + "model_config": model_config_to_dict(cfg), + "metrics": {"threshold": float(threshold)}, + }, + path, + ) + + +def test_eval_ffnn_cli_single_checkpoint_smoke( + training_artifacts, tmp_path: Path, monkeypatch +): checkpoint = tmp_path / "single.pt" _write_checkpoint(checkpoint, threshold=0.6) @@ -63,15 +100,56 @@ def test_eval_ffnn_cli_single_checkpoint_smoke(training_artifacts, tmp_path: Pat assert payload["artifact_mode"] == "single-checkpoint" assert payload["n_members"] == 1 assert payload["threshold"] == pytest.approx(0.6) + assert payload["ensemble_aggregation"] == "majority" + assert payload["ensemble_threshold"] is None assert eval_out["processed_proteins"] == 4 assert eval_out["valid_residues"] == 24 assert eval_out["pred_pos_residues"] == 0 + assert eval_out["ensemble_aggregation"] == "single-model" + assert isinstance(eval_out["threshold_grid"], list) assert metrics["precision"] == pytest.approx(0.0) assert metrics["recall"] == pytest.approx(0.0) assert metrics["f1"] == pytest.approx(0.0) -def test_eval_ffnn_cli_manifest_majority_vote_smoke(training_artifacts, tmp_path: Path, monkeypatch): +def test_eval_ffnn_cli_loads_conv_checkpoint_metadata( + training_artifacts, tmp_path: Path, monkeypatch +): + checkpoint = tmp_path / "conv_single.pt" + _write_conv_checkpoint(checkpoint, threshold=0.4) + + output_json = tmp_path / "eval_conv.json" + monkeypatch.setattr( + sys, + "argv", + [ + "eval_ffnn_cli.py", + str(checkpoint), + "--embedding-dirs", + str(training_artifacts["embedding_dir"]), + "--label-shards", + str(training_artifacts["label_shard"]), + "--batch-size", + "2", + "--num-workers", + "0", + "--output-json", + str(output_json), + ], + ) + + evaluate_ffnn_cli.main() + + payload = json.loads(output_json.read_text(encoding="utf-8")) + assert payload["model_cfg_src"] == "checkpoint" + assert payload["model_head"] == "conv1d" + assert payload["conv_channels"] == 3 + assert payload["evaluation"]["processed_proteins"] == 4 + + +def test_eval_ffnn_cli_manifest_majority_vote_smoke( + training_artifacts, tmp_path: Path, monkeypatch +): ckpt_1 = tmp_path / "fold_1.pt" ckpt_2 = tmp_path / "fold_2.pt" ckpt_3 = tmp_path / "fold_3.pt" @@ -137,9 +215,98 @@ def test_eval_ffnn_cli_manifest_majority_vote_smoke(training_artifacts, tmp_path assert payload["artifact_mode"] == "ensemble-manifest" assert payload["n_members"] == 3 assert payload["threshold"] is None + assert payload["ensemble_aggregation"] == "majority" + assert payload["ensemble_threshold"] is None assert payload["member_thresholds"] == [0.6, 0.6, 0.4] assert eval_out["votes_needed"] == 2 + assert eval_out["ensemble_aggregation"] == "majority" + assert eval_out["ensemble_threshold"] is None + assert isinstance(eval_out["threshold_grid"], list) assert eval_out["valid_residues"] == 24 assert eval_out["pred_pos_residues"] == 0 assert metrics["precision"] == pytest.approx(0.0) assert metrics["recall"] == pytest.approx(0.0) + + +def test_eval_ffnn_cli_manifest_mean_prob_smoke( + training_artifacts, tmp_path: Path, monkeypatch +): + ckpt_1 = tmp_path / "fold_1.pt" + ckpt_2 = tmp_path / "fold_2.pt" + ckpt_3 = tmp_path / "fold_3.pt" + _write_checkpoint(ckpt_1, threshold=0.6) + _write_checkpoint(ckpt_2, threshold=0.6) + _write_checkpoint(ckpt_3, threshold=0.4) + + manifest = { + "schema_version": 1, + "members": [ + { + "member_index": 1, + "fold_index": 1, + "status": "OK", + "checkpoint": str(ckpt_1), + "threshold": 0.6, + }, + { + "member_index": 2, + "fold_index": 2, + "status": "OK", + "checkpoint": str(ckpt_2), + "threshold": 0.6, + }, + { + "member_index": 3, + "fold_index": 3, + "status": "OK", + "checkpoint": str(ckpt_3), + "threshold": 0.4, + }, + ], + } + manifest_path = tmp_path / "ensemble_manifest.json" + manifest_path.write_text(json.dumps(manifest), encoding="utf-8") + + output_json = tmp_path / "eval_manifest_mean_prob.json" + monkeypatch.setattr( + sys, + "argv", + [ + "eval_ffnn_cli.py", + str(manifest_path), + "--embedding-dirs", + str(training_artifacts["embedding_dir"]), + "--label-shards", + str(training_artifacts["label_shard"]), + "--batch-size", + "2", + "--num-workers", + "0", + "--ensemble-aggregation", + "mean-prob", + "--ensemble-threshold", + "0.49", + "--output-json", + str(output_json), + ], + ) + + evaluate_ffnn_cli.main() + + payload = json.loads(output_json.read_text(encoding="utf-8")) + eval_out = payload["evaluation"] + metrics = eval_out["metrics"] + + assert payload["artifact_mode"] == "ensemble-manifest" + assert payload["n_members"] == 3 + assert payload["threshold"] is None + assert payload["ensemble_aggregation"] == "mean-prob" + assert payload["ensemble_threshold"] == pytest.approx(0.49) + assert payload["member_thresholds"] == [0.6, 0.6, 0.4] + assert eval_out["votes_needed"] is None + assert eval_out["ensemble_aggregation"] == "mean-prob" + assert eval_out["ensemble_threshold"] == pytest.approx(0.49) + assert isinstance(eval_out["threshold_grid"], list) + assert eval_out["valid_residues"] == 24 + assert eval_out["pred_pos_residues"] == 24 + assert metrics["recall"] == pytest.approx(1.0) diff --git a/tests/integration/test_prediction_cli_smoke.py b/tests/integration/test_prediction_cli_smoke.py index 202ce0a..753c420 100644 --- a/tests/integration/test_prediction_cli_smoke.py +++ b/tests/integration/test_prediction_cli_smoke.py @@ -5,6 +5,11 @@ import pytest import torch import pepseqpred.apps.prediction_cli as prediction_cli +from pepseqpred.core.models.factory import ( + PepSeqModelConfig, + build_pepseq_model, + model_config_to_dict, +) from pepseqpred.core.models.ffnn import PepSeqFFNN pytestmark = pytest.mark.integration @@ -38,7 +43,7 @@ def to(self, _device): def __call__(self, batch_tokens, repr_layers, return_contacts=False): _ = return_contacts batch_size, token_len = batch_tokens.shape - rep_dim = 3 # append_seq_len -> final emb dim is 4 + rep_dim = 3 reps = torch.ones((batch_size, token_len, rep_dim), dtype=torch.float32) return {"representations": {repr_layers[0]: reps}} @@ -46,8 +51,8 @@ def __call__(self, batch_tokens, repr_layers, return_contacts=False): def _write_checkpoint(path: Path, threshold: float = 0.5) -> None: model = PepSeqFFNN( - emb_dim=4, - hidden_sizes=(3,), + emb_dim=3, + hidden_sizes=(4,), dropouts=(0.0,), use_layer_norm=False, use_residual=False, @@ -65,6 +70,34 @@ def _write_checkpoint(path: Path, threshold: float = 0.5) -> None: ) +def _write_conv_checkpoint(path: Path, threshold: float = 0.5) -> None: + cfg = PepSeqModelConfig( + emb_dim=3, + hidden_sizes=(3,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False, + model_head="conv1d", + conv_channels=3, + conv_layers=1, + conv_kernel_size=3, + conv_dropout=0.0, + ) + model = build_pepseq_model(cfg) + for param in model.parameters(): + torch.nn.init.constant_(param, 0.0) + + torch.save( + { + "model_state_dict": model.state_dict(), + "model_config": model_config_to_dict(cfg), + "metrics": {"threshold": float(threshold)}, + }, + path, + ) + + def test_prediction_cli_smoke(monkeypatch, tmp_path: Path): fake_pretrained = types.SimpleNamespace( fake_model=lambda: (FakeESMModel(), FakeAlphabet()) @@ -108,6 +141,40 @@ def test_prediction_cli_smoke(monkeypatch, tmp_path: Path): assert set(lines[3]).issubset({"0", "1"}) +def test_prediction_cli_loads_conv_checkpoint_metadata(monkeypatch, tmp_path: Path): + fake_pretrained = types.SimpleNamespace( + fake_model=lambda: (FakeESMModel(), FakeAlphabet()) + ) + monkeypatch.setattr(prediction_cli.esm, "pretrained", fake_pretrained) + + checkpoint = tmp_path / "conv_model.pt" + _write_conv_checkpoint(checkpoint) + + fasta = tmp_path / "input.fasta" + fasta.write_text(">protein_1\nACDEFG\n", encoding="utf-8") + output_fasta = tmp_path / "predictions.fasta" + + monkeypatch.setattr( + sys, + "argv", + [ + "prediction_cli.py", + str(checkpoint), + str(fasta), + "--output-fasta", + str(output_fasta), + "--model-name", + "fake_model", + ], + ) + + prediction_cli.main() + + lines = [line.strip() for line in output_fasta.read_text( + encoding="utf-8").splitlines() if line.strip()] + assert lines == [">protein_1", "111111"] + + def test_prediction_cli_manifest_v1_majority_vote(monkeypatch, tmp_path: Path): fake_pretrained = types.SimpleNamespace( fake_model=lambda: (FakeESMModel(), FakeAlphabet()) @@ -158,6 +225,60 @@ def test_prediction_cli_manifest_v1_majority_vote(monkeypatch, tmp_path: Path): assert lines[1] == "000000" +def test_prediction_cli_manifest_mean_prob_aggregation(monkeypatch, tmp_path: Path): + fake_pretrained = types.SimpleNamespace( + fake_model=lambda: (FakeESMModel(), FakeAlphabet()) + ) + monkeypatch.setattr(prediction_cli.esm, "pretrained", fake_pretrained) + + ckpt_1 = tmp_path / "fold_1.pt" + ckpt_2 = tmp_path / "fold_2.pt" + ckpt_3 = tmp_path / "fold_3.pt" + _write_checkpoint(ckpt_1, threshold=0.6) + _write_checkpoint(ckpt_2, threshold=0.6) + _write_checkpoint(ckpt_3, threshold=0.4) + + manifest = { + "schema_version": 1, + "members": [ + {"member_index": 1, "fold_index": 1, "status": "OK", "checkpoint": str(ckpt_1), "threshold": 0.6}, + {"member_index": 2, "fold_index": 2, "status": "OK", "checkpoint": str(ckpt_2), "threshold": 0.6}, + {"member_index": 3, "fold_index": 3, "status": "OK", "checkpoint": str(ckpt_3), "threshold": 0.4}, + ] + } + manifest_path = tmp_path / "ensemble_manifest.json" + manifest_path.write_text(json.dumps(manifest), encoding="utf-8") + + fasta = tmp_path / "input.fasta" + fasta.write_text(">protein_1\nACDEFG\n", encoding="utf-8") + output_fasta = tmp_path / "predictions.fasta" + + monkeypatch.setattr( + sys, + "argv", + [ + "prediction_cli.py", + str(manifest_path), + str(fasta), + "--output-fasta", + str(output_fasta), + "--model-name", + "fake_model", + "--ensemble-aggregation", + "mean-prob", + "--ensemble-threshold", + "0.49", + ], + ) + + prediction_cli.main() + + lines = [line.strip() for line in output_fasta.read_text( + encoding="utf-8").splitlines() if line.strip()] + assert lines[0] == ">protein_1" + assert lines[1] == "111111" + + def test_prediction_cli_manifest_v2_set_index_and_k_folds(monkeypatch, tmp_path: Path): fake_pretrained = types.SimpleNamespace( fake_model=lambda: (FakeESMModel(), FakeAlphabet()) diff --git a/tests/integration/test_prepare_dataset_multisource_pipeline.py b/tests/integration/test_prepare_dataset_multisource_pipeline.py index 4aaf848..60a8bfc 100644 --- a/tests/integration/test_prepare_dataset_multisource_pipeline.py +++ b/tests/integration/test_prepare_dataset_multisource_pipeline.py @@ -8,7 +8,7 @@ import pepseqpred.apps.esm_cli as esm_cli import pepseqpred.apps.labels_cli as labels_cli -import pepseqpred.apps.train_ffnn_cli as train_cli +import pepseqpred.apps.train_cli as train_cli from pepseqpred.core.io.keys import parse_fullname from pepseqpred.core.preprocess.preparedataset import prepare_dataset from pepseqpred.core.train.split import split_ids_grouped @@ -40,7 +40,7 @@ def __init__(self): def forward(self, batch_tokens, repr_layers, return_contacts=False): _ = return_contacts batch_size, token_len = batch_tokens.shape - rep_dim = 3 # append_seq_len => final emb dim=4 + rep_dim = 3 reps = torch.ones((batch_size, token_len, rep_dim), dtype=torch.float32) return {"representations": {repr_layers[0]: reps}} @@ -328,7 +328,7 @@ def test_prepare_dataset_multisource_pipeline_smoke(monkeypatch, tmp_path: Path) sys, "argv", [ - "train_ffnn_cli.py", + "train_cli.py", "--embedding-dirs", str(emb_dir), "--label-shards", diff --git a/tests/integration/test_train_clis_inprocess.py b/tests/integration/test_train_clis_inprocess.py index f94c04e..5088c83 100644 --- a/tests/integration/test_train_clis_inprocess.py +++ b/tests/integration/test_train_clis_inprocess.py @@ -1,21 +1,23 @@ import sys import json +import csv from pathlib import Path import pytest -import pepseqpred.apps.train_ffnn_cli as train_cli -import pepseqpred.apps.train_ffnn_optuna_cli as optuna_cli +import torch +import pepseqpred.apps.train_cli as train_cli +import pepseqpred.apps.train_optuna_cli as optuna_cli pytestmark = pytest.mark.integration -def test_train_ffnn_cli_main_inprocess(training_artifacts, tmp_path: Path, monkeypatch): +def test_train_cli_main_inprocess(training_artifacts, tmp_path: Path, monkeypatch): save_dir = tmp_path / "train_out" monkeypatch.setattr( sys, "argv", [ - "train_ffnn_cli.py", + "train_cli.py", "--embedding-dirs", str(training_artifacts["embedding_dir"]), "--label-shards", @@ -50,9 +52,18 @@ def test_train_ffnn_cli_main_inprocess(training_artifacts, tmp_path: Path, monke assert (run_dirs[0] / "fully_connected.pt").exists() assert (save_dir / "runs.csv").exists() assert (save_dir / "multi_run_summary.json").exists() - - -def test_train_ffnn_cli_main_inprocess_with_val_curves( + with (save_dir / "runs.csv").open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert rows[0]["SplitStrategy"] == "size-balanced" + assert Path(rows[0]["SplitReportJson"]).exists() + assert "TrainPositiveRate" in rows[0] + assert "ValPositiveRate" in rows[0] + assert rows[0]["ThresholdPolicy"] == "max-recall-min-precision" + assert float(rows[0]["ThresholdMinPrecision"]) == pytest.approx(0.25) + assert "ThresholdPredPosFrac" in rows[0] + + +def test_train_cli_main_inprocess_with_val_curves( training_artifacts, tmp_path: Path, monkeypatch ): save_dir = tmp_path / "train_out_curves" @@ -61,7 +72,7 @@ def test_train_ffnn_cli_main_inprocess_with_val_curves( sys, "argv", [ - "train_ffnn_cli.py", + "train_cli.py", "--embedding-dirs", str(training_artifacts["embedding_dir"]), "--label-shards", @@ -108,14 +119,14 @@ def test_train_ffnn_cli_main_inprocess_with_val_curves( assert pr_plot.exists() -def test_train_ffnn_cli_ensemble_kfold_inprocess(training_artifacts, tmp_path: Path, monkeypatch): +def test_train_cli_ensemble_kfold_inprocess(training_artifacts, tmp_path: Path, monkeypatch): save_dir = tmp_path / "ensemble_out" monkeypatch.setattr( sys, "argv", [ - "train_ffnn_cli.py", + "train_cli.py", "--embedding-dirs", str(training_artifacts["embedding_dir"]), "--label-shards", @@ -163,16 +174,22 @@ def test_train_ffnn_cli_ensemble_kfold_inprocess(training_artifacts, tmp_path: P payload = json.loads(manifest_path.read_text(encoding="utf-8")) assert payload["train_mode"] == "ensemble-kfold" assert payload["n_sets"] == 2 + assert payload["split_strategy"] == "size-balanced" + assert Path(payload["split_report_json"]).exists() + assert payload["threshold_policy"] == "max-recall-min-precision" + assert payload["threshold_min_precision"] == pytest.approx(0.25) assert len(payload["sets"]) == 2 assert [(x["split_seed"], x["train_seed"]) for x in payload["sets"]] == [ (17, 101), (19, 202), ] assert all(int(x["n_members"]) == 2 for x in payload["sets"]) + assert all(x["split_strategy"] == "size-balanced" for x in payload["sets"]) + assert all(x["threshold_policy"] == "max-recall-min-precision" for x in payload["sets"]) assert all(Path(x["manifest_path"]).exists() for x in payload["sets"]) -def test_train_ffnn_cli_ensemble_kfold_with_aggregate_val_curves( +def test_train_cli_ensemble_kfold_with_aggregate_val_curves( training_artifacts, tmp_path: Path, monkeypatch ): save_dir = tmp_path / "ensemble_out_curves" @@ -181,7 +198,7 @@ def test_train_ffnn_cli_ensemble_kfold_with_aggregate_val_curves( sys, "argv", [ - "train_ffnn_cli.py", + "train_cli.py", "--embedding-dirs", str(training_artifacts["embedding_dir"]), "--label-shards", @@ -259,7 +276,7 @@ def test_train_ffnn_cli_ensemble_kfold_with_aggregate_val_curves( @pytest.mark.slow -def test_train_ffnn_optuna_cli_main_inprocess( +def test_train_optuna_cli_main_inprocess( training_artifacts, tmp_path: Path, monkeypatch ): save_dir = tmp_path / "optuna_out" @@ -269,7 +286,7 @@ def test_train_ffnn_optuna_cli_main_inprocess( sys, "argv", [ - "train_ffnn_optuna_cli.py", + "train_optuna_cli.py", "--embedding-dirs", str(training_artifacts["embedding_dir"]), "--label-shards", @@ -309,5 +326,25 @@ def test_train_ffnn_optuna_cli_main_inprocess( optuna_cli.main() - assert (save_dir / "best_trial.json").exists() + best_payload = json.loads((save_dir / "best_trial.json").read_text(encoding="utf-8")) + assert best_payload["split_strategy"] == "size-balanced" + assert Path(best_payload["split_report_json"]).exists() + assert best_payload["threshold_policy"] == "max-recall-min-precision" + assert best_payload["model_head"] == "ffnn" + best_checkpoint = save_dir / "best_model_by_score.pt" + assert best_checkpoint.exists() + checkpoint = torch.load( + best_checkpoint, + map_location="cpu", + weights_only=False, + ) + assert checkpoint["model_config"]["model_head"] == "ffnn" assert csv_path.exists() + with csv_path.open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert rows[0]["ThresholdPolicy"] == "max-recall-min-precision" + assert rows[0]["SplitStrategy"] == "size-balanced" + assert rows[0]["ModelHead"] == "ffnn" + assert "TrainPositiveRate" in rows[0] + assert "ValPositiveRate" in rows[0] + assert "ThresholdPredPosFrac" in rows[0] diff --git a/tests/integration/test_train_ffnn_cli_smoke.py b/tests/integration/test_train_ffnn_cli_smoke.py index 856b304..8ec6ef4 100644 --- a/tests/integration/test_train_ffnn_cli_smoke.py +++ b/tests/integration/test_train_ffnn_cli_smoke.py @@ -3,11 +3,12 @@ import subprocess import sys import pytest +import torch pytestmark = pytest.mark.integration -def test_train_ffnn_cli_smoke(training_artifacts, tmp_path): +def test_train_cli_smoke(training_artifacts, tmp_path): repo_root = Path(__file__).resolve().parents[2] src_path = str(repo_root / "src") env = os.environ.copy() @@ -22,7 +23,7 @@ def test_train_ffnn_cli_smoke(training_artifacts, tmp_path): cmd = [ sys.executable, "-m", - "pepseqpred.apps.train_ffnn_cli", + "pepseqpred.apps.train_cli", "--embedding-dirs", str(training_artifacts["embedding_dir"]), "--label-shards", @@ -58,3 +59,69 @@ def test_train_ffnn_cli_smoke(training_artifacts, tmp_path): assert (run_dirs[0] / "fully_connected.pt").exists() assert (save_dir / "runs.csv").exists() assert (save_dir / "multi_run_summary.json").exists() + + +def test_train_cli_conv1d_smoke(training_artifacts, tmp_path): + repo_root = Path(__file__).resolve().parents[2] + src_path = str(repo_root / "src") + env = os.environ.copy() + current_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = ( + src_path + if not current_pythonpath + else f"{src_path}{os.pathsep}{current_pythonpath}" + ) + + save_dir = tmp_path / "conv_out" + cmd = [ + sys.executable, + "-m", + "pepseqpred.apps.train_cli", + "--embedding-dirs", + str(training_artifacts["embedding_dir"]), + "--label-shards", + str(training_artifacts["label_shard"]), + "--epochs", + "1", + "--batch-size", + "2", + "--num-workers", + "0", + "--hidden-sizes", + "8", + "--dropouts", + "0.1", + "--model-head", + "conv1d", + "--conv-channels", + "3", + "--conv-layers", + "1", + "--conv-kernel-size", + "3", + "--conv-dropout", + "0.0", + "--val-frac", + "0.5", + "--split-seeds", + "11", + "--train-seeds", + "101", + "--save-path", + str(save_dir), + "--results-csv", + str(save_dir / "runs.csv"), + ] + proc = subprocess.run( + cmd, capture_output=True, text=True, cwd=repo_root, env=env + ) + assert proc.returncode == 0, proc.stderr + + run_dirs = list(save_dir.glob("run_*")) + assert run_dirs + checkpoint = torch.load( + run_dirs[0] / "fully_connected.pt", + map_location="cpu", + weights_only=True, + ) + assert checkpoint["model_config"]["model_head"] == "conv1d" diff --git a/tests/integration/test_trainer_ddp_uneven.py b/tests/integration/test_trainer_ddp_uneven.py new file mode 100644 index 0000000..14e8551 --- /dev/null +++ b/tests/integration/test_trainer_ddp_uneven.py @@ -0,0 +1,170 @@ +import logging +import socket +import sys +from datetime import timedelta +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.nn.parallel import DistributedDataParallel as DDP + +from pepseqpred.core.models.ffnn import PepSeqFFNN +from pepseqpred.core.train.trainer import Trainer, TrainerConfig + + +pytestmark = pytest.mark.integration + + +def _rank_batches(rank: int) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + batch_count = 1 if rank == 0 else 3 + batches = [] + for batch_index in range(batch_count): + x = torch.full( + (1, 2, 2), + float((rank + 1) * (batch_index + 1)), + dtype=torch.float32, + ) + y = torch.tensor([[1.0, 0.0]], dtype=torch.float32) + mask = torch.ones((1, 2), dtype=torch.long) + batches.append((x, y, mask)) + return batches + + +def _optimizer_payload(trainer: Trainer) -> list[dict[str, torch.Tensor | int]]: + payload = [] + for group in trainer.optimizer.param_groups: + for parameter in group["params"]: + state = trainer.optimizer.state[parameter] + payload.append({ + "step": int(state["step"].item()), + "exp_avg": state["exp_avg"].detach().cpu().clone(), + "exp_avg_sq": state["exp_avg_sq"].detach().cpu().clone(), + }) + return payload + + +def _uneven_ddp_worker( + rank: int, + world_size: int, + init_method: str, + output_dir: str, +) -> None: + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=30), + ) + try: + torch.manual_seed(1234) + model = PepSeqFFNN( + emb_dim=2, + hidden_sizes=(4,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False, + ) + ddp_model = DDP(model) + trainer = Trainer( + model=ddp_model, + train_loader=_rank_batches(rank), + logger=logging.getLogger(f"uneven_ddp_rank_{rank}"), + val_loader=None, + config=TrainerConfig( + epochs=2, + batch_size=1, + learning_rate=1e-2, + device="cpu", + ), + ) + + epoch_outputs = [] + for epoch in range(2): + epoch_outputs.append(trainer._run_epoch(epoch, train=True)) + result = { + "epochs": epoch_outputs, + "model_state": { + name: tensor.detach().cpu().clone() + for name, tensor in trainer.model.state_dict().items() + }, + "optimizer_state": _optimizer_payload(trainer), + } + torch.save(result, Path(output_dir) / f"rank_{rank}.pt") + dist.barrier() + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + sys.platform == "win32", + reason="PyTorch 2.4 Gloo rendezvous is unavailable in the Windows test environment", +) +def test_uneven_ddp_training_keeps_model_and_adam_state_synchronized( + tmp_path: Path, +) -> None: + world_size = 2 + output_dir = tmp_path / "rank_outputs" + output_dir.mkdir() + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + port = int(sock.getsockname()[1]) + # Use the legacy TCPStore backend for portability across CPU PyTorch builds. + init_method = f"tcp://127.0.0.1:{port}?use_libuv=0" + + mp.spawn( + _uneven_ddp_worker, + args=(world_size, init_method, str(output_dir)), + nprocs=world_size, + join=True, + ) + + rank_results = [ + torch.load( + output_dir / f"rank_{rank}.pt", + map_location="cpu", + weights_only=False, + ) + for rank in range(world_size) + ] + + for rank, result in enumerate(rank_results): + expected_real_batches = 1 if rank == 0 else 3 + expected_dummy_batches = 2 if rank == 0 else 0 + for epoch_output in result["epochs"]: + assert epoch_output["synchronized_steps"] == 3 + assert epoch_output["optimizer_steps"] == 3 + assert epoch_output["zero_valid_steps"] == 0 + assert epoch_output["real_batches"] == expected_real_batches + assert epoch_output["dummy_batches"] == expected_dummy_batches + assert epoch_output["n_residues"] == 8 + + rank_zero = rank_results[0] + rank_one = rank_results[1] + assert rank_zero["model_state"].keys() == rank_one["model_state"].keys() + for name in rank_zero["model_state"]: + torch.testing.assert_close( + rank_zero["model_state"][name], + rank_one["model_state"][name], + rtol=0.0, + atol=0.0, + ) + + assert len(rank_zero["optimizer_state"]) == len(rank_one["optimizer_state"]) + for state_zero, state_one in zip( + rank_zero["optimizer_state"], + rank_one["optimizer_state"], + ): + assert state_zero["step"] == 6 + assert state_one["step"] == 6 + torch.testing.assert_close( + state_zero["exp_avg"], state_one["exp_avg"], rtol=0.0, atol=0.0) + torch.testing.assert_close( + state_zero["exp_avg_sq"], + state_one["exp_avg_sq"], + rtol=0.0, + atol=0.0, + ) diff --git a/tests/unit/api/test_predictor_api.py b/tests/unit/api/test_predictor_api.py index bc22cdd..4fd1de6 100644 --- a/tests/unit/api/test_predictor_api.py +++ b/tests/unit/api/test_predictor_api.py @@ -53,7 +53,7 @@ def to(self, _device): def __call__(self, batch_tokens, repr_layers, return_contacts=False): _ = return_contacts batch_size, token_len = batch_tokens.shape - rep_dim = 3 # append_seq_len -> final emb dim becomes 4 + rep_dim = 3 reps = torch.ones((batch_size, token_len, rep_dim), dtype=torch.float32) return {"representations": {repr_layers[0]: reps}} @@ -61,8 +61,8 @@ def __call__(self, batch_tokens, repr_layers, return_contacts=False): def _write_checkpoint(path: Path, threshold: float = 0.37) -> None: model = PepSeqFFNN( - emb_dim=4, - hidden_sizes=(3,), + emb_dim=3, + hidden_sizes=(4,), dropouts=(0.0,), use_layer_norm=False, use_residual=False, diff --git a/tests/unit/api/test_predictor_api_edges.py b/tests/unit/api/test_predictor_api_edges.py index 229c96f..a09c9fe 100644 --- a/tests/unit/api/test_predictor_api_edges.py +++ b/tests/unit/api/test_predictor_api_edges.py @@ -180,6 +180,23 @@ def _build_model_from_checkpoint(checkpoint, **_kwargs): with pytest.raises(ValueError, match="share emb_dim"): PepSeqPredictor.from_artifact(artifact, model_name="fake_model", device="cpu") + def _build_mixed_seq_len_feature(checkpoint, **_kwargs): + token = checkpoint["token"] + seq_len_feature = "raw" if token == "a" else "inverse" + return ( + object(), + SimpleNamespace(emb_dim=4, seq_len_feature=seq_len_feature), + "state_dict", + ) + + monkeypatch.setattr( + predictor_mod, + "build_model_from_checkpoint", + _build_mixed_seq_len_feature, + ) + with pytest.raises(ValueError, match="seq_len_feature"): + PepSeqPredictor.from_artifact(artifact, model_name="fake_model", device="cpu") + def test_predictor_payload_and_wrapper_branches(monkeypatch): predictor = PepSeqPredictor( diff --git a/tests/unit/apps/test_cli_wrappers.py b/tests/unit/apps/test_cli_wrappers.py index 9dbf627..fecf56d 100644 --- a/tests/unit/apps/test_cli_wrappers.py +++ b/tests/unit/apps/test_cli_wrappers.py @@ -109,6 +109,7 @@ def fake_esm_embeddings_from_fasta(*args, **kwargs): key_delimiter="-", model_name="fake_model", max_tokens=16, + seq_len_feature="inverse", batch_size=2, num_shards=1, shard_id=0 @@ -134,6 +135,7 @@ def fake_esm_embeddings_from_fasta(*args, **kwargs): assert captured["kwargs"]["key_mode"] == "id" assert captured["kwargs"]["id_col"] == "ID" + assert captured["kwargs"]["seq_len_feature"] == "inverse" def test_esm_cli_id_family_requires_metadata(monkeypatch, tmp_path: Path): @@ -154,6 +156,7 @@ def test_esm_cli_id_family_requires_metadata(monkeypatch, tmp_path: Path): key_delimiter="-", model_name="fake_model", max_tokens=16, + seq_len_feature="none", batch_size=2, num_shards=1, shard_id=0 diff --git a/tests/unit/apps/test_train_cli_coverage.py b/tests/unit/apps/test_train_cli_coverage.py index 3872f46..10e1a17 100644 --- a/tests/unit/apps/test_train_cli_coverage.py +++ b/tests/unit/apps/test_train_cli_coverage.py @@ -9,8 +9,8 @@ import pytest import torch -import pepseqpred.apps.train_ffnn_cli as train_cli -import pepseqpred.apps.train_ffnn_optuna_cli as optuna_cli +import pepseqpred.apps.train_cli as train_cli +import pepseqpred.apps.train_optuna_cli as optuna_cli pytestmark = pytest.mark.unit @@ -67,6 +67,50 @@ def _run_main(entrypoint, argv: list[str]) -> None: sys.argv = old_argv +def _capture_protein_dataset_calls(monkeypatch, module): + calls = [] + real_dataset = module.ProteinDataset + + def _wrapped_protein_dataset(*args, **kwargs): + calls.append(dict(kwargs)) + return real_dataset(*args, **kwargs) + + monkeypatch.setattr(module, "ProteinDataset", _wrapped_protein_dataset) + return calls + + +def _assert_train_window_val_full_protein(calls): + train_val_calls = [kwargs for kwargs in calls if "protein_ids" in kwargs] + assert len(train_val_calls) == 2 + + train_kwargs, val_kwargs = train_val_calls + assert train_kwargs["window_size"] == 4 + assert train_kwargs["stride"] == 2 + assert train_kwargs["pad_last_window"] is True + assert val_kwargs["window_size"] is None + assert val_kwargs["stride"] == 1 + assert val_kwargs["pad_last_window"] is False + + +def _capture_pos_weight_resolution(monkeypatch, module, *, counts: tuple[int, int]): + calls = {"count": 0, "global": 0} + + def _compute_pos_neg_counts(loader): + assert loader is not None + calls["count"] += 1 + return counts + + def _global_pos_neg_counts(local_pos, local_neg, ddp): + _ = ddp + calls["global"] += 1 + assert (local_pos, local_neg) == counts + return local_pos, local_neg + + monkeypatch.setattr(module, "compute_pos_neg_counts", _compute_pos_neg_counts) + monkeypatch.setattr(module, "global_pos_neg_counts", _global_pos_neg_counts) + return calls + + def test_train_cli_helper_parsers_and_numeric_summary(): summary_empty = train_cli.summarize_numeric( pd.Series([float("nan"), float("inf"), -float("inf")]) @@ -89,6 +133,18 @@ def test_train_cli_helper_parsers_and_numeric_summary(): train_cli._parse_plot_formats("png,jpg") +def test_hpc_scripts_only_pass_pos_weight_when_env_is_set(): + for script in [ + Path("scripts/hpc/train.sh"), + Path("scripts/hpc/trainoptuna.sh") + ]: + text = script.read_text(encoding="utf-8") + assert "13.18999647945325" not in text + assert 'POS_WEIGHT="${POS_WEIGHT:-}"' in text + assert 'POS_WEIGHT_ARGS+=(--pos-weight "$POS_WEIGHT")' in text + assert '"${POS_WEIGHT_ARGS[@]}"' in text + + @pytest.mark.parametrize( ("legacy_flag", "legacy_value"), [ @@ -97,14 +153,14 @@ def test_train_cli_helper_parsers_and_numeric_summary(): ("--ensemble-train-seeds", "11,12"), ], ) -def test_train_ffnn_cli_rejects_removed_legacy_mode_flags( +def test_train_cli_rejects_removed_legacy_mode_flags( legacy_flag: str, legacy_value: str ): with pytest.raises(ValueError, match="Legacy train-mode flags are no longer supported"): _run_main( train_cli.main, [ - "train_ffnn_cli.py", + "train_cli.py", "--embedding-dirs", "dummy_embedding_dir", "--label-shards", @@ -115,18 +171,24 @@ def test_train_ffnn_cli_rejects_removed_legacy_mode_flags( ) -def test_train_ffnn_cli_real_no_valid_score_with_val_curve_artifacts(): +def test_train_cli_real_no_valid_score_with_val_curve_artifacts(monkeypatch): case_dir = _mk_case_dir("ffnn_no_valid") emb_dir, label_shard = _write_training_artifacts( case_dir, all_uncertain=True ) save_dir = case_dir / "train_out" + dataset_calls = _capture_protein_dataset_calls(monkeypatch, train_cli) + pos_weight_calls = _capture_pos_weight_resolution( + monkeypatch, + train_cli, + counts=(2, 8) + ) try: _run_main( train_cli.main, [ - "train_ffnn_cli.py", + "train_cli.py", "--embedding-dirs", str(emb_dir), "--label-shards", @@ -137,6 +199,10 @@ def test_train_ffnn_cli_real_no_valid_score_with_val_curve_artifacts(): "2", "--num-workers", "0", + "--window-size", + "4", + "--stride", + "2", "--hidden-sizes", "8", "--dropouts", @@ -171,27 +237,51 @@ def test_train_ffnn_cli_real_no_valid_score_with_val_curve_artifacts(): assert int(runs_df.shape[0]) == 1 assert str(runs_df.iloc[0]["BestMetricKey"]) == "f1" assert str(runs_df.iloc[0]["Status"]) == "NO_VALID_SCORE" + assert str(runs_df.iloc[0]["SplitStrategy"]) == "size-balanced" + assert str(runs_df.iloc[0]["SeqLenFeature"]) == "none" + assert Path(str(runs_df.iloc[0]["SplitReportJson"])).exists() + assert "TrainPositiveRate" in runs_df.columns + assert "ValPositiveRate" in runs_df.columns + assert str(runs_df.iloc[0]["ThresholdPolicy"]) == "max-recall-min-precision" + assert str(runs_df.iloc[0]["ThresholdStatus"]) == "no_valid_residues" + assert float(runs_df.iloc[0]["ThresholdMinPrecision"]) == pytest.approx(0.25) + assert float(runs_df.iloc[0]["ThresholdMinRecall"]) == pytest.approx(0.80) + assert float(runs_df.iloc[0]["ThresholdFixedValue"]) == pytest.approx(0.50) summary = json.loads( (save_dir / "multi_run_summary.json").read_text(encoding="utf-8") ) assert int(summary["n_runs"]) == 1 + assert summary["split_strategy"] == "size-balanced" + assert summary["seq_len_feature"] == "none" + assert Path(summary["split_report_json"]).exists() + assert summary["threshold_policy"] == "max-recall-min-precision" + assert summary["threshold_min_precision"] == pytest.approx(0.25) + assert "ThresholdPredPosFrac" in summary["metrics"] + _assert_train_window_val_full_protein(dataset_calls) + assert pos_weight_calls == {"count": 1, "global": 1} finally: shutil.rmtree(case_dir, ignore_errors=True) -def test_train_ffnn_cli_real_ensemble_manifest_generation(): +def test_train_cli_real_ensemble_manifest_generation(monkeypatch): case_dir = _mk_case_dir("ffnn_ensemble") emb_dir, label_shard = _write_training_artifacts( case_dir, all_uncertain=False ) save_dir = case_dir / "ensemble_out" + def _unexpected_auto_pos_weight(*_args, **_kwargs): + raise AssertionError("explicit --pos-weight should bypass automatic counting") + + monkeypatch.setattr(train_cli, "compute_pos_neg_counts", _unexpected_auto_pos_weight) + monkeypatch.setattr(train_cli, "global_pos_neg_counts", _unexpected_auto_pos_weight) + try: _run_main( train_cli.main, [ - "train_ffnn_cli.py", + "train_cli.py", "--embedding-dirs", str(emb_dir), "--label-shards", @@ -214,6 +304,18 @@ def test_train_ffnn_cli_real_ensemble_manifest_generation(): "17,19", "--train-seeds", "101,202", + "--pos-weight", + "3.5", + "--split-strategy", + "label-stratified", + "--threshold-policy", + "fixed", + "--threshold-min-precision", + "0.33", + "--threshold-min-recall", + "0.77", + "--threshold-fixed-value", + "0.49", "--save-path", str(save_dir), "--results-csv", @@ -228,13 +330,37 @@ def test_train_ffnn_cli_real_ensemble_manifest_generation(): ) assert payload["train_mode"] == "ensemble-kfold" assert payload["n_sets"] == 2 + assert payload["split_strategy"] == "label-stratified" + assert payload["seq_len_feature"] == "none" + assert Path(payload["split_report_json"]).exists() + assert payload["threshold_policy"] == "fixed" + assert payload["threshold_min_precision"] == pytest.approx(0.33) + assert payload["threshold_min_recall"] == pytest.approx(0.77) + assert payload["threshold_fixed_value"] == pytest.approx(0.49) assert len(payload["sets"]) == 2 assert all(int(x["n_members"]) == 2 for x in payload["sets"]) + assert all(x["threshold_policy"] == "fixed" for x in payload["sets"]) + set_manifest_path = Path(payload["sets"][0]["manifest_path"]) + set_payload = json.loads(set_manifest_path.read_text(encoding="utf-8")) + assert set_payload["threshold_policy"] == "fixed" + assert set_payload["split_strategy"] == "label-stratified" + assert set_payload["seq_len_feature"] == "none" + assert all( + member["threshold_policy"] == "fixed" + for member in set_payload["members"] + ) + runs_df = pd.read_csv(save_dir / "runs.csv") + assert set(runs_df["SplitStrategy"]) == {"label-stratified"} + assert "TrainPositiveRate" in runs_df.columns + assert "ValPositiveRate" in runs_df.columns + assert set(runs_df["ThresholdPolicy"]) == {"fixed"} + assert set(runs_df["ThresholdStatus"]) == {"ok"} + assert set(runs_df["ThresholdFixedValue"]) == {0.49} finally: shutil.rmtree(case_dir, ignore_errors=True) -def test_train_ffnn_optuna_cli_real_with_storage_and_helpers(): +def test_train_optuna_cli_real_with_storage_and_helpers(monkeypatch): assert optuna_cli._broadcast_params({"a": 1}, None) == {"a": 1} study = optuna.create_study(sampler=optuna.samplers.RandomSampler(seed=7)) @@ -286,12 +412,18 @@ def test_train_ffnn_optuna_cli_real_with_storage_and_helpers(): save_dir = case_dir / "optuna_out" csv_path = save_dir / "trials.csv" storage_uri = f"sqlite:///{(case_dir / 'study.db').as_posix()}" + dataset_calls = _capture_protein_dataset_calls(monkeypatch, optuna_cli) + pos_weight_calls = _capture_pos_weight_resolution( + monkeypatch, + optuna_cli, + counts=(3, 12) + ) try: _run_main( optuna_cli.main, [ - "train_ffnn_optuna_cli.py", + "train_optuna_cli.py", "--embedding-dirs", str(emb_dir), "--label-shards", @@ -310,6 +442,10 @@ def test_train_ffnn_optuna_cli_real_with_storage_and_helpers(): "2", "--num-workers", "0", + "--window-size", + "4", + "--stride", + "2", "--metric", "auc", "--arch-mode", @@ -328,6 +464,16 @@ def test_train_ffnn_optuna_cli_real_with_storage_and_helpers(): str(csv_path), "--study-name", "unit_realcov_study", + "--split-strategy", + "label-stratified", + "--threshold-policy", + "fixed", + "--threshold-min-precision", + "0.35", + "--threshold-min-recall", + "0.75", + "--threshold-fixed-value", + "0.49", ], ) @@ -336,7 +482,24 @@ def test_train_ffnn_optuna_cli_real_with_storage_and_helpers(): ) assert best_payload["study_name"] == "unit_realcov_study" assert best_payload["metric"] == "auc" + assert best_payload["seq_len_feature"] == "none" + assert best_payload["split_strategy"] == "label-stratified" + assert Path(best_payload["split_report_json"]).exists() + assert best_payload["threshold_policy"] == "fixed" + assert best_payload["threshold_min_precision"] == pytest.approx(0.35) + assert best_payload["threshold_min_recall"] == pytest.approx(0.75) + assert best_payload["threshold_fixed_value"] == pytest.approx(0.49) assert csv_path.exists() + trials_df = pd.read_csv(csv_path) + assert str(trials_df.iloc[0]["SplitStrategy"]) == "label-stratified" + assert str(trials_df.iloc[0]["SeqLenFeature"]) == "none" + assert "TrainPositiveRate" in trials_df.columns + assert "ValPositiveRate" in trials_df.columns + assert str(trials_df.iloc[0]["ThresholdPolicy"]) == "fixed" + assert float(trials_df.iloc[0]["ThresholdFixedValue"]) == pytest.approx(0.49) + assert "ThresholdPredPosFrac" in trials_df.columns assert (case_dir / "study.db").exists() + _assert_train_window_val_full_protein(dataset_calls) + assert pos_weight_calls == {"count": 1, "global": 1} finally: shutil.rmtree(case_dir, ignore_errors=True) diff --git a/tests/unit/apps/test_train_ffnn_run_plans.py b/tests/unit/apps/test_train_ffnn_run_plans.py index 692cb81..f2201bb 100644 --- a/tests/unit/apps/test_train_ffnn_run_plans.py +++ b/tests/unit/apps/test_train_ffnn_run_plans.py @@ -1,6 +1,6 @@ import argparse import pytest -from pepseqpred.apps.train_ffnn_cli import _build_run_plans +from pepseqpred.apps.train_cli import _build_run_plans pytestmark = pytest.mark.unit @@ -12,6 +12,7 @@ def _args(**overrides): train_seeds=None, seed=42, split_type="id-family", + split_strategy="size-balanced", val_frac=0.5, ) for key, value in overrides.items(): @@ -19,6 +20,17 @@ def _args(**overrides): return base +def _label_support(ids, family_groups): + return { + protein_id: { + "valid_residues": 10, + "positive_residues": 10 if family_groups[protein_id] in {"A", "C"} else 0, + "negative_residues": 0 if family_groups[protein_id] in {"A", "C"} else 10, + } + for protein_id in ids + } + + def test_build_run_plans_kfold_set_pairs_map_to_full_kfold_sets(): args = _args( n_folds=2, @@ -39,6 +51,7 @@ def test_build_run_plans_kfold_set_pairs_map_to_full_kfold_sets(): assert meta["train_seeds"] == [11, 22] assert meta["train_mode"] == "ensemble-kfold" assert meta["ensemble_seed_mode"] == "set-paired" + assert meta["split_strategy"] == "size-balanced" set_1 = [plan for plan in plans if plan.ensemble_set_index == 1] assert len(set_1) == 2 @@ -103,6 +116,7 @@ def test_build_run_plans_single_fold_keeps_holdout_behavior(): assert meta["n_folds"] == 1 assert meta["n_sets"] == 2 assert meta["train_mode"] == "seeded" + assert meta["split_strategy"] == "size-balanced" def test_build_run_plans_requires_n_folds_at_least_one(): @@ -111,3 +125,69 @@ def test_build_run_plans_requires_n_folds_at_least_one(): with pytest.raises(ValueError, match="--n-folds must be >= 1"): _build_run_plans(args, ids, {}) + + +def test_build_run_plans_label_stratified_single_fold_uses_support(): + args = _args( + n_folds=1, + split_seeds="7", + train_seeds="77", + split_type="id-family", + split_strategy="label-stratified", + ) + ids = ["a1", "a2", "b1", "b2", "c1", "c2", "d1", "d2"] + family_groups = {protein_id: protein_id[0].upper() for protein_id in ids} + support = _label_support(ids, family_groups) + + plans, meta = _build_run_plans( + args, + ids, + split_groups=family_groups, + family_groups=family_groups, + label_support_by_id=support, + ) + + assert len(plans) == 1 + plan = plans[0] + train_families = {family_groups[i] for i in plan.train_ids_all} + val_families = {family_groups[i] for i in plan.val_ids_all} + assert train_families.isdisjoint(val_families) + val_pos = sum(support[i]["positive_residues"] for i in plan.val_ids_all) + val_valid = sum(support[i]["valid_residues"] for i in plan.val_ids_all) + assert val_valid > 0 + assert (val_pos / val_valid) == pytest.approx(0.5) + assert meta["split_strategy"] == "label-stratified" + assert meta["n_folds"] == 1 + + +def test_build_run_plans_label_stratified_kfold_uses_support(): + args = _args( + n_folds=2, + split_seeds="7", + train_seeds="77", + split_type="id-family", + split_strategy="label-stratified", + ) + ids = ["a1", "a2", "b1", "b2", "c1", "c2", "d1", "d2"] + family_groups = {protein_id: protein_id[0].upper() for protein_id in ids} + support = _label_support(ids, family_groups) + + plans, meta = _build_run_plans( + args, + ids, + split_groups=family_groups, + family_groups=family_groups, + label_support_by_id=support, + ) + + assert len(plans) == 2 + assert meta["split_strategy"] == "label-stratified" + assert meta["n_folds"] == 2 + for plan in plans: + train_families = {family_groups[i] for i in plan.train_ids_all} + val_families = {family_groups[i] for i in plan.val_ids_all} + assert train_families.isdisjoint(val_families) + val_pos = sum(support[i]["positive_residues"] for i in plan.val_ids_all) + val_valid = sum(support[i]["valid_residues"] for i in plan.val_ids_all) + assert val_valid > 0 + assert (val_pos / val_valid) == pytest.approx(0.5) diff --git a/tests/unit/core/embeddings/test_esm2.py b/tests/unit/core/embeddings/test_esm2.py index 82a8c42..6714b37 100644 --- a/tests/unit/core/embeddings/test_esm2.py +++ b/tests/unit/core/embeddings/test_esm2.py @@ -44,7 +44,7 @@ def forward(self, batch_tokens, repr_layers, return_contacts=False): return {"representations": {repr_layers[0]: rep}} -def test_clean_seq_token_batches_and_append_len(): +def test_clean_seq_token_batches_and_seq_len_features(): assert esm2.clean_seq("acdx-*\n") == "ACDX" batches = list( @@ -58,10 +58,21 @@ def test_clean_seq_token_batches_and_append_len(): assert [x[0] for x in batches[0]] == ["a", "b"] arr = torch.ones((5, 3), dtype=torch.float32).numpy() + out_none = esm2.apply_seq_len_feature(arr, 5, "none") + assert out_none.shape == (5, 3) + out = esm2.append_seq_len(arr, 5) assert out.shape == (5, 4) assert (out[:, -1] == 5).all() + out_inverse = esm2.apply_seq_len_feature(arr, 5, "inverse") + assert out_inverse.shape == (5, 4) + assert out_inverse[0, -1] == pytest.approx(0.2) + assert (out_inverse[:, -1] == out_inverse[0, -1]).all() + + with pytest.raises(ValueError, match="seq_len_feature"): + esm2.apply_seq_len_feature(arr, 5, "bad") + def test_compute_window_embedding_cpu_paths(): token = torch.tensor([[0, 1, 1, 1, 1, 1, 2]], dtype=torch.long) @@ -115,8 +126,48 @@ def test_esm_embeddings_from_fasta_short_and_long(monkeypatch, tmp_path: Path): assert failed == [] assert len(index_df) == 2 assert set(index_df["handle"]) == {"short", "long"} + assert set(index_df["seq_len_feature"]) == {"none"} for key in index_df["id"].tolist(): assert (per_seq / f"{key}.pt").exists() + emb = torch.load(per_seq / f"{key}.pt", + map_location="cpu", weights_only=True) + assert emb.shape[1] == 3 + + +def test_esm_embeddings_from_fasta_raw_seq_len_feature(monkeypatch, tmp_path: Path): + fake_pretrained = types.SimpleNamespace( + fake_model=lambda: (FakeModel(embed_dim=3), FakeAlphabet()) + ) + monkeypatch.setattr(esm2.esm, "pretrained", fake_pretrained) + + df = pd.DataFrame( + [{"ID": "P1", "Sequence": "ACD", "viral_family": "111"}] + ) + + per_seq = tmp_path / "pts" + idx_csv = tmp_path / "idx.csv" + index_df, failed = esm2.esm_embeddings_from_fasta( + df, + id_col="ID", + seq_col="Sequence", + family_col="viral_family", + model_name="fake_model", + max_tokens=6, + batch_size=1, + per_seq_dir=per_seq, + index_csv_path=idx_csv, + key_mode="id-family", + key_delimiter="-", + seq_len_feature="raw", + logger=logging.getLogger("esm2_test") + ) + + assert failed == [] + assert index_df.loc[0, "seq_len_feature"] == "raw" + emb = torch.load(per_seq / f"{index_df.loc[0, 'id']}.pt", + map_location="cpu", weights_only=True) + assert emb.shape == (3, 4) + assert (emb[:, -1] == 3).all() def test_esm_embeddings_from_fasta_key_validation(monkeypatch, tmp_path: Path): diff --git a/tests/unit/core/models/test_model_factory.py b/tests/unit/core/models/test_model_factory.py new file mode 100644 index 0000000..6f015a2 --- /dev/null +++ b/tests/unit/core/models/test_model_factory.py @@ -0,0 +1,107 @@ +import pytest +import torch + +from pepseqpred.core.models.factory import ( + PepSeqModelConfig, + build_pepseq_model, + model_config_from_mapping, + model_config_to_dict, +) +from pepseqpred.core.models.ffnn import PepSeqConvFFNN, PepSeqFFNN + +pytestmark = pytest.mark.unit + + +def test_conv_head_preserves_residue_logit_shape(): + model = PepSeqConvFFNN( + emb_dim=4, + hidden_sizes=(8,), + dropouts=(0.0,), + conv_channels=3, + conv_layers=2, + conv_kernel_size=5, + conv_dropout=0.0, + ) + + out = model(torch.randn(2, 7, 4)) + + assert out.shape == (2, 7) + + +def test_conv_head_rejects_even_kernel_size(): + with pytest.raises(ValueError, match="positive odd"): + PepSeqConvFFNN( + emb_dim=4, + hidden_sizes=(8,), + dropouts=(0.0,), + conv_kernel_size=4, + ) + + +def test_model_factory_builds_ffnn_and_conv_heads(): + base = PepSeqModelConfig( + emb_dim=4, + hidden_sizes=(8,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False, + ) + assert isinstance(build_pepseq_model(base), PepSeqFFNN) + + conv = PepSeqModelConfig( + emb_dim=4, + hidden_sizes=(8,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False, + model_head="conv1d", + conv_channels=3, + conv_layers=1, + conv_kernel_size=3, + conv_dropout=0.0, + ) + assert isinstance(build_pepseq_model(conv), PepSeqConvFFNN) + + payload = model_config_to_dict(conv) + assert payload["hidden_sizes"] == [8] + assert payload["dropouts"] == [0.0] + assert "seq_len_feature" not in payload + assert model_config_from_mapping(payload) == conv + + +def test_model_config_seq_len_feature_serialization(): + cfg = PepSeqModelConfig( + emb_dim=4, + hidden_sizes=(8,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False, + seq_len_feature="raw", + ) + payload = model_config_to_dict(cfg) + assert payload["seq_len_feature"] == "raw" + assert model_config_from_mapping(payload) == cfg + + inverse = PepSeqModelConfig( + emb_dim=4, + hidden_sizes=(8,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False, + seq_len_feature="inverse", + ) + assert model_config_to_dict(inverse)["seq_len_feature"] == "inverse" + + omitted = dict(payload) + omitted.pop("seq_len_feature") + assert model_config_from_mapping(omitted).seq_len_feature is None + + for bad in ("none", None, "bad"): + bad_payload = dict(payload) + bad_payload["seq_len_feature"] = bad + with pytest.raises(ValueError, match="seq_len_feature"): + model_config_from_mapping(bad_payload) diff --git a/tests/unit/core/predict/test_inference.py b/tests/unit/core/predict/test_inference.py index 7a906ef..d49f21b 100644 --- a/tests/unit/core/predict/test_inference.py +++ b/tests/unit/core/predict/test_inference.py @@ -1,6 +1,11 @@ import pytest import torch from pepseqpred.core.models.ffnn import PepSeqFFNN +from pepseqpred.core.models.factory import ( + PepSeqModelConfig, + build_pepseq_model, + model_config_to_dict, +) from pepseqpred.core.predict.inference import ( FFNNModelConfig, build_model_from_checkpoint, @@ -32,6 +37,28 @@ def _make_checkpoint( return {"model_state_dict": model.state_dict(), "metrics": {"threshold": 0.37}} +def _make_conv_checkpoint(): + cfg = PepSeqModelConfig( + emb_dim=4, + hidden_sizes=(8,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False, + model_head="conv1d", + conv_channels=3, + conv_layers=1, + conv_kernel_size=3, + conv_dropout=0.0, + ) + model = build_pepseq_model(cfg) + return { + "model_state_dict": model.state_dict(), + "model_config": model_config_to_dict(cfg), + "metrics": {"threshold": 0.37}, + } + + def test_normalize_state_dict_keys_strips_module_prefix(): ckpt = _make_checkpoint() state = ckpt["model_state_dict"] @@ -74,6 +101,19 @@ def test_build_model_from_checkpoint_handles_ddp_prefixed_state(): assert y.shape == (2, 5) +def test_build_model_from_checkpoint_uses_checkpoint_model_config_for_conv(): + ckpt = _make_conv_checkpoint() + + model, cfg, cfg_src = build_model_from_checkpoint(ckpt, device="cpu") + + assert cfg_src == "checkpoint" + assert cfg.model_head == "conv1d" + assert cfg.conv_channels == 3 + with torch.inference_mode(): + y = model(torch.randn(2, 5, 4)) + assert y.shape == (2, 5) + + def test_build_model_from_checkpoint_rejects_invalid_num_classes(): ckpt = _make_checkpoint() bad_cfg = FFNNModelConfig( @@ -125,6 +165,8 @@ def test_predict_ensemble_from_embedding_uses_majority_vote(): assert out["binary_mask"] == "1110" assert out["n_members"] == 3 assert out["votes_needed"] == 2 + assert out["ensemble_aggregation"] == "majority" + assert out["ensemble_threshold"] is None def test_predict_ensemble_from_embedding_even_tie_is_negative(): @@ -143,6 +185,26 @@ def test_predict_ensemble_from_embedding_even_tie_is_negative(): assert out["votes_needed"] == 2 +def test_predict_ensemble_from_embedding_mean_prob_thresholding(): + emb = torch.zeros((2, 3), dtype=torch.float32) + models = [ + _ConstantLogitModel([2.0, -2.0]), + _ConstantLogitModel([-2.0, 2.0]), + ] + out = predict_ensemble_from_embedding( + psp_models=models, + protein_emb=emb, + device="cpu", + thresholds=[0.5, 0.5], + aggregation="mean-prob", + ensemble_threshold=0.49 + ) + assert out["binary_mask"] == "11" + assert out["votes_needed"] is None + assert out["ensemble_aggregation"] == "mean-prob" + assert out["ensemble_threshold"] == pytest.approx(0.49) + + def test_predict_ensemble_from_embedding_k1_matches_single(): emb = torch.zeros((3, 3), dtype=torch.float32) model = _ConstantLogitModel([2.0, -2.0, 2.0]) diff --git a/tests/unit/core/predict/test_inference_edges.py b/tests/unit/core/predict/test_inference_edges.py index 9c98983..fc6e85f 100644 --- a/tests/unit/core/predict/test_inference_edges.py +++ b/tests/unit/core/predict/test_inference_edges.py @@ -119,6 +119,21 @@ def test_infer_decision_threshold_edge_cases(): assert inf.infer_decision_threshold({"metrics": {"threshold": 1.1}}, default=0.5) == pytest.approx(0.5) +def test_resolve_prediction_seq_len_feature(): + cfg_none = types.SimpleNamespace(emb_dim=3) + cfg_raw = types.SimpleNamespace(emb_dim=4, seq_len_feature="raw") + cfg_inverse = types.SimpleNamespace(emb_dim=4, seq_len_feature="inverse") + + assert inf.resolve_prediction_seq_len_feature("auto", cfg_none) == "none" + assert inf.resolve_prediction_seq_len_feature("auto", cfg_raw) == "raw" + assert inf.resolve_prediction_seq_len_feature("auto", cfg_inverse) == "inverse" + assert inf.resolve_prediction_seq_len_feature("none", cfg_raw) == "none" + assert inf.resolve_prediction_seq_len_feature("raw", cfg_none) == "raw" + assert inf.resolve_prediction_seq_len_feature("inverse", cfg_none) == "inverse" + with pytest.raises(ValueError, match="seq_len_feature"): + inf.resolve_prediction_seq_len_feature("bad", cfg_none) + + def test_prediction_payload_and_member_probability_errors(): with pytest.raises(ValueError, match="Expected probs shape"): inf._build_prediction_payload( @@ -184,7 +199,31 @@ def test_embed_protein_seq_paths(monkeypatch): device="cpu", max_tokens=10 ) - assert out_short.shape == (5, 4) + assert out_short.shape == (5, 3) + + out_raw = inf.embed_protein_seq( + protein_seq="ABCDE", + esm_model=_FakeEsmModel(), + layer=1, + batch_converter=_FakeBatchConverter(), + device="cpu", + max_tokens=10, + seq_len_feature="raw" + ) + assert out_raw.shape == (5, 4) + assert (out_raw[:, -1] == 5).all() + + out_inverse = inf.embed_protein_seq( + protein_seq="ABCDE", + esm_model=_FakeEsmModel(), + layer=1, + batch_converter=_FakeBatchConverter(), + device="cpu", + max_tokens=10, + seq_len_feature="inverse" + ) + assert out_inverse.shape == (5, 4) + assert out_inverse[0, -1] == pytest.approx(0.2) monkeypatch.setattr( inf, @@ -199,7 +238,7 @@ def test_embed_protein_seq_paths(monkeypatch): device="cpu", max_tokens=1 ) - assert out_long.shape == (5, 4) + assert out_long.shape == (5, 3) class _BadShapeModel(torch.nn.Module): @@ -279,6 +318,23 @@ def test_predict_ensemble_and_predict_protein_edge_cases(monkeypatch): device="cpu", thresholds=[1.0] ) + with pytest.raises(ValueError, match="aggregation must be one of"): + inf.predict_ensemble_from_embedding( + psp_models=[_GoodShapeModel()], + protein_emb=torch.zeros((3, 3), dtype=torch.float32), + device="cpu", + thresholds=[0.5], + aggregation="bad" + ) + with pytest.raises(ValueError, match="ensemble_threshold"): + inf.predict_ensemble_from_embedding( + psp_models=[_GoodShapeModel(), _GoodShapeModel()], + protein_emb=torch.zeros((3, 3), dtype=torch.float32), + device="cpu", + thresholds=[0.5, 0.5], + aggregation="mean-prob", + ensemble_threshold=1.0 + ) model_a = _GoodShapeModel() model_b = _GoodShapeModel() diff --git a/tests/unit/core/train/test_split.py b/tests/unit/core/train/test_split.py index 0eb91c9..5623dc6 100644 --- a/tests/unit/core/train/test_split.py +++ b/tests/unit/core/train/test_split.py @@ -1,5 +1,10 @@ import pytest +import torch from pepseqpred.core.train.split import ( + build_label_support_by_id, + build_label_stratified_kfold_splits, + count_label_tensor_support, + split_ids_label_stratified, split_ids_grouped, partition_ids_weighted, build_kfold_splits, @@ -20,6 +25,143 @@ def test_split_ids_grouped_keeps_groups_intact(): assert train_groups.isdisjoint(val_groups) +def test_label_support_counts_binary_tristate_and_missing(tmp_path): + shard_path = tmp_path / "labels.pt" + torch.save( + { + "labels": { + "binary": torch.tensor([1, 0, 0, 1], dtype=torch.float32), + "tristate": torch.tensor( + [ + [1, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], + dtype=torch.uint8, + ), + "all_uncertain": torch.tensor( + [[0, 1, 0], [0, 1, 0]], + dtype=torch.uint8, + ), + } + }, + shard_path, + ) + + assert count_label_tensor_support( + torch.tensor([1, 0, 0], dtype=torch.float32) + ) == { + "valid_residues": 3, + "positive_residues": 1, + "negative_residues": 2, + } + + support = build_label_support_by_id( + ["binary", "tristate", "all_uncertain", "missing"], + { + "binary": shard_path, + "tristate": shard_path, + "all_uncertain": shard_path, + }, + ) + + assert support["binary"]["status"] == "ok" + assert support["binary"]["valid_residues"] == 4 + assert support["binary"]["positive_residues"] == 2 + assert support["binary"]["negative_residues"] == 2 + assert support["tristate"]["valid_residues"] == 2 + assert support["tristate"]["positive_residues"] == 1 + assert support["tristate"]["negative_residues"] == 1 + assert support["all_uncertain"]["valid_residues"] == 0 + assert support["missing"]["status"] == "missing_label" + + +def test_label_stratified_holdout_keeps_groups_and_balances_positive_rate(): + ids = ["a1", "a2", "b1", "b2", "c1", "c2", "d1", "d2"] + groups = {protein_id: protein_id[0].upper() for protein_id in ids} + support = { + protein_id: { + "valid_residues": 10, + "positive_residues": 10 if groups[protein_id] in {"A", "C"} else 0, + "negative_residues": 0 if groups[protein_id] in {"A", "C"} else 10, + "label_shard": "labels.pt", + "status": "ok", + } + for protein_id in ids + } + + size_train, size_val = split_ids_grouped( + ids, val_frac=0.5, seed=0, groups=groups) + strat_train, strat_val = split_ids_label_stratified( + ids, val_frac=0.5, seed=0, groups=groups, support_by_id=support) + strat_train_2, strat_val_2 = split_ids_label_stratified( + ids, val_frac=0.5, seed=0, groups=groups, support_by_id=support) + + def _pos_rate(split_ids): + pos = sum(support[protein_id]["positive_residues"] for protein_id in split_ids) + valid = sum(support[protein_id]["valid_residues"] for protein_id in split_ids) + return pos / valid + + assert (strat_train, strat_val) == (strat_train_2, strat_val_2) + assert {groups[i] for i in strat_train}.isdisjoint({groups[i] for i in strat_val}) + assert abs(_pos_rate(strat_val) - 0.5) < abs(_pos_rate(size_val) - 0.5) + assert sorted(strat_train + strat_val) == sorted(ids) + + +def test_label_stratified_kfold_keeps_groups_and_is_deterministic(): + ids = ["a1", "a2", "b1", "b2", "c1", "c2", "d1", "d2"] + groups = {protein_id: protein_id[0].upper() for protein_id in ids} + support = { + protein_id: { + "valid_residues": 10, + "positive_residues": 10 if groups[protein_id] in {"A", "C"} else 0, + "negative_residues": 0 if groups[protein_id] in {"A", "C"} else 10, + "label_shard": "labels.pt", + "status": "ok", + } + for protein_id in ids + } + + splits = build_label_stratified_kfold_splits( + ids, n_folds=2, seed=5, groups=groups, support_by_id=support) + splits_2 = build_label_stratified_kfold_splits( + ids, n_folds=2, seed=5, groups=groups, support_by_id=support) + + assert splits == splits_2 + assert len(splits) == 2 + for train_ids, val_ids in splits: + assert {groups[i] for i in train_ids}.isdisjoint({groups[i] for i in val_ids}) + assert sorted(train_ids + val_ids) == sorted(ids) + pos = sum(support[protein_id]["positive_residues"] for protein_id in val_ids) + valid = sum(support[protein_id]["valid_residues"] for protein_id in val_ids) + assert pos / valid == pytest.approx(0.5) + + +def test_label_stratified_split_edge_cases(): + ids = ["a1", "b1"] + groups = {"a1": "A", "b1": "B"} + support = { + protein_id: { + "valid_residues": 0, + "positive_residues": 0, + "negative_residues": 0, + "label_shard": "labels.pt", + "status": "ok", + } + for protein_id in ids + } + + assert split_ids_label_stratified( + ids, val_frac=0.0, seed=1, groups=groups, support_by_id=support + ) == (ids, []) + assert split_ids_label_stratified( + ids, val_frac=1.0, seed=1, groups=groups, support_by_id=support + ) == ([], ids) + with pytest.raises(ValueError, match="cannot exceed number of groups"): + build_label_stratified_kfold_splits( + ids, n_folds=3, seed=1, groups=groups, support_by_id=support) + + def test_partition_ids_weighted_non_empty(): ids = ["p1", "p2", "p3", "p4"] weights = {"p1": 100.0, "p2": 90.0, "p3": 10.0, "p4": 9.0} diff --git a/tests/unit/core/train/test_threshold.py b/tests/unit/core/train/test_threshold.py index 3885e11..85e5d3f 100644 --- a/tests/unit/core/train/test_threshold.py +++ b/tests/unit/core/train/test_threshold.py @@ -1,6 +1,11 @@ import numpy as np import pytest -from pepseqpred.core.train.threshold import find_threshold_max_recall_min_precision +from pepseqpred.core.train.threshold import ( + DEFAULT_THRESHOLD_GRID, + find_threshold_max_recall_min_precision, + select_threshold, + threshold_diagnostic_grid, +) pytestmark = pytest.mark.unit @@ -21,3 +26,73 @@ def test_threshold_respects_min_precision(): y_true, y_prob, min_precision=0.5) assert out["status"] == "ok" assert out["precision"] >= 0.5 + + +def test_threshold_min_precision_unreachable_falls_back_to_best_precision(): + y_true = np.array([1, 0, 0], dtype=np.int64) + y_prob = np.array([0.1, 0.9, 0.8], dtype=np.float64) + out = select_threshold( + y_true, + y_prob, + policy="max-recall-min-precision", + min_precision=0.8, + ) + assert out["status"] == "min_precision_unreachable" + assert out["threshold"] == pytest.approx(0.1) + assert out["precision"] == pytest.approx(1.0 / 3.0) + + +def test_threshold_best_f1_and_mcc_pick_less_conservative_operating_point(): + y_true = np.array([1, 1, 0, 0], dtype=np.int64) + y_prob = np.array([0.9, 0.4, 0.8, 0.1], dtype=np.float64) + f1_out = select_threshold(y_true, y_prob, policy="best-f1") + mcc_out = select_threshold(y_true, y_prob, policy="best-mcc") + assert f1_out["status"] == "ok" + assert f1_out["threshold"] == pytest.approx(0.4) + assert f1_out["f1"] == pytest.approx(0.8) + assert mcc_out["threshold"] == pytest.approx(0.4) + + +def test_threshold_min_recall_max_precision(): + y_true = np.array([1, 1, 0, 0], dtype=np.int64) + y_prob = np.array([0.9, 0.3, 0.8, 0.1], dtype=np.float64) + out = select_threshold( + y_true, + y_prob, + policy="min-recall-max-precision", + min_recall=1.0, + ) + assert out["status"] == "ok" + assert out["threshold"] == pytest.approx(0.3) + assert out["recall"] == pytest.approx(1.0) + assert out["precision"] == pytest.approx(2.0 / 3.0) + + +def test_threshold_fixed_policy_and_grid(): + y_true = np.array([1, 0, 1, 0], dtype=np.int64) + y_prob = np.array([0.7, 0.6, 0.2, 0.1], dtype=np.float64) + out = select_threshold( + y_true, + y_prob, + policy="fixed", + fixed_threshold=0.5, + ) + assert out["status"] == "ok" + assert out["policy"] == "fixed" + assert out["threshold"] == pytest.approx(0.5) + assert out["tp"] == 1 + assert out["fp"] == 1 + grid = threshold_diagnostic_grid(y_true, y_prob) + assert len(grid) == len(DEFAULT_THRESHOLD_GRID) + assert {row["threshold"] for row in grid} == set(DEFAULT_THRESHOLD_GRID) + + +def test_threshold_invalid_inputs_raise(): + y_true = np.array([1, 0], dtype=np.int64) + y_prob = np.array([0.9, 0.1], dtype=np.float64) + with pytest.raises(ValueError, match="Unsupported threshold policy"): + select_threshold(y_true, y_prob, policy="bogus") + with pytest.raises(ValueError, match="fixed_threshold"): + select_threshold(y_true, y_prob, policy="fixed", fixed_threshold=1.0) + with pytest.raises(ValueError, match="length mismatch"): + select_threshold(y_true, np.array([0.5], dtype=np.float64)) diff --git a/tests/unit/core/train/test_trainer_batch_step.py b/tests/unit/core/train/test_trainer_batch_step.py index 94adbea..590177b 100644 --- a/tests/unit/core/train/test_trainer_batch_step.py +++ b/tests/unit/core/train/test_trainer_batch_step.py @@ -2,6 +2,7 @@ import pytest import torch import torch.nn as nn +import pepseqpred.core.train.trainer as trainer_mod from pepseqpred.core.models.ffnn import PepSeqFFNN from pepseqpred.core.train.trainer import Trainer, TrainerConfig @@ -18,6 +19,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.zeros((x.size(0), x.size(1) + 1), device=x.device) + (self.p * 0.0) +class ConstantLogitModel(nn.Module): + def __init__(self): + super().__init__() + self.p = nn.Parameter(torch.tensor(0.0)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.zeros((x.size(0), x.size(1)), device=x.device) + self.p + + def _make_trainer(model: nn.Module) -> Trainer: return Trainer( model=model, @@ -28,7 +38,42 @@ def _make_trainer(model: nn.Module) -> Trainer: ) -def test_batch_step_zero_mask_returns_zero_n(): +def test_batch_step_zero_mask_returns_zero_n(monkeypatch): + model = PepSeqFFNN( + emb_dim=4, + hidden_sizes=(8,), + dropouts=(0.0,), + num_classes=1, + use_layer_norm=False, + use_residual=False + ) + trainer = _make_trainer(model) + + x = torch.randn(1, 3, 4) + y = torch.tensor([[1.0, 0.0, 1.0]], dtype=torch.float32) + mask = torch.zeros((1, 3), dtype=torch.long) + step_calls = 0 + + def _step(): + nonlocal step_calls + step_calls += 1 + + monkeypatch.setattr(trainer.optimizer, "step", _step) + + out = trainer._batch_step((x, y, mask), train=True) + assert out["n"] == 0 + assert out["loss"] == pytest.approx(0.0, abs=1e-12) + assert out["global_valid_count"] == 0 + assert out["optimizer_step"] is False + assert step_calls == 0 + assert all( + parameter.grad is not None + and torch.count_nonzero(parameter.grad).item() == 0 + for parameter in model.parameters() + ) + + +def test_batch_step_zero_local_mask_steps_when_global_valid(monkeypatch): model = PepSeqFFNN( emb_dim=4, hidden_sizes=(8,), @@ -42,10 +87,54 @@ def test_batch_step_zero_mask_returns_zero_n(): x = torch.randn(1, 3, 4) y = torch.tensor([[1.0, 0.0, 1.0]], dtype=torch.float32) mask = torch.zeros((1, 3), dtype=torch.long) + step_calls = 0 + + def _step(): + nonlocal step_calls + step_calls += 1 + + def _global_valid(t: torch.Tensor) -> torch.Tensor: + out = t.clone() + out[0] = 1 + out[1] = 2 + return out + + monkeypatch.setattr(trainer.optimizer, "step", _step) + monkeypatch.setattr(trainer_mod, "ddp_all_reduce_sum", _global_valid) out = trainer._batch_step((x, y, mask), train=True) assert out["n"] == 0 assert out["loss"] == pytest.approx(0.0, abs=1e-12) + assert out["global_valid_count"] == 1 + assert out["optimizer_step"] is True + assert step_calls == 1 + + +def test_batch_step_scales_training_loss_by_global_valid_count(monkeypatch): + model = ConstantLogitModel() + trainer = _make_trainer(model) + + x = torch.randn(1, 2, 4) + y = torch.zeros((1, 2), dtype=torch.float32) + mask = torch.tensor([[1, 0]], dtype=torch.long) + + def _step(): + return None + + def _global_stats(t: torch.Tensor) -> torch.Tensor: + out = t.clone() + out[0] = 4 + out[1] = 2 + return out + + monkeypatch.setattr(trainer.optimizer, "step", _step) + monkeypatch.setattr(trainer_mod, "ddp_all_reduce_sum", _global_stats) + + out = trainer._batch_step((x, y, mask), train=True) + assert out["n"] == 1 + assert out["global_valid_count"] == 4 + assert out["optimizer_step"] is True + assert model.p.grad.item() == pytest.approx(0.25, abs=1e-12) def test_batch_step_rejects_invalid_y_dim(): diff --git a/tests/unit/core/train/test_trainer_fit.py b/tests/unit/core/train/test_trainer_fit.py index 6fe6e1b..bfe7f5f 100644 --- a/tests/unit/core/train/test_trainer_fit.py +++ b/tests/unit/core/train/test_trainer_fit.py @@ -54,6 +54,80 @@ def _make_trainer(train_loader, val_loader=None, emb_dim: int = 4, epochs: int = ) +def test_make_zero_valid_dummy_batch_supports_masked_and_unmasked_batches(): + trainer = _make_trainer(_make_batches(n_batches=1), epochs=1) + x, y, mask = _make_batches(n_batches=1)[0] + + masked_dummy = trainer._make_zero_valid_dummy_batch((x, y, mask)) + unmasked_dummy = trainer._make_zero_valid_dummy_batch((x, y)) + + for dummy_x, dummy_y, dummy_mask in (masked_dummy, unmasked_dummy): + assert dummy_x.shape == (1, 1, x.size(-1)) + assert dummy_y.shape == (1, 1) + assert dummy_mask.shape == (1, 1) + assert torch.count_nonzero(dummy_x).item() == 0 + assert torch.count_nonzero(dummy_y).item() == 0 + assert torch.count_nonzero(dummy_mask).item() == 0 + + +def test_synchronized_training_batches_yields_dummy_after_local_exhaustion(monkeypatch): + batches = _make_batches(n_batches=1) + trainer = _make_trainer(batches, epochs=1) + reduced_states = iter(((2, 2), (1, 2), (0, 2))) + local_states = [] + + def _reduce_active(tensor: torch.Tensor) -> torch.Tensor: + local_states.append(tuple(int(x) for x in tensor.tolist())) + active, world = next(reduced_states) + return torch.tensor( + [active, world], device=tensor.device, dtype=tensor.dtype) + + monkeypatch.setattr(trainer_mod, "ddp_all_reduce_sum", _reduce_active) + + synchronized = list(trainer._synchronized_training_batches(batches)) + + assert local_states == [(1, 1), (0, 1), (0, 1)] + assert len(synchronized) == 2 + assert synchronized[0][0] is batches[0] + assert synchronized[0][1] is True + dummy_batch, is_real = synchronized[1] + assert is_real is False + assert dummy_batch[0].shape == (1, 1, 4) + assert torch.count_nonzero(dummy_batch[2]).item() == 0 + + +def test_synchronized_training_batches_rejects_initial_empty_rank(monkeypatch): + trainer = _make_trainer([], epochs=1) + + def _one_active_rank(tensor: torch.Tensor) -> torch.Tensor: + return torch.tensor([1, 2], device=tensor.device, dtype=tensor.dtype) + + monkeypatch.setattr(trainer_mod, "ddp_all_reduce_sum", _one_active_rank) + + with pytest.raises(RuntimeError, match="no initial training batch"): + list(trainer._synchronized_training_batches([])) + + +def test_run_epoch_train_reports_single_rank_synchronization(): + batches = _make_batches(n_batches=2) + trainer = _make_trainer(batches, epochs=1) + + out = trainer._run_epoch(0, train=True) + + assert out["synchronized_steps"] == 2 + assert out["optimizer_steps"] == 2 + assert out["zero_valid_steps"] == 0 + assert out["real_batches"] == 2 + assert out["dummy_batches"] == 0 + assert out["sync_summary"] == { + "synchronized_steps": 2, + "optimizer_steps": 2, + "zero_valid_steps": 0, + "dummy_batch_fraction": 0.0, + "per_rank": [{"rank": 0, "real_batches": 2, "dummy_batches": 0}], + } + + def test_fit_with_validation_saves_checkpoint(tmp_path: Path): trainer = _make_trainer( _make_batches(), _make_batches(), emb_dim=4, epochs=2) @@ -62,6 +136,10 @@ def test_fit_with_validation_saves_checkpoint(tmp_path: Path): assert summary["best_epoch"] >= 0 assert (tmp_path / "fully_connected.pt").exists() assert isinstance(summary["best_metrics"], dict) + metrics = summary["best_metrics"] + assert metrics["threshold_policy"] == "max-recall-min-precision" + assert metrics["threshold_min_precision"] == pytest.approx(0.25) + assert isinstance(metrics["threshold_grid"], list) def test_fit_without_validation_saves_no_val_checkpoint(tmp_path: Path): @@ -103,6 +181,21 @@ def test_run_epoch_eval_no_valid_residues(): out = trainer._run_epoch(0, train=False) assert out["eval_metrics"]["threshold_status"] == "no_valid_residues" + assert out["eval_metrics"]["threshold_policy"] == "max-recall-min-precision" + + +def test_run_epoch_eval_records_configured_threshold_policy(): + train_loader = _make_batches(n_batches=1, mask_value=1) + val_loader = _make_batches(n_batches=1, mask_value=1) + trainer = _make_trainer(train_loader, val_loader, emb_dim=4, epochs=1) + trainer.config.threshold_policy = "fixed" + trainer.config.threshold_fixed_value = 0.5 + + out = trainer._run_epoch(0, train=False) + metrics = out["eval_metrics"] + assert metrics["threshold_policy"] == "fixed" + assert metrics["threshold"] == pytest.approx(0.5) + assert metrics["threshold_status"] == "ok" class _AlwaysPruneTrial: diff --git a/tests/unit/core/train/test_weights.py b/tests/unit/core/train/test_weights.py index afa54c7..c5f3727 100644 --- a/tests/unit/core/train/test_weights.py +++ b/tests/unit/core/train/test_weights.py @@ -7,6 +7,7 @@ import pepseqpred.core.train.weights as weights_mod from pepseqpred.core.train.weights import ( compute_pos_neg_counts, + global_pos_neg_counts, global_pos_weight, pos_weight_from_label_shards ) @@ -70,6 +71,28 @@ def test_global_pos_weight_without_ddp_uses_safe_denominator(): assert global_pos_weight(local_pos=0, local_neg=9, ddp=None) == pytest.approx(9.0) +def test_global_pos_neg_counts_without_ddp_returns_local_counts(): + assert global_pos_neg_counts(local_pos=2, local_neg=5, ddp=None) == (2, 5) + + +def test_global_pos_neg_counts_with_ddp_all_reduce(monkeypatch): + original_tensor = torch.tensor + + def _cpu_tensor(data, device=None, **kwargs): + _ = device + return original_tensor(data, device=torch.device("cpu"), **kwargs) + + def _all_reduce(t, op): + assert op == weights_mod.dist.ReduceOp.SUM + t[0] += 4 + t[1] += 6 + + monkeypatch.setattr(weights_mod.torch, "tensor", _cpu_tensor) + monkeypatch.setattr(weights_mod.dist, "all_reduce", _all_reduce) + + assert global_pos_neg_counts(local_pos=1, local_neg=2, ddp={"rank": 0}) == (5, 8) + + def test_global_pos_weight_with_ddp_all_reduce(monkeypatch): original_tensor = torch.tensor