From 0ed97287014ee745ee5fa797a88e0256ae30d57c Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Thu, 9 Jul 2026 13:56:24 -0400 Subject: [PATCH 1/2] Add predict-table command to annotate CSVs with predictor scores (#231) Downstream evaluation workflows often start from an annotated benchmark table (sample_id, hit, peptide, genotype/allele columns) and just need external predictor scores appended. Add a general, additive way to do that. Library core (mhctools/annotate.py): * annotate_table(df, specs, peptide_column, allele_column, ...) -> DataFrame. Runs each predictor ONCE over the union of (peptide, allele) pairs, then a per-row lookup picks the best allele. Preserves all input columns; appends one score column per predictor plus a _best_allele provenance column. * "Best" reuses mhctools.pred.best_direction: score higher-better, affinity and percentile_rank lower-better. * Row alleles are normalized with normalize_allele_name_or_raw, so "A0201" matches a prediction emitted as "HLA-A*02:01" and exotic un-normalizable alleles round-trip (#220). Multiple alleles per cell (whitespace/comma/ semicolon) are supported. * AnnotationSpec accepts a built predictor or an alleles->predictor factory, so commandline predictors are constructed with exactly the table's alleles. * parse_annotation_spec("NAME:COLUMN:FIELD") builds a spec from the CLI registry; COLUMN and FIELD default to NAME_FIELD and "affinity". CLI (mhctools/cli/annotate_table.py): thin I/O wrapper exposed as the `mhctools predict-table` subcommand (dispatched from cli/script.main). Reads CSV/CSV.bz2, writes the annotated table, and optionally a --predictor-info sidecar (predictor, output_column, score_field, higher_is_better). Exports annotate_table / AnnotationSpec / parse_annotation_spec from the package. Binary-free tests use a deterministic fixture predictor to assert direction handling, best-allele selection, column preservation, collision/ overwrite, NaN fallbacks, and spec parsing (added to the CI public subset). Version 3.20.1 -> 3.21.0. Claude-Session: https://claude.ai/code/session_01LZahFhBSCiehXTESCYQ7wG --- .github/workflows/tests.yml | 1 + README.md | 42 +++++ mhctools/__init__.py | 6 +- mhctools/annotate.py | 322 +++++++++++++++++++++++++++++++++ mhctools/cli/annotate_table.py | 128 +++++++++++++ mhctools/cli/script.py | 10 + tests/test_annotate_table.py | 306 +++++++++++++++++++++++++++++++ 7 files changed, 814 insertions(+), 1 deletion(-) create mode 100644 mhctools/annotate.py create mode 100644 mhctools/cli/annotate_table.py create mode 100644 tests/test_annotate_table.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index e19a802..a489028 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -52,6 +52,7 @@ jobs: tests/test_netchop_parse.py \ tests/test_process_helpers.py \ tests/test_pred.py \ + tests/test_annotate_table.py \ tests/test_processing_predictor.py \ tests/test_pepsickle.py \ tests/test_bigmhc.py \ diff --git a/README.md b/README.md index cd799c2..95ec870 100644 --- a/README.md +++ b/README.md @@ -372,6 +372,48 @@ mhctools --sequence SIINFEKL SIINFEKLQ --mhc-predictor netmhc --mhc-alleles A020 mhctools --sequence AAAQQQSIINFEKL --extract-subsequences --mhc-peptide-lengths 8-10 --mhc-predictor mhcflurry --mhc-alleles A0201 ``` +### Annotate an existing table with predictor scores (`predict-table`) + +Downstream evaluation workflows often start from an annotated benchmark table +(with columns like `sample_id`, `hit`, `peptide`, and per-row genotype/allele +info) and just need external predictor scores appended. `mhctools +predict-table` reads a CSV, runs each requested predictor once, and appends one +score column per predictor — choosing the best allele per row — while +preserving every input column: + +```sh +mhctools predict-table \ + --input benchmark.csv.bz2 \ + --peptide-column peptide \ + --alleles-column hla \ + --predictor netmhcpan42-ba:netmhcpan4.2.ba:affinity \ + --predictor netmhcpan42-el:netmhcpan4.2.el:score \ + --out benchmark.with_scores.csv.bz2 +``` + +Each `--predictor` spec is `NAME[:OUTPUT_COLUMN[:FIELD]]`, where `FIELD` is +`affinity`, `score`, or `percentile_rank` (lower is better for `affinity` and +`percentile_rank`; higher for `score`). Rows may hold several alleles per cell +(whitespace-, comma-, or semicolon-separated); the best one per peptide is +chosen and recorded in a `_best_allele` provenance column. +Pass `--predictor-info info.csv` to also write a sidecar describing each +column's `score_field` and `higher_is_better`. + +The same thing from Python (I/O-free, works on any `DataFrame`): + +```python +from mhctools import annotate_table, AnnotationSpec, NetMHCpan42_BA + +annotated = annotate_table( + df, + [AnnotationSpec( + predictor=lambda alleles: NetMHCpan42_BA(alleles=alleles), + output_column="netmhcpan4.2.ba", + field="affinity")], + peptide_column="peptide", + allele_column="hla") +``` + ## Legacy API The old `predict_peptides()` and `predict_subsequences()` methods still work diff --git a/mhctools/__init__.py b/mhctools/__init__.py index d060029..f5e440a 100644 --- a/mhctools/__init__.py +++ b/mhctools/__init__.py @@ -15,6 +15,7 @@ ) from .sample import MultiSample from .tcr import TCR +from .annotate import AnnotationSpec, annotate_table, parse_annotation_spec from .iedb import ( IedbNetMHCcons, IedbNetMHCpan, @@ -79,7 +80,7 @@ def __getattr__(name): raise AttributeError( "module %r has no attribute %r" % (__name__, name)) -__version__ = "3.20.1" +__version__ = "3.21.0" __all__ = [ "Prediction", @@ -95,6 +96,9 @@ def __getattr__(name): "preds_from_rows", "MultiSample", "TCR", + "AnnotationSpec", + "annotate_table", + "parse_annotation_spec", "BindingPrediction", "BindingPredictionCollection", "IedbNetMHCcons", diff --git a/mhctools/annotate.py b/mhctools/annotate.py new file mode 100644 index 0000000..3dd72ec --- /dev/null +++ b/mhctools/annotate.py @@ -0,0 +1,322 @@ +# Copyright (c) 2016. Mount Sinai School of Medicine +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Annotate a table of peptides (and optional alleles) with predictor scores. + +Downstream evaluation workflows often start from an annotated benchmark +table (columns like ``sample_id``, ``hit``, ``peptide``, and genotype/allele +information) rather than a bare peptide list. :func:`annotate_table` runs one +or more predictors over such a table and appends one score column per +predictor, choosing the best allele per row, while preserving every input +column unchanged. + +The core function is I/O-free and works on any :class:`pandas.DataFrame`; the +``mhctools predict-table`` subcommand (see :mod:`mhctools.cli.annotate_table`) +wraps it with CSV reading/writing. + +Design notes +------------ +* Each predictor runs **once** over the union of all ``(peptide, allele)`` + pairs in the table (leveraging the predictors' allele batching), then a + per-row lookup picks the best allele. Predictors are not re-invoked per row. +* "Best" comes from :func:`mhctools.pred.best_direction`: ``score`` is + higher-better, ``affinity`` and ``percentile_rank`` are lower-better. +* Row alleles are normalized with + :func:`mhctools.allele_normalization.normalize_allele_name_or_raw` so a + cell written ``A0201`` matches a prediction emitted as ``HLA-A*02:01``, and + exotic un-normalizable alleles still round-trip (see #220). +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field as _dc_field +from typing import Optional + +import pandas as pd + +from .allele_normalization import normalize_allele_name_or_raw +from .pred import Kind, best_direction + + +# Whitespace, comma, or semicolon separate multiple alleles packed into one +# cell (e.g. a genotype column "HLA-A*02:01 HLA-B*07:02" or "A0201,B0702"). +_DEFAULT_ALLELE_SEP = re.compile(r"[\s,;]+") + + +# Output-field token -> (kind, Prediction attribute). ``kind`` is ``None`` when +# the field's "best" direction is kind-independent (``score`` is always +# higher-better, ``percentile_rank`` always lower-better); ``affinity`` maps to +# the ``value`` of ``pMHC_affinity`` predictions (IC50 nM, lower-better). +_OUTPUT_FIELDS = { + "affinity": (Kind.pMHC_affinity, "value"), + "score": (None, "score"), + "percentile_rank": (None, "percentile_rank"), + # convenience aliases + "rank": (None, "percentile_rank"), + "value": (Kind.pMHC_affinity, "value"), + "presentation": (Kind.pMHC_presentation, "score"), + "stability": (Kind.pMHC_stability, "value"), + "immunogenicity": (Kind.immunogenicity, "score"), +} + + +def output_field_tokens(): + """Sorted list of accepted output-field tokens (for CLI help / errors).""" + return sorted(_OUTPUT_FIELDS) + + +@dataclass +class AnnotationSpec: + """One predictor mapped to one output column. + + Parameters + ---------- + predictor : predictor instance, or callable ``alleles -> predictor`` + A built predictor (used as-is), or a factory that builds one given + the union of alleles in the table. Commandline predictors validate + their alleles at construction, so passing a factory lets the predictor + be built with exactly the alleles the table needs. + output_column : str + Name of the appended score column. + field : str + Which prediction field to write, one of :func:`output_field_tokens` + (``"affinity"``, ``"score"``, ``"percentile_rank"``, ...). Determines + both the value written and the best-allele direction. + add_best_allele : bool + If True (default) and the table has an allele column, also append a + provenance column naming the allele that won each row. + best_allele_column : str, optional + Name of the provenance column. Defaults to + ``f"{output_column}_best_allele"``. + """ + + predictor: object + output_column: str + field: str = "affinity" + add_best_allele: bool = True + best_allele_column: Optional[str] = None + kind: Optional[str] = _dc_field(default=None, init=False) + prediction_field: str = _dc_field(default="", init=False) + + def __post_init__(self): + if self.field not in _OUTPUT_FIELDS: + raise ValueError( + "Unknown output field %r. Available: %s" + % (self.field, ", ".join(output_field_tokens()))) + self.kind, self.prediction_field = _OUTPUT_FIELDS[self.field] + + def resolved_best_allele_column(self): + return self.best_allele_column or ("%s_best_allele" % self.output_column) + + def build_predictor(self, alleles): + """Return a predictor instance for the given union of alleles. + + A callable that is not itself a predictor (no ``predict`` method) is + treated as a factory and called with ``alleles``; anything else is + assumed to already be a built predictor and returned unchanged. + """ + predictor = self.predictor + if callable(predictor) and not hasattr(predictor, "predict"): + return predictor(alleles) + return predictor + + def direction_op(self): + """``max`` or ``min`` callable for reducing candidate predictions.""" + return max if best_direction(self.kind, self.prediction_field) == "max" \ + else min + + +def parse_annotation_spec(token): + """Parse a CLI predictor spec ``NAME:OUTPUT_COLUMN:FIELD``. + + ``OUTPUT_COLUMN`` and ``FIELD`` are optional; they default to + ``f"{NAME}_{FIELD}"`` and ``"affinity"`` respectively. ``NAME`` must be a + key in :data:`mhctools.cli.args.mhc_predictors`. + + Examples + -------- + ``"netmhcpan42-ba:netmhcpan4.2.ba:affinity"``, + ``"netmhcpan42-el:netmhcpan4.2.el:score"``, ``"mixmhcpred"``. + """ + # Import the registry lazily to keep the core module free of the CLI's + # (heavier) import graph and avoid an import cycle. + from .cli.args import mhc_predictors, _cls_accepts + + parts = [p.strip() for p in str(token).split(":")] + name = parts[0].lower() + field = parts[2].lower() if len(parts) > 2 and parts[2] else "affinity" + output_column = parts[1] if len(parts) > 1 and parts[1] else "%s_%s" % (name, field) + + if name not in mhc_predictors: + raise ValueError( + "Unknown predictor %r. Available: %s" + % (name, ", ".join(sorted(mhc_predictors.keys())))) + if field not in _OUTPUT_FIELDS: + raise ValueError( + "Unknown output field %r in spec %r. Available: %s" + % (field, token, ", ".join(output_field_tokens()))) + + cls = mhc_predictors[name] + + def factory(alleles): + if alleles is not None and _cls_accepts(cls, "alleles"): + return cls(alleles=alleles) + return cls() + + return AnnotationSpec( + predictor=factory, output_column=output_column, field=field) + + +def _split_alleles(cell, allele_sep): + """Split one table cell into a list of normalized allele names.""" + if cell is None or (isinstance(cell, float) and pd.isna(cell)): + return [] + tokens = [t for t in allele_sep.split(str(cell).strip()) if t] + # Normalize + de-dup while preserving order. + seen = {} + for token in tokens: + normalized = normalize_allele_name_or_raw(token) + seen.setdefault(normalized, None) + return list(seen) + + +def _build_lookup(results, spec): + """Index one predictor's results by (peptide, allele) and by peptide. + + Returns ``(by_pair, by_peptide)`` where ``by_pair`` maps + ``(peptide, allele) -> Prediction`` (allele-bearing predictions) and + ``by_peptide`` maps ``peptide -> Prediction`` (allele-free predictions, + e.g. processing predictors). Predictions whose target field is ``None`` + are skipped. + """ + by_pair = {} + by_peptide = {} + for peptide_result in results: + for pred in peptide_result.filter(kind=spec.kind): + if getattr(pred, spec.prediction_field) is None: + continue + if pred.allele: + by_pair[(pred.peptide, pred.allele)] = pred + else: + by_peptide[pred.peptide] = pred + return by_pair, by_peptide + + +def annotate_table( + table, + specs, + peptide_column="peptide", + allele_column=None, + allele_sep=_DEFAULT_ALLELE_SEP, + overwrite=False): + """Append predictor score columns to a table, best-allele per row. + + Parameters + ---------- + table : pandas.DataFrame + Input table; returned unmodified with new columns appended to a copy. + specs : sequence of AnnotationSpec + One entry per output column. + peptide_column : str + Column holding the peptide sequence for each row. + allele_column : str, optional + Column holding the row's allele(s). May contain several alleles per + cell (see ``allele_sep``). If omitted, predictors run allele-free and + no best-allele provenance column is written. + allele_sep : compiled regex + Splits multiple alleles packed into one cell. Defaults to splitting on + whitespace, commas, and semicolons. + overwrite : bool + If False (default), raise when an output column already exists. + + Returns + ------- + pandas.DataFrame + A copy of ``table`` with one numeric column per spec appended (plus a + ``_best_allele`` column when an allele column is given). + Rows with no usable prediction get ``NaN`` (and ``None`` best allele). + + Notes + ----- + Every allele in the table must be supported by each predictor; + commandline predictors raise ``UnsupportedAllele`` at construction + otherwise, matching mhctools' behavior elsewhere. + """ + if not isinstance(table, pd.DataFrame): + raise TypeError( + "annotate_table expects a pandas DataFrame, got %s; use the " + "predict-table CLI to read a file." % type(table).__name__) + + df = table.copy() + + if peptide_column not in df.columns: + raise KeyError( + "peptide column %r not found; columns are: %s" + % (peptide_column, list(df.columns))) + if allele_column is not None and allele_column not in df.columns: + raise KeyError( + "allele column %r not found; columns are: %s" + % (allele_column, list(df.columns))) + + specs = list(specs) + # Fail fast on any output-column collision before running predictors. + for spec in specs: + new_columns = [spec.output_column] + if spec.add_best_allele and allele_column is not None: + new_columns.append(spec.resolved_best_allele_column()) + for column in new_columns: + if column in df.columns and not overwrite: + raise ValueError( + "output column %r already exists; pass overwrite=True to " + "replace it" % column) + + peptides = df[peptide_column].astype(str).tolist() + if allele_column is not None: + row_alleles = [_split_alleles(cell, allele_sep) for cell in df[allele_column]] + else: + row_alleles = [[] for _ in peptides] + + union_peptides = list(dict.fromkeys(peptides)) + union_alleles = sorted({a for alleles in row_alleles for a in alleles}) + + for spec in specs: + predictor = spec.build_predictor(union_alleles or None) + results = predictor.predict(union_peptides) + by_pair, by_peptide = _build_lookup(results, spec) + reduce_op = spec.direction_op() + field = spec.prediction_field + + values = [] + best_alleles = [] + for peptide, alleles in zip(peptides, row_alleles): + candidates = [by_pair[(peptide, a)] for a in alleles + if (peptide, a) in by_pair] + if not candidates and not alleles: + allele_free = by_peptide.get(peptide) + if allele_free is not None: + candidates = [allele_free] + if candidates: + best = reduce_op(candidates, key=lambda p: getattr(p, field)) + values.append(getattr(best, field)) + best_alleles.append(best.allele or None) + else: + values.append(float("nan")) + best_alleles.append(None) + + df[spec.output_column] = values + if spec.add_best_allele and allele_column is not None: + df[spec.resolved_best_allele_column()] = best_alleles + + return df diff --git a/mhctools/cli/annotate_table.py b/mhctools/cli/annotate_table.py new file mode 100644 index 0000000..c68f00b --- /dev/null +++ b/mhctools/cli/annotate_table.py @@ -0,0 +1,128 @@ +# Copyright (c) 2016. Mount Sinai School of Medicine +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``mhctools predict-table`` — annotate a CSV of peptides with predictor scores. + +Thin I/O wrapper around :func:`mhctools.annotate.annotate_table`: reads a +CSV (optionally compressed), runs the requested predictors, and writes the +input table back with one appended score column per predictor. +""" + +from argparse import ArgumentParser + +import pandas as pd + +from ..annotate import ( + annotate_table, + parse_annotation_spec, + output_field_tokens, +) +from ..pred import best_direction +from ..logging import get_logger + + +logger = get_logger(__name__) + + +def make_arg_parser(): + parser = ArgumentParser( + prog="mhctools predict-table", + description=( + "Annotate a CSV table of peptides (and optional alleles) with " + "predictor score columns, preserving all input columns.")) + parser.add_argument( + "--input", + required=True, + help="Input CSV (compression inferred from extension: .csv, .bz2, .gz).") + parser.add_argument( + "--out", + required=True, + help="Output CSV path (compression inferred from extension).") + parser.add_argument( + "--peptide-column", + default="peptide", + help="Column holding the peptide sequence (default: peptide).") + parser.add_argument( + "--alleles-column", + default=None, + help=( + "Column holding the row's allele(s). Multiple alleles per cell " + "may be separated by whitespace, comma, or semicolon. Omit for " + "allele-free predictors.")) + parser.add_argument( + "--predictor", + dest="predictors", + action="append", + required=True, + metavar="NAME:OUTPUT_COLUMN:FIELD", + help=( + "Predictor spec, repeatable. OUTPUT_COLUMN and FIELD are optional " + "(default column NAME_FIELD, default field 'affinity'). FIELD is " + "one of: %s. E.g. 'netmhcpan42-ba:netmhcpan4.2.ba:affinity'." + % ", ".join(output_field_tokens()))) + parser.add_argument( + "--predictor-info", + default=None, + help=( + "Optional path to write a sidecar CSV describing each output " + "column: predictor spec, output_column, score_field, " + "higher_is_better.")) + parser.add_argument( + "--overwrite", + action="store_true", + default=False, + help="Overwrite output columns if they already exist in the input.") + return parser + + +def _write_predictor_info(path, raw_specs, specs): + rows = [] + for raw, spec in zip(raw_specs, specs): + rows.append({ + "predictor": raw, + "output_column": spec.output_column, + "score_field": spec.field, + "higher_is_better": + best_direction(spec.kind, spec.prediction_field) == "max", + }) + pd.DataFrame(rows).to_csv(path, index=False) + print("Wrote predictor info: %s" % path) + + +def main(args_list=None): + args = make_arg_parser().parse_args(args_list) + + df = pd.read_csv(args.input) + logger.info("Read %d rows, %d columns from %s", + len(df), len(df.columns), args.input) + + specs = [parse_annotation_spec(token) for token in args.predictors] + + annotated = annotate_table( + df, + specs, + peptide_column=args.peptide_column, + allele_column=args.alleles_column, + overwrite=args.overwrite) + + annotated.to_csv(args.out, index=False) + print("Wrote: %s (%d rows, %d columns)" + % (args.out, len(annotated), len(annotated.columns))) + + if args.predictor_info: + _write_predictor_info(args.predictor_info, args.predictors, specs) + + +if __name__ == "__main__": + main() diff --git a/mhctools/cli/script.py b/mhctools/cli/script.py index bcbf8cf..bee4ecb 100644 --- a/mhctools/cli/script.py +++ b/mhctools/cli/script.py @@ -152,7 +152,17 @@ def main(args_list=None): --mhc-alleles HLA-A0201 H2-Db \ --mhc-predictor netmhc \ --output-csv epitope.csv + + The ``predict-table`` subcommand annotates an existing CSV of peptides + (and optional alleles) with predictor score columns; see + ``mhctools predict-table --help``. """ + if args_list is None: + args_list = sys.argv[1:] + if args_list and args_list[0] == "predict-table": + from .annotate_table import main as annotate_table_main + return annotate_table_main(args_list[1:]) + args = parse_args(args_list) binding_predictions = run_predictor(args) df = binding_predictions.to_dataframe() diff --git a/tests/test_annotate_table.py b/tests/test_annotate_table.py new file mode 100644 index 0000000..5ccc44e --- /dev/null +++ b/tests/test_annotate_table.py @@ -0,0 +1,306 @@ +# Copyright (c) 2016. Mount Sinai School of Medicine +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Binary-free tests for mhctools.annotate.annotate_table. + +Uses a deterministic in-memory predictor so best-allele selection and +direction handling can be asserted exactly (RandomBindingPredictor returns +fresh random scores on each call, so it can only be used for a structural +smoke test). +""" + +import math + +import pandas as pd +import pytest + +from mhctools import ( + annotate_table, + AnnotationSpec, + RandomBindingPredictor, +) +from mhctools.annotate import parse_annotation_spec, output_field_tokens +from mhctools.pred import Kind, PeptideResult, Prediction + + +# A fixed (peptide, allele) -> (affinity nM, score, percentile_rank) table. +# Chosen so the best allele differs by field, to catch direction bugs: +# SIINFEKL: A*02:01 has lower IC50 (better affinity) but LOWER score; +# B*07:02 has higher IC50 but HIGHER score. +_FIXTURE = { + ("SIINFEKL", "HLA-A*02:01"): (100.0, 0.20, 5.0), + ("SIINFEKL", "HLA-B*07:02"): (500.0, 0.90, 1.0), + ("GILGFVFTL", "HLA-A*02:01"): (50.0, 0.80, 0.5), + ("GILGFVFTL", "HLA-B*07:02"): (9000.0, 0.05, 40.0), +} + + +class _FixturePredictor: + """Deterministic predictor returning predictions from ``_FIXTURE``. + + Only implements the ``predict`` method that ``annotate_table`` calls. + Emits pMHC_affinity predictions carrying value (IC50), score, and rank. + """ + + def __init__(self, alleles): + self.alleles = list(alleles) if alleles else [] + + def predict(self, peptides): + results = [] + for peptide in peptides: + preds = [] + for allele in self.alleles: + key = (peptide, allele) + if key not in _FIXTURE: + continue + affinity, score, rank = _FIXTURE[key] + preds.append(Prediction( + kind=Kind.pMHC_affinity, + peptide=peptide, + allele=allele, + score=score, + value=affinity, + percentile_rank=rank, + predictor_name="fixture")) + results.append(PeptideResult(preds=tuple(preds))) + return results + + +def _factory(alleles): + return _FixturePredictor(alleles) + + +def _table(): + return pd.DataFrame({ + "sample_id": ["s1", "s2"], + "hit": [1, 0], + "peptide": ["SIINFEKL", "GILGFVFTL"], + "hla": ["HLA-A*02:01 HLA-B*07:02", "A0201,B0702"], + }) + + +def test_affinity_picks_lowest_ic50(): + out = annotate_table( + _table(), + [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + row = out[out.peptide == "SIINFEKL"].iloc[0] + # A*02:01 has the lower IC50 (100 < 500) -> better affinity + assert row["aff"] == 100.0 + assert row["aff_best_allele"] == "HLA-A*02:01" + row2 = out[out.peptide == "GILGFVFTL"].iloc[0] + assert row2["aff"] == 50.0 + assert row2["aff_best_allele"] == "HLA-A*02:01" + + +def test_score_picks_highest_and_differs_from_affinity(): + out = annotate_table( + _table(), + [AnnotationSpec(_factory, "sc", field="score")], + allele_column="hla") + row = out[out.peptide == "SIINFEKL"].iloc[0] + # B*07:02 has the higher score (0.90 > 0.20) even though worse affinity + assert row["sc"] == 0.90 + assert row["sc_best_allele"] == "HLA-B*07:02" + + +def test_percentile_rank_picks_lowest(): + out = annotate_table( + _table(), + [AnnotationSpec(_factory, "pr", field="percentile_rank")], + allele_column="hla") + row = out[out.peptide == "SIINFEKL"].iloc[0] + # B*07:02 has the lower rank (1.0 < 5.0) -> better + assert row["pr"] == 1.0 + assert row["pr_best_allele"] == "HLA-B*07:02" + + +def test_preserves_input_columns_and_order(): + df = _table() + out = annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + # original columns come first, unchanged, then the two appended columns + assert list(out.columns) == list(df.columns) + ["aff", "aff_best_allele"] + pd.testing.assert_frame_equal(out[df.columns], df) + + +def test_does_not_mutate_input(): + df = _table() + before = df.copy() + annotate_table(df, [AnnotationSpec(_factory, "aff")], allele_column="hla") + pd.testing.assert_frame_equal(df, before) + + +def test_multiple_specs_appended_independently(): + out = annotate_table( + _table(), + [AnnotationSpec(_factory, "aff", field="affinity"), + AnnotationSpec(_factory, "sc", field="score", add_best_allele=False)], + allele_column="hla") + assert "aff" in out.columns and "aff_best_allele" in out.columns + assert "sc" in out.columns + # add_best_allele=False suppresses the provenance column + assert "sc_best_allele" not in out.columns + + +def test_missing_allele_yields_nan(): + df = pd.DataFrame({ + "peptide": ["SIINFEKL"], + "hla": ["HLA-C*07:01"], # not in the fixture + }) + out = annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + assert math.isnan(out.iloc[0]["aff"]) + assert out.iloc[0]["aff_best_allele"] is None + + +def test_unknown_peptide_yields_nan(): + df = pd.DataFrame({ + "peptide": ["WWWWWWWWW"], # not in the fixture + "hla": ["HLA-A*02:01"], + }) + out = annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + assert math.isnan(out.iloc[0]["aff"]) + + +def test_empty_allele_cell_yields_nan(): + df = pd.DataFrame({ + "peptide": ["SIINFEKL", "GILGFVFTL"], + "hla": ["HLA-A*02:01", None], + }) + out = annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + assert out.iloc[0]["aff"] == 100.0 + assert math.isnan(out.iloc[1]["aff"]) + assert out.iloc[1]["aff_best_allele"] is None + + +def test_column_collision_raises(): + df = _table() + df["aff"] = 0.0 + with pytest.raises(ValueError, match="already exists"): + annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + + +def test_column_collision_overwrite_allowed(): + df = _table() + df["aff"] = 0.0 + out = annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla", overwrite=True) + assert out.iloc[0]["aff"] == 100.0 + + +def test_missing_peptide_column_raises(): + with pytest.raises(KeyError, match="peptide"): + annotate_table( + pd.DataFrame({"seq": ["SIINFEKL"]}), + [AnnotationSpec(_factory, "aff")]) + + +def test_missing_allele_column_raises(): + with pytest.raises(KeyError, match="nope"): + annotate_table( + _table(), [AnnotationSpec(_factory, "aff")], + allele_column="nope") + + +def test_non_dataframe_raises(): + with pytest.raises(TypeError, match="DataFrame"): + annotate_table("not a df", [AnnotationSpec(_factory, "aff")]) + + +def test_custom_best_allele_column_name(): + out = annotate_table( + _table(), + [AnnotationSpec(_factory, "aff", field="affinity", + best_allele_column="winner")], + allele_column="hla") + assert "winner" in out.columns + assert "aff_best_allele" not in out.columns + + +def test_prebuilt_predictor_instance_used_directly(): + # Passing a built predictor (not a factory) should work: it is used as-is. + predictor = _FixturePredictor(["HLA-A*02:01", "HLA-B*07:02"]) + out = annotate_table( + _table(), [AnnotationSpec(predictor, "aff", field="affinity")], + allele_column="hla") + assert out[out.peptide == "SIINFEKL"].iloc[0]["aff"] == 100.0 + + +# --- spec parsing ----------------------------------------------------------- + +def test_parse_spec_full(): + spec = parse_annotation_spec("netmhcpan42-ba:mycol:affinity") + assert spec.output_column == "mycol" + assert spec.field == "affinity" + assert spec.kind == Kind.pMHC_affinity + assert spec.prediction_field == "value" + + +def test_parse_spec_defaults_column_and_field(): + spec = parse_annotation_spec("mixmhcpred") + assert spec.field == "affinity" + assert spec.output_column == "mixmhcpred_affinity" + + +def test_parse_spec_default_field_only(): + spec = parse_annotation_spec("netmhcpan42-el:elcol") + assert spec.output_column == "elcol" + assert spec.field == "affinity" + + +def test_parse_spec_unknown_predictor(): + with pytest.raises(ValueError, match="Unknown predictor"): + parse_annotation_spec("does-not-exist:c:affinity") + + +def test_parse_spec_unknown_field(): + with pytest.raises(ValueError, match="Unknown output field"): + parse_annotation_spec("mixmhcpred:c:bogus") + + +def test_bad_field_in_spec_object_raises(): + with pytest.raises(ValueError, match="Unknown output field"): + AnnotationSpec(_factory, "col", field="nonsense") + + +def test_output_field_tokens_include_primary_three(): + tokens = output_field_tokens() + for expected in ("affinity", "score", "percentile_rank"): + assert expected in tokens + + +# --- integration-flavored smoke test with the real random predictor --------- + +def test_random_predictor_smoke(): + df = _table() + out = annotate_table( + df, + [AnnotationSpec(lambda a: RandomBindingPredictor(alleles=a), + "rand", field="affinity")], + allele_column="hla") + assert list(out.columns) == list(df.columns) + ["rand", "rand_best_allele"] + # every row got a numeric prediction and a best allele from its own set + for _, row in out.iterrows(): + assert not math.isnan(row["rand"]) + assert row["rand_best_allele"] in {"HLA-A*02:01", "HLA-B*07:02"} From 179d7b8473f6333dae8c6b35e418aa91676b2cfe Mon Sep 17 00:00:00 2001 From: Alex Rubinsteyn Date: Thu, 9 Jul 2026 14:19:26 -0400 Subject: [PATCH 2/2] annotate_table: guard duplicate columns, normalize peptide join (review fixes) Fixes two issues from the #234 review: * Cross-spec output-column collision: the pre-run check only compared each spec against the *existing* input columns, so two specs targeting the same output (or best-allele) column passed and then silently clobbered each other. Now also reject a column planned by more than one spec, regardless of `overwrite` (overwrite replaces an existing column; it can't make two specs coexist in one column). * Peptide join was exact-string: a predictor that upper-cases/strips the peptides it echoes, against a table written lower-case or with stray whitespace, produced silent all-NaN. Peptides are now stripped + upper-cased on both the union sent to the predictor and the per-row lookup key (amino-acid sequences are canonically upper-case, so nothing is lost; the input peptide column is preserved verbatim). Also make the allele-free (by-peptide) fallback fire whenever the (peptide, allele) lookup misses, not only for rows with no alleles, so a processing/allele-free predictor still fills in scores when the table happens to carry an allele column. Only allele-free predictors populate the by-peptide index, so this never rescues a binding predictor's unsupported-allele miss (covered by a regression test). Adds 9 tests: cross-spec duplicate output/best-allele/cross-name collisions (incl. under overwrite), case- and whitespace-insensitive peptide matching with the input column preserved, the allele-free by-peptide path (with and without an allele column), and the binding-predictor unsupported-allele NaN regression guard. 33 passing. Claude-Session: https://claude.ai/code/session_01LZahFhBSCiehXTESCYQ7wG --- mhctools/annotate.py | 34 ++++++++- tests/test_annotate_table.py | 139 +++++++++++++++++++++++++++++++++++ 2 files changed, 169 insertions(+), 4 deletions(-) diff --git a/mhctools/annotate.py b/mhctools/annotate.py index 3dd72ec..3a4ca78 100644 --- a/mhctools/annotate.py +++ b/mhctools/annotate.py @@ -35,7 +35,8 @@ * Row alleles are normalized with :func:`mhctools.allele_normalization.normalize_allele_name_or_raw` so a cell written ``A0201`` matches a prediction emitted as ``HLA-A*02:01``, and - exotic un-normalizable alleles still round-trip (see #220). + exotic un-normalizable alleles still round-trip (see #220). Peptides are + upper-cased and stripped on both sides of the join for the same reason. """ from __future__ import annotations @@ -179,6 +180,19 @@ def factory(alleles): predictor=factory, output_column=output_column, field=field) +def _normalize_peptide(peptide): + """Canonicalize a peptide for matching: strip whitespace, uppercase. + + Predictions are joined back to rows by exact peptide string, and some + predictors uppercase/strip the peptides they echo. Normalizing both the + values sent to the predictor and the per-row lookup key the same way keeps + a table written ``siinfekl`` (or with stray whitespace) from silently + missing every prediction. Amino-acid sequences are canonically uppercase, + so this never loses information. + """ + return str(peptide).strip().upper() + + def _split_alleles(cell, allele_sep): """Split one table cell into a list of normalized allele names.""" if cell is None or (isinstance(cell, float) and pd.isna(cell)): @@ -271,7 +285,11 @@ def annotate_table( % (allele_column, list(df.columns))) specs = list(specs) - # Fail fast on any output-column collision before running predictors. + # Fail fast on any output-column collision before running predictors: + # against existing input columns (unless overwrite), and against columns + # any other spec plans to write (always — two specs writing the same + # column would silently clobber each other, which overwrite can't fix). + planned_columns = set() for spec in specs: new_columns = [spec.output_column] if spec.add_best_allele and allele_column is not None: @@ -281,8 +299,12 @@ def annotate_table( raise ValueError( "output column %r already exists; pass overwrite=True to " "replace it" % column) + if column in planned_columns: + raise ValueError( + "output column %r is produced by more than one spec" % column) + planned_columns.add(column) - peptides = df[peptide_column].astype(str).tolist() + peptides = [_normalize_peptide(p) for p in df[peptide_column]] if allele_column is not None: row_alleles = [_split_alleles(cell, allele_sep) for cell in df[allele_column]] else: @@ -303,7 +325,11 @@ def annotate_table( for peptide, alleles in zip(peptides, row_alleles): candidates = [by_pair[(peptide, a)] for a in alleles if (peptide, a) in by_pair] - if not candidates and not alleles: + if not candidates: + # Allele-free predictors (e.g. processing) emit allele-less + # predictions indexed by peptide only. Only such predictors + # populate ``by_peptide``, so this fallback never masks a + # genuine unsupported-allele miss from a binding predictor. allele_free = by_peptide.get(peptide) if allele_free is not None: candidates = [allele_free] diff --git a/tests/test_annotate_table.py b/tests/test_annotate_table.py index 5ccc44e..ffd4195 100644 --- a/tests/test_annotate_table.py +++ b/tests/test_annotate_table.py @@ -81,6 +81,36 @@ def _factory(alleles): return _FixturePredictor(alleles) +# A fixed peptide -> processing score table for the allele-free path. +_PROCESSING = {"SIINFEKL": 0.80, "GILGFVFTL": 0.30} + + +class _ProcessingFixturePredictor: + """Deterministic allele-free predictor (like a processing predictor). + + Emits antigen_processing predictions with no allele, so annotate_table + routes them through the by-peptide lookup rather than the (peptide, + allele) lookup. + """ + + def __init__(self, alleles=None): + pass + + def predict(self, peptides): + results = [] + for peptide in peptides: + preds = [] + if peptide in _PROCESSING: + preds.append(Prediction( + kind=Kind.antigen_processing, + peptide=peptide, + allele="", + score=_PROCESSING[peptide], + predictor_name="processing-fixture")) + results.append(PeptideResult(preds=tuple(preds))) + return results + + def _table(): return pd.DataFrame({ "sample_id": ["s1", "s2"], @@ -247,6 +277,115 @@ def test_prebuilt_predictor_instance_used_directly(): assert out[out.peptide == "SIINFEKL"].iloc[0]["aff"] == 100.0 +# --- duplicate output-column guard (cross-spec) ----------------------------- + +def test_duplicate_output_column_across_specs_raises(): + with pytest.raises(ValueError, match="more than one spec"): + annotate_table( + _table(), + [AnnotationSpec(_factory, "dup", field="affinity"), + AnnotationSpec(_factory, "dup", field="score")], + allele_column="hla") + + +def test_duplicate_output_column_across_specs_raises_even_with_overwrite(): + # overwrite lets you replace an *existing* column, but two specs writing + # the same new column would still clobber each other -> always an error. + with pytest.raises(ValueError, match="more than one spec"): + annotate_table( + _table(), + [AnnotationSpec(_factory, "dup", field="affinity"), + AnnotationSpec(_factory, "dup", field="score")], + allele_column="hla", overwrite=True) + + +def test_duplicate_best_allele_column_across_specs_raises(): + # distinct score columns, but the same explicit provenance column + with pytest.raises(ValueError, match="more than one spec"): + annotate_table( + _table(), + [AnnotationSpec(_factory, "a", field="affinity", + best_allele_column="prov"), + AnnotationSpec(_factory, "b", field="score", + best_allele_column="prov")], + allele_column="hla") + + +def test_output_column_collides_with_other_spec_best_allele_raises(): + # spec A's score column name equals spec B's provenance column name + with pytest.raises(ValueError, match="more than one spec"): + annotate_table( + _table(), + [AnnotationSpec(_factory, "b_best_allele", field="affinity"), + AnnotationSpec(_factory, "b", field="score")], + allele_column="hla") + + +# --- peptide normalization (case / whitespace insensitive matching) --------- + +def test_peptide_matched_case_insensitively(): + df = pd.DataFrame({"peptide": ["siinfekl"], "hla": ["HLA-A*02:01"]}) + out = annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + # matched the fixture (keyed uppercase) despite the lowercase cell + assert out.iloc[0]["aff"] == 100.0 + # the input peptide column is preserved verbatim, not uppercased + assert out.iloc[0]["peptide"] == "siinfekl" + + +def test_peptide_matched_ignoring_surrounding_whitespace(): + df = pd.DataFrame({"peptide": [" SIINFEKL "], "hla": ["HLA-A*02:01"]}) + out = annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + assert out.iloc[0]["aff"] == 100.0 + assert out.iloc[0]["peptide"] == " SIINFEKL " + + +# --- allele-free predictor path (by-peptide lookup) ------------------------- + +def test_allele_free_predictor_uses_by_peptide_lookup(): + df = pd.DataFrame({"peptide": ["SIINFEKL", "GILGFVFTL", "WWWWWWWWW"]}) + out = annotate_table( + df, + [AnnotationSpec(lambda alleles: _ProcessingFixturePredictor(), + "proc", field="score")]) + # no allele column -> no provenance column + assert "proc" in out.columns + assert "proc_best_allele" not in out.columns + assert out.iloc[0]["proc"] == 0.80 + assert out.iloc[1]["proc"] == 0.30 + # peptide absent from the fixture -> NaN + assert math.isnan(out.iloc[2]["proc"]) + + +def test_allele_free_predictor_works_even_with_allele_column_present(): + # An allele-free predictor emits allele-less predictions; when a (peptide, + # allele) lookup finds nothing, we fall back to the by-peptide prediction, + # so the score is still filled in rather than silently NaN. + df = pd.DataFrame({"peptide": ["SIINFEKL"], "hla": ["HLA-A*02:01"]}) + out = annotate_table( + df, + [AnnotationSpec(lambda alleles: _ProcessingFixturePredictor(), + "proc", field="score")], + allele_column="hla") + assert out.iloc[0]["proc"] == 0.80 + # the winning prediction had no allele -> provenance is None + assert out.iloc[0]["proc_best_allele"] is None + + +def test_binding_predictor_unsupported_allele_still_nan(): + # Regression guard: the by-peptide fallback must NOT rescue a binding + # predictor's unsupported-allele miss (it never populates by_peptide). + df = pd.DataFrame({"peptide": ["SIINFEKL"], "hla": ["HLA-C*07:01"]}) + out = annotate_table( + df, [AnnotationSpec(_factory, "aff", field="affinity")], + allele_column="hla") + assert math.isnan(out.iloc[0]["aff"]) + assert out.iloc[0]["aff_best_allele"] is None + + # --- spec parsing ----------------------------------------------------------- def test_parse_spec_full():