diff --git a/scripts/eval/classify_altloc_regions.py b/scripts/eval/classify_altloc_regions.py index 0d1de1de..af29cda7 100644 --- a/scripts/eval/classify_altloc_regions.py +++ b/scripts/eval/classify_altloc_regions.py @@ -5,17 +5,18 @@ 1. ``side_chain_only`` : altloc atoms exist in the residue, but none of its backbone atoms have altlocs. -2. ``small_loop`` : a contiguous backbone altloc span whose mean per-residue - backbone lDDT score (defined below) between altlocs is above - ``--loop-lddt-threshold`` (default 0.75). -3. ``large_loop`` : a contiguous backbone-altloc span whose mean per-residue - backbone lDDT score between altlocs is below ``--loop-lddt-threshold``. +2. ``small_loop`` : a contiguous backbone-altloc span whose loop score is above the threshold. +3. ``large_loop`` : a contiguous backbone-altloc span whose loop score is below the threshold. 4. ``domain_shift`` : a single contiguous backbone-altloc span longer than - ``--domain-shift-min-span`` residues (default 50). Classified before the - loop lDDT test. + ``--domain-shift-min-span`` residues (default 50). Classified before loop + scoring. -Score definition (important, slightly different from canonical lDDT): +Loop scores are calculated for every combination of altloc pairs when more than +two altlocs are present. The selected scoring strategy defines how per-residue +scores are reduced into a pair score, how pair scores are reduced into the final +loop score, and how that loop score is thresholded. +LDDT SCORER: For a given pair of altlocs, the score is the **equal-weighted arithmetic mean** of per residue backbone lDDT scores across the span: @@ -32,19 +33,31 @@ has the same neighbor count. The 0.75 default is calibrated for this specific calculation. -Altloc pairing: when > 2 altlocs are present, the score above is +RMSD SCORER: + For a given pair of altlocs, the score is the maximum of per residue + RMSD scores across the span: + + score = max(score_k) + + Each ``score_k`` is the standard all-atom RMSD from + :class:`sampleworks.metrics.rmsd.AllAtomRMSD`. + +`Altloc pairing`: when > 2 altlocs are present, the scores above are computed for every combination of altloc pairs and the span is classified by the *minimum* score over pair combinations. -Use ``find_altloc_selections.py --min-span 1`` to ensure single-residue side only +Use ``find_altloc_selections.py --min-span 1`` to ensure single-residue side chain only selections. """ import argparse import json +import operator import re -from collections.abc import Mapping +from collections.abc import Callable, Mapping +from dataclasses import dataclass from pathlib import Path +from typing import Any import numpy as np import pandas as pd @@ -55,8 +68,10 @@ ATOMWORKS_COMPARISON_OPS, get_mask_from_old_selection_string, parse_selection_string, + selection_to_residues, ) from sampleworks.metrics.lddt import AllAtomLDDT +from sampleworks.metrics.rmsd import AllAtomRMSD from sampleworks.utils.atom_array_utils import ( BACKBONE_ATOM_TYPES, BLANK_ALTLOC_IDS, @@ -77,10 +92,11 @@ "end_res", "span_length", "classification", - "worst_pair_mean_backbone_lddt", + "score", + "score_metric", "n_backbone_altloc_residues", "n_altlocs", - "pair_lddts", + "pair_scores", ] @@ -111,40 +127,85 @@ def _chain_from_selection(selection: str) -> str | None: return chain_id -def _mean_residue_lddt_for_pair( - gt_array: AtomArray | AtomArrayStack | None, - pred_array: AtomArray | AtomArrayStack | None, +@dataclass(frozen=True) +class SpanScorer: + """Metric specific behavior for classification.""" + + metric_name: str + metric_compute: Callable[..., dict[str, Any]] + residue_scores_key: str + pair_score_reducer: Callable[[list[float]], float] + worse_pair_score: Callable[[list[float]], float] + is_small_loop_score: Callable[[float, float], bool] + default_threshold: float + + +def _score_pair_with_scorer( + scorer: SpanScorer, + gt_array: AtomArrayStack | AtomArray | None, + pred_array: AtomArrayStack | AtomArray | None, chain: str, residues: list[int], ) -> float: - """Equal weighted arithmetic mean of per residue lDDT across the span.""" + """Calculate a scorer's pair score over a residue span.""" if gt_array is None or pred_array is None or not residues: return float("nan") res_clause = " or ".join(f"res_id == {r}" for r in residues) selection = f"chain_id == '{chain}' and ({res_clause}) and atom_name in ['C','CA','N','O']" try: - result = AllAtomLDDT().compute( + result = scorer.metric_compute( predicted_atom_array_stack=pred_array, ground_truth_atom_array_stack=gt_array, selection=selection, ) except Exception as e: - logger.warning(f"lDDT compute failed for chain {chain} residues {residues}: {e}") + logger.warning( + f"{scorer.metric_name} computation failed for chain {chain} residues {residues}: {e}" + ) + return float("nan") + + residue_scores = result.get(scorer.residue_scores_key, {}) + if not isinstance(residue_scores, dict): + logger.warning( + f"{scorer.metric_name} result did not contain a residue score dictionary " + f"under key '{scorer.residue_scores_key}'" + ) return float("nan") - residue_scores = result.get("residue_lddt_scores", {}) keys = [f"{chain}{r}" for r in residues] missing = [k for k in keys if k not in residue_scores] if missing: logger.warning( - f"lDDT result missing residues {missing} for chain {chain}. This means the result" - f"averaged only over the {len(keys) - len(missing)} residues it returned" + f"{scorer.metric_name} result missing residues {missing} for chain {chain}. " + f"This means the result was reduced only over the " + f"{len(keys) - len(missing)} residues it returned" ) - flat = [residue_scores[k][0] for k in keys if k in residue_scores] - if not flat: - return float("nan") - return float(np.mean(flat)) + + flat = [float(residue_scores[k][0]) for k in keys if k in residue_scores and residue_scores[k]] + return float(scorer.pair_score_reducer(flat)) if flat else float("nan") + + +LOOP_SCORERS = { + "lddt": SpanScorer( + metric_name="lddt", + metric_compute=AllAtomLDDT().compute, + residue_scores_key="residue_lddt_scores", + pair_score_reducer=np.mean, + worse_pair_score=min, + is_small_loop_score=operator.gt, + default_threshold=0.75, + ), + "rmsd": SpanScorer( + metric_name="rmsd", + metric_compute=AllAtomRMSD(superimpose=False).compute, + residue_scores_key="residue_rmsd_scores", + pair_score_reducer=max, + worse_pair_score=max, + is_small_loop_score=operator.lt, + default_threshold=1.0, + ), +} def _classify_selection( @@ -158,16 +219,17 @@ def _classify_selection( structure_altloc_mask: np.ndarray, structure_backbone_mask: np.ndarray, domain_shift_min_span: int, - loop_lddt_threshold: float, + scorer: SpanScorer, + loop_score_threshold: float, ) -> tuple[dict, set[tuple[str, int]]] | None: """Classify one contiguous altloc selection into a conformational type. 1. If the span has no backbone altlocs anywhere, it is classified as ``side_chain_only``. 2. Else if the longest contiguous backbone altloc run exceeds ``domain_shift_min_span``, it is classified as ``domain_shift``. - 3. Else compute the per residue backbone lDDT for every altloc pair over + 3. Else compute the per residue metric for every altloc pair over the backbone altloc residues in the span and take the minimum - pair mean. Compare against ``loop_lddt_threshold``, if it is above is is classified as + pair mean. Compare against ``loop_score_threshold``, if it is above is is classified as ``small_loop``. If it is below, it is classified as ``large_loop``. Returns ``(row_dict, covered_altloc_residues)`` on success or ``None`` if the @@ -175,9 +237,9 @@ def _classify_selection( ``row_dict`` has the keys: ``protein``, ``selection``, ``chain``, ``start_res``, ``end_res``, - ``span_length``, ``classification``, ``worst_pair_mean_backbone_lddt``, - ``n_backbone_altloc_residues``, ``n_altlocs``, and ``pair_lddts`` (a - JSON encoded ``{pair_label: mean_lddt}`` map so the dict can be loaded + ``span_length``, ``classification``, ``score``, ``score_metric``, + ``n_backbone_altloc_residues``, ``n_altlocs``, and ``pair_scores`` (a + JSON encoded ``{pair_label: score}`` map so the dict can be loaded through the CSV intact via ``json.loads``). ``covered_altloc_residues`` is the set of ``(chain_id, res_id)`` pairs in the @@ -243,8 +305,9 @@ def _classify_selection( "n_backbone_altloc_residues": n_backbone, "n_altlocs": len(altloc_ids), # JSON encoded so the pair calculation can be loaded back through the CSV - "pair_lddts": json.dumps({}), - "worst_pair_mean_backbone_lddt": float("nan"), + "pair_scores": json.dumps({}), + "score_metric": scorer.metric_name, + "score": float("nan"), "classification": "", } @@ -258,28 +321,31 @@ def _classify_selection( row["classification"] = "domain_shift" return row, covered_altloc_residues - # Loop classification via pairwise lDDT across all altloc pairs - pair_lddts: dict[str, float] = {} + # Loop classification via pairwise strategy scores across all altloc pairs + pair_scores: dict[str, float] = {} for i in range(len(altloc_ids)): for j in range(i + 1, len(altloc_ids)): pair = pair_arrays.get((altloc_ids[i], altloc_ids[j])) gt, pred = pair if pair is not None else (None, None) - pair_lddts[f"{altloc_ids[i]}-{altloc_ids[j]}"] = _mean_residue_lddt_for_pair( - gt, pred, chain, backbone_altloc_res_ids + pair_scores[f"{altloc_ids[i]}-{altloc_ids[j]}"] = _score_pair_with_scorer( + scorer, gt, pred, chain, backbone_altloc_res_ids ) - row["pair_lddts"] = json.dumps(pair_lddts) + row["pair_scores"] = json.dumps(pair_scores) - finite_vals = [v for v in pair_lddts.values() if np.isfinite(v)] + finite_vals = [v for v in pair_scores.values() if np.isfinite(v)] if not finite_vals: raise RuntimeError( - f"[{protein}] could not compute lDDT for any altloc pair in span " - f"'{selection_str}' (backbone-altloc residues: {backbone_altloc_res_ids}). " + f"[{protein}] could not compute {scorer.metric_name} for any altloc pair " + f"in span '{selection_str}' " + f"(backbone-altloc residues: {backbone_altloc_res_ids}). " "Refusing to emit an indeterminate classification." ) - worst = float(min(finite_vals)) - row["worst_pair_mean_backbone_lddt"] = worst - row["classification"] = "small_loop" if worst > loop_lddt_threshold else "large_loop" + worst = float(scorer.worse_pair_score(finite_vals)) + row["score"] = worst + row["classification"] = ( + "small_loop" if scorer.is_small_loop_score(worst, loop_score_threshold) else "large_loop" + ) return row, covered_altloc_residues @@ -287,7 +353,8 @@ def _process_structure( input_row: pd.Series, cif_root: Path | None, domain_shift_min_span: int, - loop_lddt_threshold: float, + scoring_strategy: SpanScorer, + loop_score_threshold: float, ) -> list[dict]: protein = str(input_row["protein"]) cif_path = resolve_cif_path(input_row, cif_root) @@ -314,15 +381,25 @@ def _process_structure( structure_altloc_mask = ~np.isin(atom_array.altloc_id, list(BLANK_ALTLOC_IDS)) structure_backbone_mask = np.isin(atom_array.atom_name, BACKBONE_ATOM_TYPES) + # (chain, res_id) pairs that carry any altloc. Used to skip the redundant combined + # selection below and to check coverage after classification + all_altloc_res_ids: set[tuple[str, int]] = { + (str(c), int(r)) + for c, r in zip( + atom_array.chain_id[structure_altloc_mask], + atom_array.res_id[structure_altloc_mask], + ) + } + rows: list[dict] = [] classified_res_ids: set[tuple[str, int]] = set() for selection_str in [s.strip() for s in selection_field.split(";") if s.strip()]: - # find_altloc_selections.py appends a combined all altloc selection - # (atomworks-style with " or " clauses) at the end of each row. That one is - # a union over every span we already processed individually, so skip it. - # NOTE: This will need to be addressed when we - # migrate to atomworks-style selections for everything - if " or " in selection_str: + # find_altloc_selections.py may append a combined selection unioning every altloc + # span (atomworks-style "res_id == .. or .."). Skip only that redundant total union + # to avoid double-counting. True discontinuous "or" selections are still + # classified + covered_res_ids = selection_to_residues(atom_array, selection_str) + if covered_res_ids and covered_res_ids == all_altloc_res_ids: continue out = _classify_selection( atom_array=atom_array, @@ -333,7 +410,8 @@ def _process_structure( structure_altloc_mask=structure_altloc_mask, structure_backbone_mask=structure_backbone_mask, domain_shift_min_span=domain_shift_min_span, - loop_lddt_threshold=loop_lddt_threshold, + scorer=scoring_strategy, + loop_score_threshold=loop_score_threshold, ) if out is None: continue @@ -343,13 +421,6 @@ def _process_structure( # residues across all classified spans should equal total unique # (chain, res_id) pairs that carry any altloc in the structure. - all_altloc_res_ids: set[tuple[str, int]] = { - (str(c), int(r)) - for c, r in zip( - atom_array.chain_id[structure_altloc_mask], - atom_array.res_id[structure_altloc_mask], - ) - } if classified_res_ids != all_altloc_res_ids: missing = all_altloc_res_ids - classified_res_ids extra = classified_res_ids - all_altloc_res_ids @@ -364,6 +435,14 @@ def _process_structure( def main(args: argparse.Namespace) -> None: + """Run altloc region classification.""" + scoring_strategy = LOOP_SCORERS[args.loop_score_metric] + loop_score_threshold = ( + args.loop_score_threshold + if args.loop_score_threshold is not None + else scoring_strategy.default_threshold + ) + input_df = pd.read_csv(args.input_csv) required = {"protein", "selection"} missing = required - set(input_df.columns) @@ -377,7 +456,8 @@ def main(args: argparse.Namespace) -> None: input_row=row, cif_root=args.cif_root, domain_shift_min_span=args.domain_shift_min_span, - loop_lddt_threshold=args.loop_lddt_threshold, + scoring_strategy=scoring_strategy, + loop_score_threshold=loop_score_threshold, ) ) @@ -411,6 +491,20 @@ def main(args: argparse.Namespace) -> None: ) parser.add_argument("--output-file", type=Path, required=True) parser.add_argument("--domain-shift-min-span", type=int, default=50) - parser.add_argument("--loop-lddt-threshold", type=float, default=0.75) + parser.add_argument( + "--loop-score-metric", + choices=sorted(LOOP_SCORERS), + default="lddt", + help="Residue-level metric strategy used for small_loop / large_loop classification.", + ) + parser.add_argument( + "--loop-score-threshold", + type=float, + default=None, + help=( + "Threshold for the selected loop scoring strategy. Defaults to the scorer specific" + "threshold when omitted." + ), + ) args = parser.parse_args() main(args)