diff --git a/fitness/README.md b/fitness/README.md index 4acdd8a..5eca058 100644 --- a/fitness/README.md +++ b/fitness/README.md @@ -34,9 +34,9 @@ The following directories are gitignored. Download them from ## Metrics -- **Spearman** — absolute Spearman correlation between predicted and experimental scores -- **AUC** — AUROC using median-threshold binarisation (max of AUC and 1−AUC) -- **MCC** — absolute Matthews correlation coefficient (median-threshold binarisation) +- **Spearman** — signed Spearman correlation between predicted and experimental scores. A negative value means the model ranks variants the wrong way round, which the metric reports rather than hides +- **AUC** — AUROC of the model's continuous scores against assay labels binarised at their median. Below 0.5 means the model ranks variants the wrong way round +- **MCC** — Matthews correlation coefficient, with both the assay labels and the model's predictions binarised at their medians, signed ## Usage diff --git a/fitness/analyze_fill_strategies.py b/fitness/analyze_fill_strategies.py new file mode 100644 index 0000000..fa36fee --- /dev/null +++ b/fitness/analyze_fill_strategies.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +""" +Compare the masked-marginal fill strategies across checkpoints. + +Reads the prediction files written by fitness/baselines/masked_lm, one per assay +per checkpoint, each holding every computed strategy as its own column, and +reports the benchmark metric: signed Spearman per assay, averaged within each +ncRNA category, then a macro mean over the categories so that the 26 ribozyme +assays do not swamp the 3 tRNA and 2 aptamer assays. + +Also reports where the strategies differ, whether they reorder the checkpoints, +how the comparison responds to dropping the category weighting, and a paired +bootstrap over assays. This is the code behind the sensitivity section of +leaderboard/fitness/README.md. + +The folders and columns come from merge_scoring_files.SCORE_COLS, so the +registry is not duplicated here. + +Example: + python fitness/analyze_fill_strategies.py \\ + --predictions_folder path/to/model_predictions \\ + --ref_sheet fitness/reference_sheet_final.csv \\ + --output_folder leaderboard/fitness +""" + +import argparse +import os +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy import stats + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from fitness.merge_scoring_files import ( # noqa: E402 + FOUR_FILL_MODELS, + SCORE_COLS, + resolve_source, +) + +NCRNA = ["Ribozyme", "tRNA", "Aptamer"] +STRATEGIES = ["wt_fill", "mask_fill", "mut_fill", "match_fill"] + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--predictions_folder", required=True, + help="Folder holding one subfolder per prediction set") + parser.add_argument("--ref_sheet", required=True, + help="Reference sheet with DMS_ID and RNA_TYPE") + parser.add_argument("--output_folder", default=None, + help="Where to write the per-assay and macro CSVs (default: no CSVs)") + parser.add_argument("--bootstrap_draws", type=int, default=2000, + help="Paired bootstrap draws, 0 to skip (default: 2000)") + parser.add_argument("--seed", type=int, default=0, help="Bootstrap seed (default: 0)") + parser.add_argument("--allow_incomplete", action="store_true", + help="Report on checkpoints missing some assays instead of skipping them") + return parser.parse_args() + + +def checkpoints(): + """Group the registry's four-fill entries by checkpoint.""" + grouped = {} + for entry in FOUR_FILL_MODELS: + for strategy in STRATEGIES: + suffix = f"_{strategy}" + if entry.endswith(suffix): + grouped.setdefault(entry[: -len(suffix)], {})[strategy] = entry + break + return grouped + + +def collect(args, expected, rna_type): + """Per-assay signed Spearman for every checkpoint and strategy.""" + records = [] + for name, entries in checkpoints().items(): + columns, folders = {}, set() + for strategy, entry in entries.items(): + folder, columns[strategy] = resolve_source(SCORE_COLS, entry) + folders.add(folder) + if len(folders) != 1: + raise ValueError(f"{name}: its strategies are registered across {folders}") + directory = os.path.join(args.predictions_folder, folders.pop()) + if not os.path.isdir(directory): + continue + found = sorted(f[:-4] for f in os.listdir(directory) if f.endswith(".csv")) + unexpected = sorted(set(found) - set(expected)) + if unexpected: + raise ValueError(f"{name}: unexpected assays in {directory}: {unexpected}") + if set(found) != set(expected): + missing = sorted(set(expected) - set(found)) + print(f"{name:18s} incomplete: {len(found)}/{len(expected)} assays, " + f"missing {missing[:4]}") + if not args.allow_incomplete: + continue + for assay in found: + frame = pd.read_csv(os.path.join(directory, f"{assay}.csv")) + absent = [s for s, c in columns.items() if c not in frame.columns] + if absent: + raise KeyError(f"{directory}/{assay}.csv has no columns for {absent}") + # Every strategy must cover the same variants, or the comparison + # between them is confounded by coverage rather than by method. + masks = [frame[columns[s]].notna().to_numpy() for s in STRATEGIES] + if not all((m == masks[0]).all() for m in masks[1:]): + raise ValueError(f"{directory}/{assay}.csv: strategies cover different variants") + for strategy in STRATEGIES: + usable = frame[["DMS_score", columns[strategy]]].dropna() + if len(usable) < 3: + raise ValueError(f"{assay}: only {len(usable)} usable variants") + rho = stats.spearmanr(usable["DMS_score"], usable[columns[strategy]]).correlation + if not np.isfinite(rho): + raise ValueError(f"{assay}: Spearman undefined for {strategy}") + records.append({"model": name, "strategy": strategy, "assay": assay, + "RNA_TYPE": rna_type[assay], "n": len(usable), "spearman": rho}) + return pd.DataFrame(records) + + +def macro_table(per_assay): + rows = [] + for (model, strategy), group in per_assay.groupby(["model", "strategy"]): + by_type = {t: group[group.RNA_TYPE == t]["spearman"].mean() for t in NCRNA} + rows.append({"model": model, "strategy": strategy, **by_type, + "macro_3ncRNA": float(np.mean([by_type[t] for t in NCRNA]))}) + return pd.DataFrame(rows) + + +def bootstrap(per_assay, order, draws, seed): + """ + Paired bootstrap over assays, resampled within each category. + + The three categories are not resampled: they define the benchmark's estimand + rather than sampling from a population. Each assay keeps its strategies + together, so the comparison is paired. + """ + rng = np.random.default_rng(seed) + paired = per_assay.pivot_table(index=["model", "assay", "RNA_TYPE"], + columns="strategy", values="spearman").reset_index() + print(f"\n=== paired bootstrap, {draws} draws, seed {seed} ===") + header = " ".join(f"{'wt-fill minus ' + s.replace('_', '-'):>26}" + for s in ["mut_fill", "mask_fill", "match_fill"]) + print(f"{'checkpoint':18s} {header}") + for model in order: + sub = paired[paired.model == model] + by_type = {t: sub[sub.RNA_TYPE == t] for t in NCRNA} + samples = {s: np.empty(draws) for s in STRATEGIES} + for d in range(draws): + drawn = {t: frame.iloc[rng.integers(0, len(frame), len(frame))] + for t, frame in by_type.items()} + for strategy in STRATEGIES: + samples[strategy][d] = np.mean([drawn[t][strategy].mean() for t in NCRNA]) + cells = [] + for other in ["mut_fill", "mask_fill", "match_fill"]: + delta = samples["wt_fill"] - samples[other] + lo, hi = np.percentile(delta, [2.5, 97.5]) + cells.append(f"{delta.mean():+.4f} [{lo:+.4f},{hi:+.4f}]" + ("*" if lo > 0 or hi < 0 else " ")) + print(f"{model:18s} " + " ".join(f"{c:>26}" for c in cells)) + print(" * = the 95% interval excludes zero. The tRNA and aptamer categories hold") + print(" 3 and 2 assays, so these intervals are wide by construction.") + + +def main(): + args = parse_args() + ref = pd.read_csv(args.ref_sheet, encoding="utf-8-sig") + rna_type = dict(zip(ref["DMS_ID"], ref["RNA_TYPE"])) + expected = sorted(d for d, t in rna_type.items() if t in NCRNA) + print(f"{len(expected)} ncRNA assays expected per checkpoint") + + per_assay = collect(args, expected, rna_type) + if per_assay.empty: + raise SystemExit("No complete checkpoint found in the predictions folder") + macro = macro_table(per_assay) + wide = macro.pivot(index="model", columns="strategy", values="macro_3ncRNA")[STRATEGIES] + order = list(wide.sort_values("wt_fill", ascending=False).index) + + print("\n=== signed Spearman, macro over the 3 ncRNA categories ===") + print(wide.loc[order].round(4).to_string()) + + print("\n=== ordering under each strategy, best first ===") + for strategy in STRATEGIES: + print(f" {strategy:10s} " + " > ".join(wide[strategy].sort_values(ascending=False).index)) + + print("\n=== where the strategies differ ===") + spread = pd.DataFrame([ + {"model": m, **{c: macro[macro.model == m].set_index("strategy")[c].pipe( + lambda s: s.max() - s.min()) for c in NCRNA}} for m in order]).set_index("model") + print(spread.round(4).to_string()) + print(" mean spread: " + ", ".join(f"{c} {spread[c].mean():.4f}" for c in NCRNA)) + + print("\n=== without the category weighting ===") + flat = per_assay.groupby(["model", "strategy"])["spearman"].mean().unstack()[STRATEGIES] + print(flat.loc[order].round(4).to_string()) + print(f" wt-fill best on {(wide.loc[order].idxmax(axis=1) == 'wt_fill').sum()}/{len(order)} " + f"under the macro metric, {(flat.loc[order].idxmax(axis=1) == 'wt_fill').sum()}/{len(order)} " + "under a flat mean over all assays") + + if args.bootstrap_draws: + bootstrap(per_assay, order, args.bootstrap_draws, args.seed) + + if args.output_folder: + out = Path(args.output_folder) + out.mkdir(parents=True, exist_ok=True) + per_assay.to_csv(out / "fill_strategy_per_assay.csv", index=False) + macro.to_csv(out / "fill_strategy_macro.csv", index=False) + print(f"\nwrote {out}/fill_strategy_per_assay.csv and {out}/fill_strategy_macro.csv") + + +if __name__ == "__main__": + main() diff --git a/fitness/baselines/AIDO_RNA/score_aido_rna.sh b/fitness/baselines/AIDO_RNA/score_aido_rna.sh index e693982..1caaf57 100644 --- a/fitness/baselines/AIDO_RNA/score_aido_rna.sh +++ b/fitness/baselines/AIDO_RNA/score_aido_rna.sh @@ -8,10 +8,11 @@ export model_name="genbio-ai/AIDO.RNA-1.6B" export reference_sheet="reference_sheet.csv" -# Write predictions under a folder named "aido_rna" so they line up with the -# "aido_rna" entry in fitness/merge_scoring_files.py (which reads the -# aido_rna_score column from model_predictions/aido_rna/). -export output_scores_dir="path/to/model_predictions/aido_rna" +# Write predictions under a folder named "aido_rna_4fill". One run writes all +# four fill strategies into it as aido_rna_score_wt_fill and so on, which the +# aido_rna_wt_fill, aido_rna_mask_fill, aido_rna_mut_fill and aido_rna_match_fill +# entries in fitness/merge_scoring_files.py read. +export output_scores_dir="path/to/model_predictions/aido_rna_4fill" export dms_data_dir="path/to/dms/data/dir" # Reference-sheet row to score. Set by a Slurm array job (0-69), or defaults @@ -19,6 +20,15 @@ export dms_data_dir="path/to/dms/data/dir" # tRNA, aptamer); scoring only those leaves aido_rna without mRNA predictions, # so read its aggregate from performance_fitness.py --type ncRNA. Under # --type all its All_Mean is NaN by design. + +# Masked-marginal fill strategy. The default computes all four (wt-fill, +# mask-fill, mut-fill, match-fill), which share their contexts and so cost only +# about 19% more unique context examples than mut-fill alone, and writes one +# column per strategy named +# {COLUMN}_{strategy}. Pass --strategies mut-fill (or any single strategy) to +# write the historical bare {COLUMN} column instead. See +# fitness/baselines/masked_lm/strategies.py for the formulas. + DMS_index=${SLURM_ARRAY_TASK_ID:-0} python score_aido_rna_single_dms.py \ diff --git a/fitness/baselines/AIDO_RNA/score_aido_rna_single_dms.py b/fitness/baselines/AIDO_RNA/score_aido_rna_single_dms.py index 3dd38ea..c4de072 100644 --- a/fitness/baselines/AIDO_RNA/score_aido_rna_single_dms.py +++ b/fitness/baselines/AIDO_RNA/score_aido_rna_single_dms.py @@ -1,363 +1,72 @@ #!/usr/bin/env python3 """ -Script to run AIDO.RNA inference on DMS assay sequences. -Takes a reference sheet and row ID to process specific assays. +Score DMS assay sequences with AIDO.RNA. AIDO.RNA (https://huggingface.co/genbio-ai/AIDO.RNA-1.6B) is an encoder-only transformer pretrained with a masked language modelling objective on 42M -non-coding RNA sequences from RNAcentral. We score variants with the -masked-marginal log-likelihood ratio, the same zero-shot proxy used by the other -masked RNA language-model baselines (RiNALMo, Orthrus): - - score(variant) = sum_i [ log P(mut_i | variant context, pos_i masked) - - log P(wt_i | variant context, pos_i masked) ] - -summed over the variant's mutated positions ``i``. Each position is masked in -the variant's own sequence, so for a multi-mutant the remaining mutations stay -in the context. - -Masking one position of one variant is a single forward pass, but two variants -that differ only at the masked position share the same masked context. Those -contexts are deduplicated before inference, which is exact and cuts the number -of forward passes by up to 3x on single-substitution libraries. +non-coding RNA sequences from RNAcentral. It is scored with the shared +masked-marginal engine in ``fitness/baselines/masked_lm``, which offers the four +fill strategies (``wt-fill``, ``mask-fill``, ``mut-fill``, ``match-fill``) that +differ in what the model sees at a variant's other mutated positions. The model code is the official implementation released by GenBio AI in the ``modelgenerator`` package (``pip install --no-deps modelgenerator``); only torch and transformers are needed on top of it. """ -import argparse import os import sys from pathlib import Path -import numpy as np -import pandas as pd import torch -from scipy.stats import spearmanr -from tqdm.auto import tqdm - -from modelgenerator.huggingface_models.rnabert import ( - RNABertForMaskedLM, - RNABertTokenizer, -) - -# AIDO.RNA accepts both alphabets and mirrors whichever one the context uses: on -# T-form input it puts ~0 probability on U and vice versa. Masked negative -# log-likelihood on the wild-type ncRNA constructs is consistently lower in -# T-form (e.g. 0.2197 vs 0.2223 on Domingo_2018_tRNA), so U is folded to T. -BASES = "ACGT" -MASK_CHAR = "#" - - -def preprocess_sequence(sequence: str) -> str: - """ - Preprocess an RNA/DNA sequence for AIDO.RNA: - - Convert to uppercase - - Convert RNA (U) to DNA (T) - - Remove any whitespace - """ - return sequence.strip().upper().replace("U", "T") - - -def parse_args(): - """Parse command line arguments.""" - parser = argparse.ArgumentParser( - description="Run AIDO.RNA (masked LM) inference on DMS assay sequences." - ) - parser.add_argument( - "--row_id", - type=int, - required=True, - help="Row ID in the reference sheet to process", - ) - parser.add_argument( - "--ref_sheet", - type=str, - required=True, - help="Path to reference sheet containing DMS_ID column", - ) - parser.add_argument( - "--dms_dir_path", - type=str, - required=True, - help="Directory containing DMS CSV files", - ) - parser.add_argument( - "--output_dir_path", - type=str, - required=True, - help="Directory to save output files", - ) - parser.add_argument( - "--device", - type=str, - default="cuda:0" if torch.cuda.is_available() else "cpu", - help="Device to run inference on (default: cuda:0 if available, else cpu)", - ) - parser.add_argument( - "--model_name", - type=str, - default="genbio-ai/AIDO.RNA-1.6B", - help="AIDO.RNA model to use (default: genbio-ai/AIDO.RNA-1.6B)", - ) - parser.add_argument( - "--dtype", - type=str, - default="bfloat16", - choices=["bfloat16", "float32"], - help="Weight/compute dtype. bfloat16 is ~2x faster and was used for the " - "released scores; float32 gives finer logits and fewer tied variants " - "(default: bfloat16)", - ) - parser.add_argument( - "--batch_size", - type=int, - default=512, - help="Maximum masked contexts scored per forward pass (default: 512)", - ) - parser.add_argument( - "--max_batch_tokens", - type=int, - default=49152, - help="Cap on batch_size x sequence length per forward pass, so that long " - "assays automatically use a smaller batch (default: 49152, which peaks " - "around 10 GB of GPU memory for the 1.6B model)", - ) - return parser.parse_args() - - -def load_reference_data(ref_sheet_path: str, row_id: int) -> str: - """ - Load reference sheet and get DMS_ID for specified row. - - Raises: - ValueError: If row_id is not found or DMS_ID is missing - """ - try: - ref_df = pd.read_csv(ref_sheet_path) - if row_id >= len(ref_df): - raise ValueError( - f"Row ID {row_id} exceeds number of rows in reference sheet" - ) - - dms_id = ref_df.loc[row_id, "DMS_ID"] - if pd.isna(dms_id): - raise ValueError(f"DMS_ID is missing for row {row_id}") - - return str(dms_id) - - except FileNotFoundError: - raise FileNotFoundError(f"Reference sheet not found: {ref_sheet_path}") - except KeyError: - raise KeyError("Reference sheet must contain 'DMS_ID' column") - - -def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: - """ - Load DMS data for specified DMS_ID. - - Raises: - FileNotFoundError: If DMS file is not found - """ - dms_file = Path(dms_dir_path) / f"{dms_id}.csv" - if not dms_file.exists(): - raise FileNotFoundError(f"DMS file not found: {dms_file}") - - df = pd.read_csv(dms_file) - required_cols = ["mutant", "DMS_score", "sequence"] - missing_cols = [col for col in required_cols if col not in df.columns] - if missing_cols: - raise ValueError(f"Missing required columns in DMS file: {missing_cols}") - - return df - - -def parse_mutations(mutant_str: str) -> list: - """ - Parse a mutation string such as ``"A4T,A5G"`` into a list of - ``(pos0, wt_base, mut_base)`` tuples with 0-based positions, in the DNA - alphabet used by the model. - - Raises: - ValueError: for non-substitution edits (indels) or unknown bases, so - the caller can score the affected variant as NaN. - """ - mutations = [] - for token in str(mutant_str).replace(" ", "").split(","): - if not token: - continue - wt_base = token[0].upper().replace("U", "T") - mut_base = token[-1].upper().replace("U", "T") - pos = int(token[1:-1]) - 1 # 1-based -> 0-based - if wt_base not in BASES or mut_base not in BASES: - raise ValueError(f"Unsupported mutation token: {token}") - mutations.append((pos, wt_base, mut_base)) - return mutations +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -def build_masking_tasks(mutants: list, sequences: list) -> tuple: - """ - Expand each variant into one masked-scoring task per mutated position, and - deduplicate identical masked contexts. +from masked_lm import MaskedLMAdapter, main # noqa: E402 - A masked context is the variant's own sequence with the scored position - replaced by a mask placeholder. Two variants that differ only at that - position produce the same context and therefore the same log-probabilities, - so the forward pass is shared. The context also determines the masked - position, so it does not need to be tracked separately. - Returns: - contexts: list of masked context strings, one per forward pass - ctx_tasks: list parallel to ``contexts``; entry k holds the - ``(row_idx, wt_base, mut_base)`` tuples scored from context k - scores: per-variant score array, pre-filled with 0.0 for scorable - variants and NaN for wild-type / unparseable / out-of-range rows +class AIDORNAAdapter(MaskedLMAdapter): """ - scores = np.full(len(sequences), np.nan, dtype=float) - ctx_index = {} - contexts = [] - ctx_tasks = [] - - for i, (mutant_str, seq) in enumerate(zip(mutants, sequences)): - if pd.isna(mutant_str): - continue # wild-type row: leave as NaN - try: - mutations = parse_mutations(mutant_str) - if not mutations: - continue - row_tasks = [] - for pos, wt_base, mut_base in mutations: - if pos < 0 or pos >= len(seq): - raise ValueError(f"Mutation position {pos + 1} outside sequence") - if seq[pos] != mut_base: - raise ValueError( - f"Sequence has {seq[pos]} at position {pos + 1}, " - f"expected the mutant base {mut_base}" - ) - row_tasks.append((seq[:pos] + MASK_CHAR + seq[pos + 1 :], wt_base, mut_base)) - except (ValueError, IndexError) as err: - print(f"Skipping variant {mutant_str}: {err}") - continue # unsupported edit: leave as NaN - for context, wt_base, mut_base in row_tasks: - k = ctx_index.get(context) - if k is None: - k = len(contexts) - ctx_index[context] = k - contexts.append(context) - ctx_tasks.append([]) - ctx_tasks[k].append((i, wt_base, mut_base)) - scores[i] = 0.0 # scorable: accumulate per-position deltas below + AIDO.RNA: DNA alphabet, ``[CLS]`` and ``[SEP]``, attention mask. - return contexts, ctx_tasks, scores - - -def run_inference( - model, - tokenizer, - contexts: list, - ctx_tasks: list, - scores: np.ndarray, - device: str, - batch_size: int, - max_batch_tokens: int, -) -> np.ndarray: + AIDO.RNA accepts both alphabets and mirrors whichever one the context uses: + on T-form input it puts about zero probability on U and vice versa. Masked + negative log-likelihood on the wild-type ncRNA constructs is consistently + lower in T-form (0.2197 vs 0.2223 on Domingo_2018_tRNA), so U is folded to T. """ - Score masked contexts with AIDO.RNA's masked LM head and accumulate the - per-position log-likelihood ratios into each variant's score. - """ - base_ids = {b: tokenizer.convert_tokens_to_ids(b) for b in BASES} - if tokenizer.unk_token_id in base_ids.values(): - raise ValueError(f"Tokenizer does not cover the {BASES} alphabet: {base_ids}") - cls_id = tokenizer.cls_token_id - sep_id = tokenizer.sep_token_id - mask_id = tokenizer.mask_token_id - pad_id = tokenizer.pad_token_id - unk_id = tokenizer.unk_token_id # any base outside ACGT, e.g. an N in a construct - - # Sequence length is constant within an assay, so a single batch size is - # enough; still derive it from the length so long assays stay in memory. - seq_len = max(len(c) for c in contexts) - batch_size = max(1, min(batch_size, max_batch_tokens // (seq_len + 2))) - print(f"Using batch size {batch_size} for sequence length {seq_len}") - - for start in tqdm( - range(0, len(contexts), batch_size), desc="Scoring", unit="batch" - ): - batch = contexts[start : start + batch_size] - max_len = max(len(c) for c in batch) + 2 # [CLS] ... [SEP] - - input_ids = torch.full((len(batch), max_len), pad_id, dtype=torch.long) - attention_mask = torch.zeros((len(batch), max_len), dtype=torch.long) - mask_pos = [] - for b, context in enumerate(batch): - ids = [cls_id] - for ch in context: - ids.append(mask_id if ch == MASK_CHAR else base_ids.get(ch, unk_id)) - ids.append(sep_id) - input_ids[b, : len(ids)] = torch.tensor(ids, dtype=torch.long) - attention_mask[b, : len(ids)] = 1 - # Dedup is only exact if a context masks exactly one position, which is - # what makes the context string identify the position unambiguously. - if context.count(MASK_CHAR) != 1: - raise ValueError(f"Context does not mask exactly one position: {context}") - mask_pos.append(context.index(MASK_CHAR) + 1) # +1 for [CLS] - - input_ids = input_ids.to(device) - attention_mask = attention_mask.to(device) - rows = torch.arange(len(batch), device=device) - cols = torch.tensor(mask_pos, device=device) - - with torch.inference_mode(): - logits = model(input_ids=input_ids, attention_mask=attention_mask).logits - log_probs = torch.log_softmax(logits[rows, cols].float(), dim=-1) - log_probs = log_probs.cpu().numpy() - - for b in range(len(batch)): - row_log_probs = log_probs[b] - for row_idx, wt_base, mut_base in ctx_tasks[start + b]: - scores[row_idx] += ( - row_log_probs[base_ids[mut_base]] - row_log_probs[base_ids[wt_base]] - ) - - return scores - -def main(): - args = parse_args() - - # Create output directory if it doesn't exist - output_dir = Path(args.output_dir_path) - output_dir.mkdir(parents=True, exist_ok=True) - - try: - # Load DMS ID from reference sheet - dms_id = load_reference_data(args.ref_sheet, args.row_id) - print(f"Processing DMS ID: {dms_id}") - - # Load DMS data - dms_df = load_dms_data(args.dms_dir_path, dms_id) - - # Preprocess sequences (uppercase, RNA -> DNA) - print("Preprocessing sequences...") - sequences = [preprocess_sequence(seq) for seq in dms_df["sequence"].tolist()] - - # Expand variants into deduplicated masked contexts - contexts, ctx_tasks, scores = build_masking_tasks( - dms_df["mutant"].tolist(), sequences + name = "AIDO.RNA" + bases = "ACGT" + score_column = "aido_rna_score" + n_special_tokens = 2 # [CLS] and [SEP] + default_max_batch_tokens = 49152 + + @staticmethod + def add_arguments(parser): + parser.add_argument( + "--model_name", + type=str, + default="genbio-ai/AIDO.RNA-1.6B", + help="HuggingFace model id or local path of the AIDO.RNA checkpoint " + "(default: genbio-ai/AIDO.RNA-1.6B)", + ) + parser.add_argument( + "--dtype", + type=str, + default="bfloat16", + choices=["bfloat16", "float32"], + help="Torch dtype for the model weights. bfloat16 is the default and " + "is what the released predictions were produced with; fp32 changed " + "the aptamer Spearman by at most 0.0006 when it was checked", ) - n_tasks = sum(len(t) for t in ctx_tasks) - print( - f"Prepared {n_tasks} masked positions across " - f"{int(np.sum(~np.isnan(scores)))} scorable variants " - f"(of {len(sequences)} total), deduplicated to {len(contexts)} " - f"forward passes" + + def load(self, args): + from modelgenerator.huggingface_models.rnabert import ( + RNABertForMaskedLM, + RNABertTokenizer, ) - if not contexts: - raise ValueError("No scorable variants found") - # Initialize model. The tokenizer vocabulary ships with modelgenerator. - print(f"Initializing AIDO.RNA model ({args.model_name})...") + # The tokenizer vocabulary ships with modelgenerator. vocab_file = os.path.join( os.path.dirname(__import__("modelgenerator").__file__), "huggingface_models", @@ -365,53 +74,25 @@ def main(): "vocab.txt", ) tokenizer = RNABertTokenizer(vocab_file, version="v2") - model = RNABertForMaskedLM.from_pretrained( + self.model = RNABertForMaskedLM.from_pretrained( args.model_name, torch_dtype=getattr(torch, args.dtype) ) - model = model.to(args.device) - model.eval() + self.model = self.model.to(args.device).eval() + self.device = args.device - # Run masked-marginal inference - print("Running inference...") - sequence_scores = run_inference( - model, - tokenizer, - contexts, - ctx_tasks, - scores, - args.device, - args.batch_size, - args.max_batch_tokens, - ) - - # Add scores to DataFrame - score_column = "aido_rna_score" - dms_df[score_column] = sequence_scores - - # Calculate Spearman correlation (ignoring unscored variants) - correlation, pvalue = spearmanr( - dms_df["DMS_score"], dms_df[score_column], nan_policy="omit" - ) - - # Save results - output_file = output_dir / f"{dms_id}.csv" - dms_df.to_csv(output_file, index=False) - print(f"Saved results to: {output_file}") - - # Print summary statistics - print("\nSummary:") - print(f"Number of sequences: {len(sequences)}") - print( - f"Spearman correlation with DMS scores: {correlation:.3f} (p-value: {pvalue:.2e})" - ) - print(f"Output saved to: {output_file}") + self.base_ids = {b: tokenizer.convert_tokens_to_ids(b) for b in self.bases} + self.mask_id = tokenizer.mask_token_id + self.pad_id = tokenizer.pad_token_id + self.unk_id = tokenizer.unk_token_id # any base outside ACGT, e.g. an N + self.prefix_ids = [tokenizer.cls_token_id] + self.suffix_ids = [tokenizer.sep_token_id] - except Exception as e: - print(f"Error: {str(e)}", file=sys.stderr) - sys.exit(1) + def logits_at(self, input_ids, attention_mask, rows, cols): + outputs = self.model(input_ids=input_ids, attention_mask=attention_mask) + return outputs.logits[rows, cols] if __name__ == "__main__": - main() + main(AIDORNAAdapter()) # python score_aido_rna_single_dms.py --row_id 0 --ref_sheet reference_sheet.csv --dms_dir_path fitness_processed_assays --output_dir_path aido_rna_output --model_name genbio-ai/AIDO.RNA-1.6B diff --git a/fitness/baselines/RNAGenesis/score_rnagenesis.sh b/fitness/baselines/RNAGenesis/score_rnagenesis.sh index e261c33..217f9f2 100644 --- a/fitness/baselines/RNAGenesis/score_rnagenesis.sh +++ b/fitness/baselines/RNAGenesis/score_rnagenesis.sh @@ -14,10 +14,11 @@ export model_dir="path/to/local/RNAGenesis" export reference_sheet="reference_sheet.csv" -# Write predictions under a folder named "rnagenesis" so they line up with the -# "rnagenesis" entry in fitness/merge_scoring_files.py (which reads the -# rnagenesis_score column from model_predictions/rnagenesis/). -export output_scores_dir="path/to/model_predictions/rnagenesis" +# Write predictions under a folder named "rnagenesis_4fill". One run writes all +# four fill strategies into it as rnagenesis_score_wt_fill and so on, which the +# rnagenesis_wt_fill, rnagenesis_mask_fill, rnagenesis_mut_fill and rnagenesis_match_fill +# entries in fitness/merge_scoring_files.py read. +export output_scores_dir="path/to/model_predictions/rnagenesis_4fill" export dms_data_dir="path/to/dms/data/dir" # Reference-sheet row to score. Set by a Slurm array job (0-69), or defaults @@ -25,6 +26,15 @@ export dms_data_dir="path/to/dms/data/dir" # tRNA, aptamer); scoring only those leaves rnagenesis without mRNA predictions, # so read its aggregate from performance_fitness.py --type ncRNA. Under # --type all its All_Mean is NaN by design. + +# Masked-marginal fill strategy. The default computes all four (wt-fill, +# mask-fill, mut-fill, match-fill), which share their contexts and so cost only +# about 19% more unique context examples than mut-fill alone, and writes one +# column per strategy named +# {COLUMN}_{strategy}. Pass --strategies mut-fill (or any single strategy) to +# write the historical bare {COLUMN} column instead. See +# fitness/baselines/masked_lm/strategies.py for the formulas. + DMS_index=${SLURM_ARRAY_TASK_ID:-0} python score_rnagenesis_single_dms.py \ diff --git a/fitness/baselines/RNAGenesis/score_rnagenesis_single_dms.py b/fitness/baselines/RNAGenesis/score_rnagenesis_single_dms.py index 8292707..ee93976 100644 --- a/fitness/baselines/RNAGenesis/score_rnagenesis_single_dms.py +++ b/fitness/baselines/RNAGenesis/score_rnagenesis_single_dms.py @@ -1,22 +1,15 @@ #!/usr/bin/env python3 """ -Script to run RNAGenesis inference on DMS assay sequences. -Takes a reference sheet and row ID to process specific assays. +Score DMS assay sequences with RNAGenesis. RNAGenesis (https://github.com/zaixizhang/RNAGenesis) is a generalist RNA foundation model. Its released encoder (https://huggingface.co/Zaixi/RNAGenesis) is an xTrimoPGLM-style bidirectional transformer pretrained with a masked language-modelling objective on RNAcentral, exposed as -``xTrimoPGLMForMaskedLM``. We score variants with the masked-marginal -log-likelihood ratio, the same zero-shot proxy used by the other masked RNA -language-model baselines (RiNALMo, Orthrus, AIDO.RNA): - - score(variant) = sum_i [ log P(mut_i | variant context, pos_i masked) - - log P(wt_i | variant context, pos_i masked) ] - -summed over the variant's mutated positions ``i``. Each position is masked in -the variant's own sequence, so for a multi-mutant the remaining mutations stay -in the context. +``xTrimoPGLMForMaskedLM``. It is scored with the shared masked-marginal engine +in ``fitness/baselines/masked_lm``, which offers the four fill strategies +(``wt-fill``, ``mask-fill``, ``mut-fill``, ``match-fill``) that differ in what +the model sees at a variant's other mutated positions. Two properties of the released checkpoint drive the implementation: @@ -28,52 +21,32 @@ ``convert_tokens_to_ids`` iterates a string character by character. The author's own ``run.py`` depends on that character iteration and adds no special tokens, so input ids are built directly from ``tokenizer.model``: - one id per nucleotide, no CLS/EOS. + one id per nucleotide, no CLS or EOS. The mask token is ``tMASK``, the token-level mask of the xTrimoPGLM family. It was confirmed empirically: on wild-type ncRNA constructs it gives a lower masked -negative log-likelihood than gMASK or sMASK, and puts >99.9% of the predicted -mass on A/C/G/U. - -Masking one position of one variant is a single forward pass, but two variants -that differ only at the masked position share the same masked context. Those -contexts are deduplicated before inference, which is exact and cuts the number -of forward passes by up to 3x on single-substitution libraries. +negative log-likelihood than gMASK or sMASK, and puts over 99.9% of the +predicted mass on A/C/G/U. """ -import argparse import sys from pathlib import Path -import numpy as np -import pandas as pd import torch -from scipy.stats import spearmanr -from tqdm.auto import tqdm -from transformers import AutoModelForMaskedLM -# RNAGenesis uses the RNA alphabet: its vocabulary has U and no T. -BASES = "ACGU" -MASK_TOKEN = "tMASK" -UNK_TOKEN = "N" -MASK_CHAR = "#" +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from masked_lm import MaskedLMAdapter, main # noqa: E402 -def preprocess_sequence(sequence: str) -> str: - """ - Preprocess an RNA/DNA sequence for RNAGenesis: - - Convert to uppercase - - Convert DNA (T) to RNA (U) - - Remove any whitespace - """ - return sequence.strip().upper().replace("T", "U") +MASK_TOKEN = "tMASK" +UNK_TOKEN = "N" def load_vocab(model_path: str) -> dict: """ - Read the model's own vocabulary file and return a token -> id mapping. + Read the model's own vocabulary file and return a token to id mapping. - The HuggingFace tokenizer wrapper is bypassed on purpose, see module + The HuggingFace tokenizer wrapper is bypassed on purpose, see the module docstring. """ vocab_file = Path(model_path) / "tokenizer.model" @@ -86,351 +59,59 @@ def load_vocab(model_path: str) -> dict: return {token: index for index, token in enumerate(tokens)} -def parse_args(): - """Parse command line arguments.""" - parser = argparse.ArgumentParser( - description="Run RNAGenesis (masked LM) inference on DMS assay sequences." - ) - parser.add_argument( - "--row_id", - type=int, - required=True, - help="Row ID in the reference sheet to process", - ) - parser.add_argument( - "--ref_sheet", - type=str, - required=True, - help="Path to reference sheet containing DMS_ID column", - ) - parser.add_argument( - "--dms_dir_path", - type=str, - required=True, - help="Directory containing DMS CSV files", - ) - parser.add_argument( - "--output_dir_path", - type=str, - required=True, - help="Directory to save output files", - ) - parser.add_argument( - "--device", - type=str, - default="cuda:0" if torch.cuda.is_available() else "cpu", - help="Device to run inference on (default: cuda:0 if available, else cpu)", - ) - parser.add_argument( - "--model_name", - type=str, - required=True, - help="Local directory holding the RNAGenesis checkpoint. The released " - "repo omits quantization.py, which the modelling code imports, so a " - "local copy with that file added is required.", - ) - parser.add_argument( - "--dtype", - type=str, - default="bfloat16", - choices=["bfloat16", "float32"], - help="Weight/compute dtype. bfloat16 is ~2x faster and was used for the " - "released scores; float32 gives finer logits and fewer tied variants " - "(default: bfloat16)", - ) - parser.add_argument( - "--batch_size", - type=int, - default=256, - help="Maximum masked contexts scored per forward pass (default: 256)", - ) - parser.add_argument( - "--max_batch_tokens", - type=int, - default=32768, - help="Cap on batch_size x sequence length per forward pass, so that long " - "assays automatically use a smaller batch (default: 32768)", - ) - return parser.parse_args() - - -def load_reference_data(ref_sheet_path: str, row_id: int) -> str: - """ - Load reference sheet and get DMS_ID for specified row. - - Raises: - ValueError: If row_id is not found or DMS_ID is missing - """ - try: - ref_df = pd.read_csv(ref_sheet_path) - if row_id >= len(ref_df): - raise ValueError( - f"Row ID {row_id} exceeds number of rows in reference sheet" - ) - - dms_id = ref_df.loc[row_id, "DMS_ID"] - if pd.isna(dms_id): - raise ValueError(f"DMS_ID is missing for row {row_id}") - - return str(dms_id) - - except FileNotFoundError: - raise FileNotFoundError(f"Reference sheet not found: {ref_sheet_path}") - except KeyError: - raise KeyError("Reference sheet must contain 'DMS_ID' column") - - -def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: - """ - Load DMS data for specified DMS_ID. - - Raises: - FileNotFoundError: If DMS file is not found - """ - dms_file = Path(dms_dir_path) / f"{dms_id}.csv" - if not dms_file.exists(): - raise FileNotFoundError(f"DMS file not found: {dms_file}") - - df = pd.read_csv(dms_file) - required_cols = ["mutant", "DMS_score", "sequence"] - missing_cols = [col for col in required_cols if col not in df.columns] - if missing_cols: - raise ValueError(f"Missing required columns in DMS file: {missing_cols}") - - return df - - -def parse_mutations(mutant_str: str) -> list: - """ - Parse a mutation string such as ``"A4U,A5G"`` into a list of - ``(pos0, wt_base, mut_base)`` tuples with 0-based positions, in the RNA - alphabet used by the model. - - Raises: - ValueError: for non-substitution edits (indels) or unknown bases, so - the caller can score the affected variant as NaN. - """ - mutations = [] - for token in str(mutant_str).replace(" ", "").split(","): - if not token: - continue - wt_base = token[0].upper().replace("T", "U") - mut_base = token[-1].upper().replace("T", "U") - pos = int(token[1:-1]) - 1 # 1-based -> 0-based - if wt_base not in BASES or mut_base not in BASES: - raise ValueError(f"Unsupported mutation token: {token}") - mutations.append((pos, wt_base, mut_base)) - return mutations - - -def build_masking_tasks(mutants: list, sequences: list) -> tuple: - """ - Expand each variant into one masked-scoring task per mutated position, and - deduplicate identical masked contexts. - - A masked context is the variant's own sequence with the scored position - replaced by a mask placeholder. Two variants that differ only at that - position produce the same context and therefore the same log-probabilities, - so the forward pass is shared. The context also determines the masked - position, so it does not need to be tracked separately. - - Returns: - contexts: list of masked context strings, one per forward pass - ctx_tasks: list parallel to ``contexts``; entry k holds the - ``(row_idx, wt_base, mut_base)`` tuples scored from context k - scores: per-variant score array, pre-filled with 0.0 for scorable - variants and NaN for wild-type / unparseable / out-of-range rows - """ - scores = np.full(len(sequences), np.nan, dtype=float) - ctx_index = {} - contexts = [] - ctx_tasks = [] - - for i, (mutant_str, seq) in enumerate(zip(mutants, sequences)): - if pd.isna(mutant_str): - continue # wild-type row: leave as NaN - try: - mutations = parse_mutations(mutant_str) - if not mutations: - continue - row_tasks = [] - for pos, wt_base, mut_base in mutations: - if pos < 0 or pos >= len(seq): - raise ValueError(f"Mutation position {pos + 1} outside sequence") - if seq[pos] != mut_base: - raise ValueError( - f"Sequence has {seq[pos]} at position {pos + 1}, " - f"expected the mutant base {mut_base}" - ) - row_tasks.append((seq[:pos] + MASK_CHAR + seq[pos + 1 :], wt_base, mut_base)) - except (ValueError, IndexError) as err: - print(f"Skipping variant {mutant_str}: {err}") - continue # unsupported edit: leave as NaN - for context, wt_base, mut_base in row_tasks: - k = ctx_index.get(context) - if k is None: - k = len(contexts) - ctx_index[context] = k - contexts.append(context) - ctx_tasks.append([]) - ctx_tasks[k].append((i, wt_base, mut_base)) - scores[i] = 0.0 # scorable: accumulate per-position deltas below - - return contexts, ctx_tasks, scores - - -def run_inference( - model, - vocab: dict, - contexts: list, - ctx_tasks: list, - scores: np.ndarray, - device: str, - batch_size: int, - max_batch_tokens: int, -) -> np.ndarray: - """ - Score masked contexts with the RNAGenesis masked LM head and accumulate the - per-position log-likelihood ratios into each variant's score. - - Input ids are one per nucleotide with no special tokens, matching the - reference usage shipped with the checkpoint, so the masked position index in - the token sequence equals the position in the RNA sequence. - """ - base_ids = {b: vocab[b] for b in BASES} - mask_id = vocab[MASK_TOKEN] - pad_id = vocab[""] - unk_id = vocab[UNK_TOKEN] # any base outside ACGU, e.g. an N in a construct - - seq_len = max(len(c) for c in contexts) - batch_size = max(1, min(batch_size, max_batch_tokens // seq_len)) - print(f"Using batch size {batch_size} for sequence length {seq_len}") - - for start in tqdm( - range(0, len(contexts), batch_size), desc="Scoring", unit="batch" - ): - batch = contexts[start : start + batch_size] - max_len = max(len(c) for c in batch) - - input_ids = torch.full((len(batch), max_len), pad_id, dtype=torch.long) - attention_mask = torch.zeros((len(batch), max_len), dtype=torch.long) - mask_pos = [] - for b, context in enumerate(batch): - ids = [ - mask_id if ch == MASK_CHAR else base_ids.get(ch, unk_id) - for ch in context - ] - input_ids[b, : len(ids)] = torch.tensor(ids, dtype=torch.long) - attention_mask[b, : len(ids)] = 1 - # Dedup is only exact if a context masks exactly one position, which - # is what makes the context string identify the position. - if context.count(MASK_CHAR) != 1: - raise ValueError(f"Context does not mask exactly one position: {context}") - mask_pos.append(context.index(MASK_CHAR)) # no CLS offset - - input_ids = input_ids.to(device) - attention_mask = attention_mask.to(device) - rows = torch.arange(len(batch), device=device) - cols = torch.tensor(mask_pos, device=device) - - with torch.inference_mode(): - logits = model(input_ids=input_ids, attention_mask=attention_mask).logits - log_probs = torch.log_softmax(logits[rows, cols].float(), dim=-1) - log_probs = log_probs.cpu().numpy() - - for b in range(len(batch)): - row_log_probs = log_probs[b] - for row_idx, wt_base, mut_base in ctx_tasks[start + b]: - scores[row_idx] += ( - row_log_probs[base_ids[mut_base]] - row_log_probs[base_ids[wt_base]] - ) - - return scores - - -def main(): - args = parse_args() - - # Create output directory if it doesn't exist - output_dir = Path(args.output_dir_path) - output_dir.mkdir(parents=True, exist_ok=True) - - try: - # Load DMS ID from reference sheet - dms_id = load_reference_data(args.ref_sheet, args.row_id) - print(f"Processing DMS ID: {dms_id}") - - # Load DMS data - dms_df = load_dms_data(args.dms_dir_path, dms_id) - - # Preprocess sequences (uppercase, DNA -> RNA) - print("Preprocessing sequences...") - sequences = [preprocess_sequence(seq) for seq in dms_df["sequence"].tolist()] - - # Expand variants into deduplicated masked contexts - contexts, ctx_tasks, scores = build_masking_tasks( - dms_df["mutant"].tolist(), sequences +class RNAGenesisAdapter(MaskedLMAdapter): + """RNAGenesis: RNA alphabet, no special tokens, attention mask.""" + + name = "RNAGenesis" + bases = "ACGU" + score_column = "rnagenesis_score" + n_special_tokens = 0 # the reference usage adds none + default_batch_size = 256 + default_max_batch_tokens = 32768 + + @staticmethod + def add_arguments(parser): + parser.add_argument( + "--model_name", + type=str, + required=True, + help="Local path of the RNAGenesis encoder checkpoint directory, " + "which must contain tokenizer.model", ) - n_tasks = sum(len(t) for t in ctx_tasks) - print( - f"Prepared {n_tasks} masked positions across " - f"{int(np.sum(~np.isnan(scores)))} scorable variants " - f"(of {len(sequences)} total), deduplicated to {len(contexts)} " - f"forward passes" + parser.add_argument( + "--dtype", + type=str, + default="bfloat16", + choices=["bfloat16", "float32"], + help="Torch dtype for the model weights. bfloat16 is the default and " + "is what the released predictions were produced with", ) - if not contexts: - raise ValueError("No scorable variants found") - # Initialize model. The vocabulary is read from the checkpoint directly. - print(f"Initializing RNAGenesis model ({args.model_name})...") - vocab = load_vocab(args.model_name) - model = AutoModelForMaskedLM.from_pretrained( - args.model_name, trust_remote_code=True, torch_dtype=getattr(torch, args.dtype) - ) - model = model.to(args.device) - model.eval() - - # Run masked-marginal inference - print("Running inference...") - sequence_scores = run_inference( - model, - vocab, - contexts, - ctx_tasks, - scores, - args.device, - args.batch_size, - args.max_batch_tokens, - ) + def load(self, args): + from transformers import AutoModelForMaskedLM - # Add scores to DataFrame - score_column = "rnagenesis_score" - dms_df[score_column] = sequence_scores - - # Calculate Spearman correlation (ignoring unscored variants) - correlation, pvalue = spearmanr( - dms_df["DMS_score"], dms_df[score_column], nan_policy="omit" + vocab = load_vocab(args.model_name) + self.model = AutoModelForMaskedLM.from_pretrained( + args.model_name, + trust_remote_code=True, + torch_dtype=getattr(torch, args.dtype), ) + self.model = self.model.to(args.device).eval() + self.device = args.device - # Save results - output_file = output_dir / f"{dms_id}.csv" - dms_df.to_csv(output_file, index=False) - print(f"Saved results to: {output_file}") - - # Print summary statistics - print("\nSummary:") - print(f"Number of sequences: {len(sequences)}") - print( - f"Spearman correlation with DMS scores: {correlation:.3f} (p-value: {pvalue:.2e})" - ) - print(f"Output saved to: {output_file}") + self.base_ids = {b: vocab[b] for b in self.bases} + self.mask_id = vocab[MASK_TOKEN] + self.unk_id = vocab[UNK_TOKEN] + self.pad_id = vocab.get("", self.unk_id) + self.prefix_ids = [] # the reference usage adds no special tokens + self.suffix_ids = [] - except Exception as e: - print(f"Error: {str(e)}", file=sys.stderr) - sys.exit(1) + def logits_at(self, input_ids, attention_mask, rows, cols): + outputs = self.model(input_ids=input_ids, attention_mask=attention_mask) + return outputs.logits[rows, cols] if __name__ == "__main__": - main() + main(RNAGenesisAdapter()) # python score_rnagenesis_single_dms.py --row_id 0 --ref_sheet reference_sheet.csv --dms_dir_path fitness_processed_assays --output_dir_path rnagenesis_output --model_name /path/to/rnagenesis_checkpoint diff --git a/fitness/baselines/RNA_FM/score_rna_fm.sh b/fitness/baselines/RNA_FM/score_rna_fm.sh index 7a996fc..a5cd7d2 100644 --- a/fitness/baselines/RNA_FM/score_rna_fm.sh +++ b/fitness/baselines/RNA_FM/score_rna_fm.sh @@ -7,15 +7,25 @@ export checkpoint_path="path/to/RNA-FM_pretrained.pth" export reference_sheet="reference_sheet.csv" -# Write predictions under a folder named "RNA-FM" so they line up with the -# "RNA-FM" entry in fitness/merge_scoring_files.py (which reads the -# RNA_FM_scores column from model_predictions/RNA-FM/). -export output_scores_dir="path/to/model_predictions/RNA-FM" +# Write predictions under a folder named "rna_fm_4fill". One run writes all +# four fill strategies into it as RNA_FM_scores_wt_fill and so on, which the +# rna_fm_wt_fill, rna_fm_mask_fill, rna_fm_mut_fill and rna_fm_match_fill +# entries in fitness/merge_scoring_files.py read. +export output_scores_dir="path/to/model_predictions/rna_fm_4fill" export dms_data_dir="path/to/dms/data/dir" # Reference-sheet row to score. Set by a Slurm array job (0-69), or defaults # to 0 when run directly. Rows 0-8,11-32 are the 31 ncRNA assays; read an # ncRNA-only aggregate with performance_fitness.py --type ncRNA. + +# Masked-marginal fill strategy. The default computes all four (wt-fill, +# mask-fill, mut-fill, match-fill), which share their contexts and so cost only +# about 19% more unique context examples than mut-fill alone, and writes one +# column per strategy named +# {COLUMN}_{strategy}. Pass --strategies mut-fill (or any single strategy) to +# write the historical bare {COLUMN} column instead. See +# fitness/baselines/masked_lm/strategies.py for the formulas. + DMS_index=${SLURM_ARRAY_TASK_ID:-0} python score_rna_fm_single_dms.py \ diff --git a/fitness/baselines/RNA_FM/score_rna_fm_single_dms.py b/fitness/baselines/RNA_FM/score_rna_fm_single_dms.py index c532ba4..99d3b55 100644 --- a/fitness/baselines/RNA_FM/score_rna_fm_single_dms.py +++ b/fitness/baselines/RNA_FM/score_rna_fm_single_dms.py @@ -1,430 +1,79 @@ #!/usr/bin/env python3 """ -Script to run RNA-FM inference on DMS assay sequences. -Takes a reference sheet and row ID to process specific assays. +Score DMS assay sequences with RNA-FM. -RNA-FM is scored with the masked-marginal convention used by the other masked -RNA language-model baselines (RiNALMo, Orthrus, AIDO.RNA, RNAGenesis): - - score(variant) = sum_i [ log P(mut_i | variant context, pos_i masked) - - log P(wt_i | variant context, pos_i masked) ] - -Each mutated position is masked one at a time in the variant's OWN sequence, so -the remaining mutations of a multi-mutant stay in the context. +RNA-FM is scored with the shared masked-marginal engine in +``fitness/baselines/masked_lm``, which offers the four fill strategies +(``wt-fill``, ``mask-fill``, ``mut-fill``, ``match-fill``) that differ in what +the model sees at a variant's other mutated positions. See that package's +docstrings for the formulas and their provenance. This replaces the earlier ``compute_fitness.py``, which offered two strategies and whose released predictions match neither of the masked ones. Its ``masked-marginals`` masked ALL of a variant's mutated positions at once in the -WILD-TYPE sequence, so no mutation was ever visible in the context, and its +WILD-TYPE sequence, which is this package's ``mask-fill``, and its ``wt-marginals`` did not mask at all. The released RNA-FM predictions are reproduced to floating-point noise by ``wt-marginals`` (Pearson 1.000000, max -abs difference under 2e-5 on the two assays checked) and not by either masked +abs difference under 2e-5 on the two assays checked) and not by any masked strategy (Pearson 0.28 to 0.89), even though its run script requested -``masked-marginals``. Conventions agree on single substitutions and diverge on -everything else, and 99.4% of the ncRNA benchmark's variants are multi-mutants. - -Masking one position of one variant is a single forward pass, but two variants -that differ only at the masked position share the same masked context. Those -contexts are deduplicated before inference, which is exact and cuts the number -of forward passes by up to 3x on single-substitution libraries. +``masked-marginals``. """ -import argparse import sys from pathlib import Path -import numpy as np -import pandas as pd import torch -from scipy.stats import spearmanr -from tqdm.auto import tqdm - -# RNA-FM uses the RNA alphabet: A, C, G, U. -BASES = "ACGU" -MASK_CHAR = "#" - - -def preprocess_sequence(sequence: str) -> str: - """ - Preprocess an RNA/DNA sequence for RNA-FM: - - Convert to uppercase - - Convert DNA (T) to RNA (U) - - Remove any whitespace - """ - return sequence.strip().upper().replace("T", "U") - - -def parse_args(): - """Parse command line arguments.""" - parser = argparse.ArgumentParser( - description="Run RNA-FM (masked LM) inference on DMS assay sequences." - ) - parser.add_argument( - "--row_id", - type=int, - required=True, - help="Row ID in the reference sheet to process", - ) - parser.add_argument( - "--ref_sheet", - type=str, - required=True, - help="Path to reference sheet containing DMS_ID column", - ) - parser.add_argument( - "--dms_dir_path", - type=str, - required=True, - help="Directory containing DMS CSV files", - ) - parser.add_argument( - "--output_dir_path", - type=str, - required=True, - help="Directory to save output files", - ) - parser.add_argument( - "--device", - type=str, - default="cuda:0" if torch.cuda.is_available() else "cpu", - help="Device to run inference on (default: cuda:0 if available, else cpu)", - ) - parser.add_argument( - "--checkpoint_path", - type=str, - required=True, - help="Path to the RNA-FM_pretrained.pth weights file. Note this differs " - "from compute_fitness.py's --model_location, which is the cloned RNA-FM " - "module directory added to sys.path; here the package is expected to be " - "installed (pip install rna-fm) and only the checkpoint is passed.", - ) - parser.add_argument( - "--batch_size", - type=int, - default=512, - help="Maximum masked contexts scored per forward pass (default: 512)", - ) - parser.add_argument( - "--max_tokens", - type=int, - default=1024, - help="RNA-FM position limit including and . Contexts longer " - "than this are windowed around the scored position (default: 1024)", - ) - parser.add_argument( - "--max_batch_tokens", - type=int, - default=65536, - help="Cap on batch_size x sequence length per forward pass, so that long " - "assays automatically use a smaller batch (default: 65536)", - ) - return parser.parse_args() - - -def load_reference_data(ref_sheet_path: str, row_id: int) -> str: - """ - Load reference sheet and get DMS_ID for specified row. - - Raises: - ValueError: If row_id is not found or DMS_ID is missing - """ - try: - ref_df = pd.read_csv(ref_sheet_path) - if row_id >= len(ref_df): - raise ValueError( - f"Row ID {row_id} exceeds number of rows in reference sheet" - ) - - dms_id = ref_df.loc[row_id, "DMS_ID"] - if pd.isna(dms_id): - raise ValueError(f"DMS_ID is missing for row {row_id}") - - return str(dms_id) - - except FileNotFoundError: - raise FileNotFoundError(f"Reference sheet not found: {ref_sheet_path}") - except KeyError: - raise KeyError("Reference sheet must contain 'DMS_ID' column") - - -def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: - """ - Load DMS data for specified DMS_ID. - - Raises: - FileNotFoundError: If DMS file is not found - """ - dms_file = Path(dms_dir_path) / f"{dms_id}.csv" - if not dms_file.exists(): - raise FileNotFoundError(f"DMS file not found: {dms_file}") - - df = pd.read_csv(dms_file) - required_cols = ["mutant", "DMS_score", "sequence"] - missing_cols = [col for col in required_cols if col not in df.columns] - if missing_cols: - raise ValueError(f"Missing required columns in DMS file: {missing_cols}") - - return df - - -def parse_mutations(mutant_str: str) -> list: - """ - Parse a mutation string such as ``"A4U,A5G"`` into a list of - ``(pos0, wt_base, mut_base)`` tuples with 0-based positions, in the RNA - alphabet used by the model. - - Raises: - ValueError: for non-substitution edits (indels) or unknown bases, so - the caller can score the affected variant as NaN. - """ - mutations = [] - for token in str(mutant_str).replace(" ", "").split(","): - if not token: - continue - wt_base = token[0].upper().replace("T", "U") - mut_base = token[-1].upper().replace("T", "U") - pos = int(token[1:-1]) - 1 # 1-based -> 0-based - if wt_base not in BASES or mut_base not in BASES: - raise ValueError(f"Unsupported mutation token: {token}") - mutations.append((pos, wt_base, mut_base)) - return mutations - - -def window_context(context: str, max_tokens: int) -> str: - """ - Trim a masked context to fit RNA-FM's position limit, centred on the mask. - - RNA-FM accepts at most ``max_tokens`` tokens including and . The - ncRNA assays are all far shorter than that, but the mRNA-coding constructs - reach several thousand bases, so a window centred on the scored position is - taken, mirroring the windowing in this directory's compute_fitness.py. - Mutations that fall outside the window are lost from the context, which is - inherent to windowing. - """ - budget = max_tokens - 2 # room for and - if len(context) <= budget: - return context - pos = context.index(MASK_CHAR) - start = max(0, pos - budget // 2) - end = min(len(context), start + budget) - start = max(0, end - budget) - return context[start:end] - -def build_masking_tasks(mutants: list, sequences: list, max_tokens: int) -> tuple: - """ - Expand each variant into one masked-scoring task per mutated position, and - deduplicate identical masked contexts. +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - A masked context is the variant's own sequence with the scored position - replaced by a mask placeholder. Two variants that differ only at that - position produce the same context and therefore the same log-probabilities, - so the forward pass is shared. The context also determines the masked - position, so it does not need to be tracked separately. +from masked_lm import MaskedLMAdapter, main # noqa: E402 - Returns: - contexts: list of masked context strings, one per forward pass - ctx_tasks: list parallel to ``contexts``; entry k holds the - ``(row_idx, wt_base, mut_base)`` tuples scored from context k - scores: per-variant score array, pre-filled with 0.0 for scorable - variants and NaN for wild-type / unparseable / out-of-range rows - """ - scores = np.full(len(sequences), np.nan, dtype=float) - ctx_index = {} - contexts = [] - ctx_tasks = [] - for i, (mutant_str, seq) in enumerate(zip(mutants, sequences)): - if pd.isna(mutant_str): - continue # wild-type row: leave as NaN - try: - mutations = parse_mutations(mutant_str) - if not mutations: - continue - row_tasks = [] - for pos, wt_base, mut_base in mutations: - if pos < 0 or pos >= len(seq): - raise ValueError(f"Mutation position {pos + 1} outside sequence") - if seq[pos] != mut_base: - raise ValueError( - f"Sequence has {seq[pos]} at position {pos + 1}, " - f"expected the mutant base {mut_base}" - ) - context = window_context( - seq[:pos] + MASK_CHAR + seq[pos + 1 :], max_tokens - ) - row_tasks.append((context, wt_base, mut_base)) - except (ValueError, IndexError) as err: - print(f"Skipping variant {mutant_str}: {err}") - continue # unsupported edit: leave as NaN - for context, wt_base, mut_base in row_tasks: - k = ctx_index.get(context) - if k is None: - k = len(contexts) - ctx_index[context] = k - contexts.append(context) - ctx_tasks.append([]) - ctx_tasks[k].append((i, wt_base, mut_base)) - scores[i] = 0.0 # scorable: accumulate per-position deltas below +class RNAFMAdapter(MaskedLMAdapter): + """RNA-FM: RNA alphabet, ```` and ````, no attention mask.""" - return contexts, ctx_tasks, scores + name = "RNA-FM" + bases = "ACGU" + score_column = "RNA_FM_scores" + max_tokens = 1024 # position limit including and + n_special_tokens = 2 - -def run_inference( - model, - alphabet, - contexts: list, - ctx_tasks: list, - scores: np.ndarray, - device: str, - batch_size: int, - max_batch_tokens: int, -) -> np.ndarray: - """ - Score masked contexts with RNA-FM's masked LM head and accumulate the - per-position log-likelihood ratios into each variant's score. - - Token layout follows RNA-FM's batch converter: ```` + sequence + - ````, so the masked position index is offset by one. - """ - base_ids = {b: alphabet.get_idx(b) for b in BASES} - mask_id = alphabet.mask_idx - pad_id = alphabet.padding_idx - cls_id = alphabet.cls_idx - eos_id = alphabet.eos_idx - unk_id = alphabet.get_idx("N") # any base outside ACGU - - seq_len = max(len(c) for c in contexts) - batch_size = max(1, min(batch_size, max_batch_tokens // (seq_len + 2))) - print(f"Using batch size {batch_size} for sequence length {seq_len}") - - for start in tqdm( - range(0, len(contexts), batch_size), desc="Scoring", unit="batch" - ): - batch = contexts[start : start + batch_size] - max_len = max(len(c) for c in batch) + 2 # ... - - input_ids = torch.full((len(batch), max_len), pad_id, dtype=torch.long) - mask_pos = [] - for b, context in enumerate(batch): - ids = [cls_id] - for ch in context: - ids.append(mask_id if ch == MASK_CHAR else base_ids.get(ch, unk_id)) - ids.append(eos_id) - input_ids[b, : len(ids)] = torch.tensor(ids, dtype=torch.long) - # Dedup is only exact if a context masks exactly one position, which - # is what makes the context string identify the position. - if context.count(MASK_CHAR) != 1: - raise ValueError(f"Context does not mask exactly one position: {context}") - mask_pos.append(context.index(MASK_CHAR) + 1) # +1 for - - input_ids = input_ids.to(device) - rows = torch.arange(len(batch), device=device) - cols = torch.tensor(mask_pos, device=device) - - with torch.inference_mode(): - logits = model(input_ids)["logits"] - log_probs = torch.log_softmax(logits[rows, cols].float(), dim=-1) - log_probs = log_probs.cpu().numpy() - - for b in range(len(batch)): - row_log_probs = log_probs[b] - for row_idx, wt_base, mut_base in ctx_tasks[start + b]: - scores[row_idx] += ( - row_log_probs[base_ids[mut_base]] - row_log_probs[base_ids[wt_base]] - ) - - return scores - - -def main(): - args = parse_args() - - # Create output directory if it doesn't exist - output_dir = Path(args.output_dir_path) - output_dir.mkdir(parents=True, exist_ok=True) - - try: - # Load DMS ID from reference sheet - dms_id = load_reference_data(args.ref_sheet, args.row_id) - print(f"Processing DMS ID: {dms_id}") - - # Load DMS data - dms_df = load_dms_data(args.dms_dir_path, dms_id) - - # Preprocess sequences (uppercase, DNA -> RNA) - print("Preprocessing sequences...") - sequences = [preprocess_sequence(seq) for seq in dms_df["sequence"].tolist()] - - # Expand variants into deduplicated masked contexts - contexts, ctx_tasks, scores = build_masking_tasks( - dms_df["mutant"].tolist(), sequences, args.max_tokens + @staticmethod + def add_arguments(parser): + parser.add_argument( + "--checkpoint_path", + type=str, + required=True, + help="Path to the RNA-FM_pretrained.pth weights file. Note this differs " + "from compute_fitness.py's --model_location, which is the cloned RNA-FM " + "module directory added to sys.path; here the package is expected to be " + "installed (pip install rna-fm) and only the checkpoint is passed.", ) - n_tasks = sum(len(t) for t in ctx_tasks) - print( - f"Prepared {n_tasks} masked positions across " - f"{int(np.sum(~np.isnan(scores)))} scorable variants " - f"(of {len(sequences)} total), deduplicated to {len(contexts)} " - f"forward passes" - ) - if not contexts: - raise ValueError("No scorable variants found") - # Initialize model. The released checkpoint pickles an argparse.Namespace, - # which torch>=2.6 refuses to unpickle under the weights_only default, so - # that one class is allowlisted rather than disabling the check entirely. - print(f"Initializing RNA-FM model ({args.checkpoint_path})...") + def load(self, args): + # The released checkpoint pickles an argparse.Namespace, which torch>=2.6 + # refuses to unpickle under the weights_only default, so that one class is + # allowlisted rather than disabling the check entirely. import argparse as _argparse torch.serialization.add_safe_globals([_argparse.Namespace]) import fm - model, alphabet = fm.pretrained.rna_fm_t12(args.checkpoint_path) - model = model.to(args.device) - model.eval() - - # Run masked-marginal inference - print("Running inference...") - sequence_scores = run_inference( - model, - alphabet, - contexts, - ctx_tasks, - scores, - args.device, - args.batch_size, - args.max_batch_tokens, - ) - - # Add scores to DataFrame. The column name matches the original RNA-FM - # baseline so these predictions are a drop-in replacement. - score_column = "RNA_FM_scores" - dms_df[score_column] = sequence_scores - - # Calculate Spearman correlation (ignoring unscored variants) - correlation, pvalue = spearmanr( - dms_df["DMS_score"], dms_df[score_column], nan_policy="omit" - ) - - # Save results - output_file = output_dir / f"{dms_id}.csv" - dms_df.to_csv(output_file, index=False) - print(f"Saved results to: {output_file}") - - # Print summary statistics - print("\nSummary:") - print(f"Number of sequences: {len(sequences)}") - print( - f"Spearman correlation with DMS scores: {correlation:.3f} (p-value: {pvalue:.2e})" - ) - print(f"Output saved to: {output_file}") + self.model, alphabet = fm.pretrained.rna_fm_t12(args.checkpoint_path) + self.model = self.model.to(args.device).eval() + self.device = args.device + self.base_ids = {b: alphabet.get_idx(b) for b in self.bases} + self.mask_id = alphabet.mask_idx + self.pad_id = alphabet.padding_idx + self.unk_id = alphabet.get_idx("N") # any base outside ACGU + self.prefix_ids = [alphabet.cls_idx] + self.suffix_ids = [alphabet.eos_idx] - except Exception as e: - print(f"Error: {str(e)}", file=sys.stderr) - sys.exit(1) + def logits_at(self, input_ids, attention_mask, rows, cols): + return self.model(input_ids)["logits"][rows, cols] if __name__ == "__main__": - main() + main(RNAFMAdapter()) # python score_rna_fm_single_dms.py --row_id 0 --ref_sheet reference_sheet.csv --dms_dir_path fitness_processed_assays --output_dir_path rna_fm_output --checkpoint_path RNA-FM_pretrained.pth diff --git a/fitness/baselines/RiNALMo/score_rinalmo.sh b/fitness/baselines/RiNALMo/score_rinalmo.sh index d0f0821..6ff7c30 100644 --- a/fitness/baselines/RiNALMo/score_rinalmo.sh +++ b/fitness/baselines/RiNALMo/score_rinalmo.sh @@ -8,15 +8,25 @@ export checkpoint_path="path/to/rinalmo_giga_pretrained.pt" export reference_sheet="reference_sheet.csv" -# Write predictions under a folder named "rinalmo" so they line up with the -# "rinalmo" entry in fitness/merge_scoring_files.py (which reads the -# logit_scores column from model_predictions/rinalmo/). -export output_scores_dir="path/to/model_predictions/rinalmo" +# Write predictions under a folder named "rinalmo_4fill". One run writes all +# four fill strategies into it as logit_scores_wt_fill and so on, which the +# rinalmo_wt_fill, rinalmo_mask_fill, rinalmo_mut_fill and rinalmo_match_fill +# entries in fitness/merge_scoring_files.py read. +export output_scores_dir="path/to/model_predictions/rinalmo_4fill" export dms_data_dir="path/to/dms/data/dir" # Reference-sheet row to score. Set by a Slurm array job (0-69), or defaults # to 0 when run directly. Rows 0-8,11-32 are the 31 ncRNA assays; read an # ncRNA-only aggregate with performance_fitness.py --type ncRNA. + +# Masked-marginal fill strategy. The default computes all four (wt-fill, +# mask-fill, mut-fill, match-fill), which share their contexts and so cost only +# about 19% more unique context examples than mut-fill alone, and writes one +# column per strategy named +# {COLUMN}_{strategy}. Pass --strategies mut-fill (or any single strategy) to +# write the historical bare {COLUMN} column instead. See +# fitness/baselines/masked_lm/strategies.py for the formulas. + DMS_index=${SLURM_ARRAY_TASK_ID:-0} python score_rinalmo_single_dms.py \ diff --git a/fitness/baselines/RiNALMo/score_rinalmo_single_dms.py b/fitness/baselines/RiNALMo/score_rinalmo_single_dms.py index 05d835f..bfb61b9 100644 --- a/fitness/baselines/RiNALMo/score_rinalmo_single_dms.py +++ b/fitness/baselines/RiNALMo/score_rinalmo_single_dms.py @@ -1,21 +1,15 @@ #!/usr/bin/env python3 """ -Script to run RiNALMo inference on DMS assay sequences. -Takes a reference sheet and row ID to process specific assays. +Score DMS assay sequences with RiNALMo. -This rescores RiNALMo with the masked-marginal convention used by the other -masked RNA language-model baselines (Orthrus, AIDO.RNA, RNAGenesis, and the -rescored RNA-FM): - - score(variant) = sum_i [ log P(mut_i | variant context, pos_i masked) - - log P(wt_i | variant context, pos_i masked) ] - -Each mutated position is masked one at a time in the variant's OWN sequence, so -the remaining mutations of a multi-mutant stay in the context. +RiNALMo is scored with the shared masked-marginal engine in +``fitness/baselines/masked_lm``, which offers the four fill strategies +(``wt-fill``, ``mask-fill``, ``mut-fill``, ``match-fill``) that differ in what +the model sees at a variant's other mutated positions. This replaces the earlier ``compute_fitness.py``, which already masked one -position at a time in the variant's own sequence but read the wrong logit for the -mutant base. RiNALMo's alphabet is DNA-based (A C G T), and while +position at a time in the variant's own sequence but read the wrong logit for +the mutant base. RiNALMo's alphabet is DNA-based (A C G T), and while ``Alphabet.encode`` folds U to T internally, ``Alphabet.get_idx('U')`` returns . That script looks the mutant base up in a U-form sequence and the wild-type base up in a T-form sequence, so every mutation to U scores against @@ -25,428 +19,75 @@ carry it, though the match is not exact so at least one other difference remains. That script also stripped N from the wild type inside the scoring loop without adjusting coordinates, which would shift positions on any construct -containing N. This script does the same masking with correct T-form lookups on -both sides. - -Masking one position of one variant is a single forward pass, but two variants -that differ only at the masked position share the same masked context. Those -contexts are deduplicated before inference, which is exact and cuts the number -of forward passes by up to 3x on single-substitution libraries. +containing N. The shared engine's ``check_alphabet`` refuses to run if any base +maps to the unknown token, so this class of bug cannot recur silently. + +Precision note: the transformer body must run in bfloat16 because the +checkpoint's flash-attention kernels accept only fp16 or bf16, but bf16 logits +are too coarse for this score. The model's logits are large enough that bf16 +spacing quantises log-ratio differences onto multiples of about 0.125, tying +about 8% of single mutants at exactly 0. The masked LM head is therefore +recomputed in fp32 from the bf16 representation, which keeps the ranking usable. """ -import argparse import sys from pathlib import Path -import numpy as np -import pandas as pd import torch -from scipy.stats import spearmanr -from tqdm.auto import tqdm - -# RiNALMo uses the DNA alphabet: A, C, G, T. Alphabet.encode folds U to T. -BASES = "ACGT" -MASK_CHAR = "#" - - -def preprocess_sequence(sequence: str) -> str: - """ - Preprocess an RNA/DNA sequence for RiNALMo: - - Convert to uppercase - - Convert RNA (U) to DNA (T), matching the model's alphabet - - Remove any whitespace - """ - return sequence.strip().upper().replace("U", "T") - - -def parse_args(): - """Parse command line arguments.""" - parser = argparse.ArgumentParser( - description="Run RiNALMo (masked LM) inference on DMS assay sequences." - ) - parser.add_argument( - "--row_id", - type=int, - required=True, - help="Row ID in the reference sheet to process", - ) - parser.add_argument( - "--ref_sheet", - type=str, - required=True, - help="Path to reference sheet containing DMS_ID column", - ) - parser.add_argument( - "--dms_dir_path", - type=str, - required=True, - help="Directory containing DMS CSV files", - ) - parser.add_argument( - "--output_dir_path", - type=str, - required=True, - help="Directory to save output files", - ) - parser.add_argument( - "--device", - type=str, - default="cuda:0", - help="CUDA device to run inference on (default: cuda:0). RiNALMo's giga " - "checkpoint stores the flash-attention module layout, which is CUDA only, " - "so there is no working CPU path.", - ) - parser.add_argument( - "--checkpoint_path", - type=str, - required=True, - help="Path to the RiNALMo giga-v1 checkpoint (.pt), e.g. " - "rinalmo_giga_pretrained.pt from the project's Zenodo record.", - ) - parser.add_argument( - "--batch_size", - type=int, - default=512, - help="Maximum masked contexts scored per forward pass (default: 512)", - ) - parser.add_argument( - "--max_tokens", - type=int, - default=1024, - help="RiNALMo position limit including and . Contexts longer " - "than this are windowed around the scored position (default: 1024)", - ) - parser.add_argument( - "--max_batch_tokens", - type=int, - default=65536, - help="Cap on batch_size x sequence length per forward pass, so that long " - "assays automatically use a smaller batch (default: 65536)", - ) - return parser.parse_args() - - -def load_reference_data(ref_sheet_path: str, row_id: int) -> str: - """ - Load reference sheet and get DMS_ID for specified row. - - Raises: - ValueError: If row_id is not found or DMS_ID is missing - """ - try: - ref_df = pd.read_csv(ref_sheet_path) - if row_id >= len(ref_df): - raise ValueError( - f"Row ID {row_id} exceeds number of rows in reference sheet" - ) - - dms_id = ref_df.loc[row_id, "DMS_ID"] - if pd.isna(dms_id): - raise ValueError(f"DMS_ID is missing for row {row_id}") - - return str(dms_id) - - except FileNotFoundError: - raise FileNotFoundError(f"Reference sheet not found: {ref_sheet_path}") - except KeyError: - raise KeyError("Reference sheet must contain 'DMS_ID' column") - - -def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: - """ - Load DMS data for specified DMS_ID. - - Raises: - FileNotFoundError: If DMS file is not found - """ - dms_file = Path(dms_dir_path) / f"{dms_id}.csv" - if not dms_file.exists(): - raise FileNotFoundError(f"DMS file not found: {dms_file}") - - df = pd.read_csv(dms_file) - required_cols = ["mutant", "DMS_score", "sequence"] - missing_cols = [col for col in required_cols if col not in df.columns] - if missing_cols: - raise ValueError(f"Missing required columns in DMS file: {missing_cols}") - - return df - - -def parse_mutations(mutant_str: str) -> list: - """ - Parse a mutation string such as ``"A4U,A5G"`` into a list of - ``(pos0, wt_base, mut_base)`` tuples with 0-based positions, in the RNA - alphabet used by the model. - - Raises: - ValueError: for non-substitution edits (indels) or unknown bases, so - the caller can score the affected variant as NaN. - """ - mutations = [] - for token in str(mutant_str).replace(" ", "").split(","): - if not token: - continue - wt_base = token[0].upper().replace("U", "T") - mut_base = token[-1].upper().replace("U", "T") - pos = int(token[1:-1]) - 1 # 1-based -> 0-based - if wt_base not in BASES or mut_base not in BASES: - raise ValueError(f"Unsupported mutation token: {token}") - mutations.append((pos, wt_base, mut_base)) - return mutations - - -def window_context(context: str, max_tokens: int) -> str: - """ - Trim a masked context to fit RiNALMo's position limit, centred on the mask. - - RiNALMo accepts at most ``max_tokens`` tokens including and . The - ncRNA assays are all far shorter than that, but the mRNA-coding constructs - reach several thousand bases, so a window centred on the scored position is - taken, mirroring the windowing in this directory's compute_fitness.py. - Mutations that fall outside the window are lost from the context, which is - inherent to windowing. - """ - budget = max_tokens - 2 # room for and - if len(context) <= budget: - return context - pos = context.index(MASK_CHAR) - start = max(0, pos - budget // 2) - end = min(len(context), start + budget) - start = max(0, end - budget) - return context[start:end] - - -def build_masking_tasks(mutants: list, sequences: list, max_tokens: int) -> tuple: - """ - Expand each variant into one masked-scoring task per mutated position, and - deduplicate identical masked contexts. - A masked context is the variant's own sequence with the scored position - replaced by a mask placeholder. Two variants that differ only at that - position produce the same context and therefore the same log-probabilities, - so the forward pass is shared. The context also determines the masked - position, so it does not need to be tracked separately. +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - Returns: - contexts: list of masked context strings, one per forward pass - ctx_tasks: list parallel to ``contexts``; entry k holds the - ``(row_idx, wt_base, mut_base)`` tuples scored from context k - scores: per-variant score array, pre-filled with 0.0 for scorable - variants and NaN for wild-type / unparseable / out-of-range rows - """ - scores = np.full(len(sequences), np.nan, dtype=float) - ctx_index = {} - contexts = [] - ctx_tasks = [] +from masked_lm import MaskedLMAdapter, main # noqa: E402 - for i, (mutant_str, seq) in enumerate(zip(mutants, sequences)): - if pd.isna(mutant_str): - continue # wild-type row: leave as NaN - try: - mutations = parse_mutations(mutant_str) - if not mutations: - continue - row_tasks = [] - for pos, wt_base, mut_base in mutations: - if pos < 0 or pos >= len(seq): - raise ValueError(f"Mutation position {pos + 1} outside sequence") - if seq[pos] != mut_base: - raise ValueError( - f"Sequence has {seq[pos]} at position {pos + 1}, " - f"expected the mutant base {mut_base}" - ) - context = window_context( - seq[:pos] + MASK_CHAR + seq[pos + 1 :], max_tokens - ) - row_tasks.append((context, wt_base, mut_base)) - except (ValueError, IndexError) as err: - print(f"Skipping variant {mutant_str}: {err}") - continue # unsupported edit: leave as NaN - for context, wt_base, mut_base in row_tasks: - k = ctx_index.get(context) - if k is None: - k = len(contexts) - ctx_index[context] = k - contexts.append(context) - ctx_tasks.append([]) - ctx_tasks[k].append((i, wt_base, mut_base)) - scores[i] = 0.0 # scorable: accumulate per-position deltas below - return contexts, ctx_tasks, scores +class RiNALMoAdapter(MaskedLMAdapter): + """RiNALMo: DNA alphabet, ```` and ````, CUDA only.""" + name = "RiNALMo" + bases = "ACGT" + score_column = "logit_scores" + max_tokens = 1024 + n_special_tokens = 2 + requires_cuda = True # the giga checkpoint stores flash-attention modules -def run_inference( - model, - alphabet, - contexts: list, - ctx_tasks: list, - scores: np.ndarray, - device: str, - batch_size: int, - max_batch_tokens: int, -) -> np.ndarray: - """ - Score masked contexts with RiNALMo's masked LM head and accumulate the - per-position log-likelihood ratios into each variant's score. - - Token layout follows RiNALMo's Alphabet.encode: ```` + sequence + - ````, so the masked position index is offset by one. - """ - base_ids = {b: alphabet.get_idx(b) for b in BASES} - if alphabet.unk_idx in base_ids.values(): - raise ValueError(f"Alphabet does not cover {BASES}: {base_ids}") - mask_id = alphabet.mask_idx - pad_id = alphabet.pad_idx - cls_id = alphabet.cls_idx - eos_id = alphabet.eos_idx - unk_id = alphabet.get_idx("N") # any base outside ACGT - - seq_len = max(len(c) for c in contexts) - batch_size = max(1, min(batch_size, max_batch_tokens // (seq_len + 2))) - print(f"Using batch size {batch_size} for sequence length {seq_len}") - - for start in tqdm( - range(0, len(contexts), batch_size), desc="Scoring", unit="batch" - ): - batch = contexts[start : start + batch_size] - max_len = max(len(c) for c in batch) + 2 # ... - - input_ids = torch.full((len(batch), max_len), pad_id, dtype=torch.long) - mask_pos = [] - for b, context in enumerate(batch): - ids = [cls_id] - for ch in context: - ids.append(mask_id if ch == MASK_CHAR else base_ids.get(ch, unk_id)) - ids.append(eos_id) - input_ids[b, : len(ids)] = torch.tensor(ids, dtype=torch.long) - # Dedup is only exact if a context masks exactly one position, which - # is what makes the context string identify the position. - if context.count(MASK_CHAR) != 1: - raise ValueError(f"Context does not mask exactly one position: {context}") - mask_pos.append(context.index(MASK_CHAR) + 1) # +1 for - - input_ids = input_ids.to(device) - rows = torch.arange(len(batch), device=device) - cols = torch.tensor(mask_pos, device=device) - - # The transformer body must run in bfloat16 because the checkpoint's - # flash-attention kernels accept only fp16/bf16, but bf16 logits are too - # coarse for this score: the model's logits are large enough that bf16 - # spacing quantises log-ratio differences onto multiples of ~0.125, tying - # ~8% of single mutants at exactly 0. So the masked LM head is recomputed - # in fp32 from the bf16 representation, which keeps the ranking usable. - with torch.inference_mode(): - with torch.autocast("cuda", dtype=torch.bfloat16): - representation = model(input_ids)["representation"] - logits = model.lm_mask_head(representation.float()) - log_probs = torch.log_softmax(logits[rows, cols].float(), dim=-1) - log_probs = log_probs.cpu().numpy() - - for b in range(len(batch)): - row_log_probs = log_probs[b] - for row_idx, wt_base, mut_base in ctx_tasks[start + b]: - scores[row_idx] += ( - row_log_probs[base_ids[mut_base]] - row_log_probs[base_ids[wt_base]] - ) - - return scores - - -def main(): - args = parse_args() - - # Create output directory if it doesn't exist - output_dir = Path(args.output_dir_path) - output_dir.mkdir(parents=True, exist_ok=True) - - try: - if not args.device.startswith("cuda") or not torch.cuda.is_available(): - raise ValueError( - "RiNALMo requires a CUDA device: the giga checkpoint stores the " - "flash-attention module layout, which has no CPU implementation." - ) - - # Load DMS ID from reference sheet - dms_id = load_reference_data(args.ref_sheet, args.row_id) - print(f"Processing DMS ID: {dms_id}") - - # Load DMS data - dms_df = load_dms_data(args.dms_dir_path, dms_id) - - # Preprocess sequences (uppercase, RNA -> DNA) - print("Preprocessing sequences...") - sequences = [preprocess_sequence(seq) for seq in dms_df["sequence"].tolist()] - - # Expand variants into deduplicated masked contexts - contexts, ctx_tasks, scores = build_masking_tasks( - dms_df["mutant"].tolist(), sequences, args.max_tokens - ) - n_tasks = sum(len(t) for t in ctx_tasks) - print( - f"Prepared {n_tasks} masked positions across " - f"{int(np.sum(~np.isnan(scores)))} scorable variants " - f"(of {len(sequences)} total), deduplicated to {len(contexts)} " - f"forward passes" + @staticmethod + def add_arguments(parser): + parser.add_argument( + "--checkpoint_path", + type=str, + required=True, + help="Path to the RiNALMo giga-v1 checkpoint (.pt), e.g. " + "rinalmo_giga_pretrained.pt from the project's Zenodo record.", ) - if not contexts: - raise ValueError("No scorable variants found") - # Initialize model. The released checkpoint stores the flash-attention - # module layout, so flash attention must stay enabled. - print(f"Initializing RiNALMo model ({args.checkpoint_path})...") + def load(self, args): + # The released checkpoint stores the flash-attention module layout, so + # flash attention must stay enabled. from rinalmo.config import model_config from rinalmo.data.alphabet import Alphabet from rinalmo.model.model import RiNALMo config = model_config("giga") - model = RiNALMo(config) - model.load_state_dict(torch.load(args.checkpoint_path, weights_only=True)) - alphabet = Alphabet(**config["alphabet"]) - model = model.to(args.device) - model.eval() - - # Run masked-marginal inference - print("Running inference...") - sequence_scores = run_inference( - model, - alphabet, - contexts, - ctx_tasks, - scores, - args.device, - args.batch_size, - args.max_batch_tokens, - ) + self.model = RiNALMo(config) + self.model.load_state_dict(torch.load(args.checkpoint_path, weights_only=True)) + self.model = self.model.to(args.device).eval() + self.device = args.device - # Add scores to DataFrame. The column name matches the original RiNALMo - # baseline so these predictions are a drop-in replacement. - score_column = "logit_scores" - dms_df[score_column] = sequence_scores - - # Calculate Spearman correlation (ignoring unscored variants) - correlation, pvalue = spearmanr( - dms_df["DMS_score"], dms_df[score_column], nan_policy="omit" - ) - - # Save results - output_file = output_dir / f"{dms_id}.csv" - dms_df.to_csv(output_file, index=False) - print(f"Saved results to: {output_file}") - - # Print summary statistics - print("\nSummary:") - print(f"Number of sequences: {len(sequences)}") - print( - f"Spearman correlation with DMS scores: {correlation:.3f} (p-value: {pvalue:.2e})" - ) - print(f"Output saved to: {output_file}") + alphabet = Alphabet(**config["alphabet"]) + self.base_ids = {b: alphabet.get_idx(b) for b in self.bases} + self.mask_id = alphabet.mask_idx + self.pad_id = alphabet.pad_idx + self.unk_id = alphabet.unk_idx + self.prefix_ids = [alphabet.cls_idx] + self.suffix_ids = [alphabet.eos_idx] - except Exception as e: - print(f"Error: {str(e)}", file=sys.stderr) - sys.exit(1) + def logits_at(self, input_ids, attention_mask, rows, cols): + with torch.autocast("cuda", dtype=torch.bfloat16): + representation = self.model(input_ids)["representation"] + return self.model.lm_mask_head(representation.float())[rows, cols] if __name__ == "__main__": - main() + main(RiNALMoAdapter()) -# python score_rinalmo_single_dms.py --row_id 0 --ref_sheet reference_sheet.csv --dms_dir_path fitness_processed_assays --output_dir_path rinalmo_output --checkpoint_path giga-v1.pt +# python score_rinalmo_single_dms.py --row_id 0 --ref_sheet reference_sheet.csv --dms_dir_path fitness_processed_assays --output_dir_path rinalmo_output --checkpoint_path rinalmo_giga_pretrained.pt diff --git a/fitness/baselines/masked_lm/README.md b/fitness/baselines/masked_lm/README.md new file mode 100644 index 0000000..27f10fc --- /dev/null +++ b/fitness/baselines/masked_lm/README.md @@ -0,0 +1,73 @@ +# Masked-marginal scoring + +Shared scoring code for the benchmark's masked RNA language models. A model's +script in `fitness/baselines//` supplies an alphabet, a tokenization and +one forward pass; everything else lives here. + +## The four fill strategies + +Every masked-marginal score computes a log-odds at each mutated position of a +variant and sums over its mutations. The strategies differ in exactly one thing: +what fills the variant's OTHER mutated positions while the scored position is +masked. With `M` the mutated positions, `x^wt` and `x^mt` the wild-type and +variant sequences, `x_-i` a mask at `i` and `x_-M` masks at every position in `M` +(Meier et al. 2021, ESM-1v supplement, Appendix A): + +| strategy | fill at the other mutated sites | score | +|:--|:--|:--| +| `wt-fill` | wild-type bases | `sum_i log p(mt_i \| x^wt_-i) - log p(wt_i \| x^wt_-i)` | +| `mask-fill` | masks | `sum_i log p(mt_i \| x^wt_-M) - log p(wt_i \| x^wt_-M)` | +| `mut-fill` | mutant bases | `sum_i log p(mt_i \| x^mt_-i) - log p(wt_i \| x^mt_-i)` | +| `match-fill` | the allele being scored | `sum_i log p(mt_i \| x^mt_-i) - log p(wt_i \| x^wt_-i)` | + +`wt-fill` is what the ESM authors' released code and ProteinGym's baseline +implement under the name `masked-marginals`; `mask-fill` is the formula written +in the ESM paper. The name "masked marginals" is overloaded, and `wt-marginals` +is a different method again: one unmasked forward pass, no masking at all. + +All four are identical on single mutants, because a variant with one mutation has +no other mutated positions. They diverge on multi-mutants, which are 99.4% of the +non-coding benchmark. `match-fill` alone mixes two contexts, so it is a +difference of two conditionals rather than a log-odds ratio. + +## Cost + +One run computes all four. Their contexts overlap, so on the 31 non-coding assays +the four together need 2,929,196 unique context examples against 2,458,521 for +`mut-fill` alone, about 19% more. Four separate runs would cost 2.6 times as much +and still could not produce `match-fill`, which needs the per-position halves +rather than the final scores. + +## Files + +| file | contents | +|:--|:--| +| `strategies.py` | mutation parsing, wild-type recovery, per-variant validation, the four strategies as contexts and terms | +| `engine.py` | windowing, batching, the gather, and accumulation | +| `adapter.py` | the model interface | +| `runner.py` | command line, assay input, wild-type cross-check, output and manifest | + +## Adding a model + +Subclass `MaskedLMAdapter`, declare the alphabet, output column, special-token +count and batching defaults, implement `add_arguments`, `load` and `logits_at`, +and call `runner.main`. `fitness/baselines/RNA_FM/score_rna_fm_single_dms.py` is +the shortest example. Two rules the engine enforces rather than assumes: the +declared special-token count must match the loaded tokenizer, and context +positions must map to token positions by a constant shift. + +## Output + +One CSV per assay, holding the assay dataframe plus `{column}_{strategy}` for +each strategy computed, and a `{DMS_ID}.manifest.json` recording the strategies, +alphabet, checkpoint, dtype, counts and a hash of the scoring source. Requesting a +single strategy also writes the model's historical bare column, so existing +commands keep producing the files they used to. + +## Tests + +`tests/test_masked_lm.py` runs on CPU against a stand-in model and needs no +checkpoint. It compares every strategy with a per-variant reference +implementation, checks that the four agree on single mutants and differ on +multi-mutants, and pins each adapter's alphabet, column, batching and dtype +defaults, which are what the released predictions were produced with. diff --git a/fitness/baselines/masked_lm/__init__.py b/fitness/baselines/masked_lm/__init__.py new file mode 100644 index 0000000..a8a3b51 --- /dev/null +++ b/fitness/baselines/masked_lm/__init__.py @@ -0,0 +1,35 @@ +""" +Shared masked-marginal scoring for the benchmark's masked RNA language models. + +The four fill strategies live in ``strategies``, batched inference in +``engine``, the model interface in ``adapter``, and the command line entry +point in ``runner``. A model's scoring script is an adapter plus a call to +``runner.main``. +""" + +from .adapter import MaskedLMAdapter +from .engine import accumulate_scores, window_contexts +from .runner import main +from .strategies import ( + MASK_CHAR, + STRATEGIES, + TaskTable, + build_tasks, + normalize_strategy, + parse_mutations, + recover_wild_type, +) + +__all__ = [ + "MASK_CHAR", + "STRATEGIES", + "MaskedLMAdapter", + "TaskTable", + "accumulate_scores", + "build_tasks", + "main", + "normalize_strategy", + "parse_mutations", + "recover_wild_type", + "window_contexts", +] diff --git a/fitness/baselines/masked_lm/adapter.py b/fitness/baselines/masked_lm/adapter.py new file mode 100644 index 0000000..aafabe2 --- /dev/null +++ b/fitness/baselines/masked_lm/adapter.py @@ -0,0 +1,145 @@ +""" +The model-specific half of masked-marginal scoring. + +An adapter supplies an alphabet, a tokenization, and one forward pass. The +engine supplies everything else, so adding a masked RNA language model to the +benchmark means writing the subclass below and nothing more. +""" + + +class MaskedLMAdapter: + """ + Interface between a masked language model and the scoring engine. + + Class attributes: + name: model name, for log messages. + bases: the four bases in the model's own alphabet, ``"ACGU"`` or + ``"ACGT"``. Assay sequences and mutation strings are folded to it. + score_column: stem of the output column, kept per model so that the + prediction files stay drop-in replacements for the released ones. + max_tokens: position limit including special tokens, or None if the + model has no fixed limit. When set, a ``--max_tokens`` flag is + offered and contexts longer than the limit are windowed. + default_batch_size, default_max_batch_tokens: batching defaults. + + Instance attributes, set by ``load``: + device, prefix_ids, suffix_ids, pad_id, mask_id, unk_id, base_ids. + """ + + name = "masked LM" + bases = "ACGU" + score_column = "score" + max_tokens = None + default_batch_size = 512 + default_max_batch_tokens = 65536 + # Whether contexts of different lengths may share a padded batch. All four + # current models are scored on fixed-length constructs, and RNA-FM and + # RiNALMo are run without an explicit attention mask, so this is False + # unless a model is known to be padding invariant. + allows_mixed_length_batches = False + requires_cuda = False + # Number of special tokens the encoding adds, declared statically so that + # the window-limit guard can run before the model is loaded. Checked against + # the loaded tokenizer in ``check_alphabet``. + n_special_tokens = 0 + # Whether context positions map to token positions by a constant shift, which + # is what lets the engine gather without calling ``token_position`` per term. + # An adapter that overrides ``token_position`` non-linearly must clear this. + constant_token_offset = True + + prefix_ids = () + suffix_ids = () + pad_id = 0 + mask_id = 0 + unk_id = 0 + base_ids = {} + device = "cpu" + + @staticmethod + def add_arguments(parser): + """Add the model's own command line arguments, such as a checkpoint.""" + raise NotImplementedError + + def load(self, args): + """Construct the model and tokenizer and fill in the token ids.""" + raise NotImplementedError + + def logits_at(self, input_ids, attention_mask, rows, cols): + """ + Run one batch and return the raw logits at the requested positions. + + Args: + input_ids: LongTensor ``(B, L)`` already on the device, padded with + ``pad_id`` and carrying the model's special tokens. + attention_mask: LongTensor ``(B, L)``, 1 on real tokens. Models that + do not take an attention mask ignore it. + rows, cols: LongTensor ``(N,)`` each, indexing the positions to read. + A context can appear more than once, once per masked position. + + Returns: + Tensor ``(N, vocab)`` of logits. The engine casts to float32 and + takes the log-softmax, so returning raw logits keeps the numerics + identical across models. + """ + raise NotImplementedError + + def canonicalize_sequence(self, sequence: str) -> str: + """ + Uppercase, strip whitespace, and fold to the model's own alphabet. + + RNA-FM and RNAGenesis read U, RiNALMo and AIDO.RNA read T, and a + mismatch puts near-zero probability on every allele, so the folding + direction is derived from ``bases`` rather than written out per model. + Symbols outside the alphabet, such as N, are left alone and are encoded + as the unknown token. + """ + sequence = str(sequence).strip().upper() + if "T" in self.bases: + return sequence.replace("U", "T") + return sequence.replace("T", "U") + + def encode_context(self, context: str, mask_char: str) -> list: + """ + Turn a masked context string into model input ids. + + The default is one token per character between the model's leading and + trailing special tokens, which is what all four single-nucleotide models + do. A model whose tokenizer does anything else overrides this together + with ``token_position``. + """ + ids = list(self.prefix_ids) + for char in context: + ids.append(self.mask_id if char == mask_char else self.base_ids.get(char, self.unk_id)) + ids.extend(self.suffix_ids) + return ids + + def token_position(self, context_position: int) -> int: + """ + Map a position in the context string to a position in the input ids. + + With one token per nucleotide this is just the number of leading special + tokens, but the engine goes through this method rather than assuming it, + since an off-by-one here reads a neighbouring nucleotide's distribution + and still produces plausible scores. + """ + return len(self.prefix_ids) + context_position + + def check_alphabet(self): + """Fail loudly if the tokenizer does not cover the model's alphabet.""" + n_special = len(self.prefix_ids) + len(self.suffix_ids) + if n_special != self.n_special_tokens: + raise ValueError( + f"{self.name} declares {self.n_special_tokens} special tokens but " + f"its loaded tokenizer adds {n_special}. The window-limit guard " + "runs before loading and would use the wrong budget" + ) + missing = [b for b in self.bases if b not in self.base_ids] + if missing: + raise ValueError(f"Tokenizer does not cover {missing} of {self.bases}") + if self.unk_id in self.base_ids.values(): + raise ValueError( + f"A base maps to the unknown token id {self.unk_id}: {self.base_ids}. " + "This is the RiNALMo lookup bug; the released RiNALMo scores carry it." + ) + if self.mask_id in self.base_ids.values(): + raise ValueError(f"A base maps to the mask token id: {self.base_ids}") diff --git a/fitness/baselines/masked_lm/engine.py b/fitness/baselines/masked_lm/engine.py new file mode 100644 index 0000000..c8ed23f --- /dev/null +++ b/fitness/baselines/masked_lm/engine.py @@ -0,0 +1,194 @@ +""" +Batched inference over a deduplicated context bank. + +The engine owns everything that is the same for every masked RNA language model: +padding a batch of masked contexts into token ids, running the model, taking +log-softmax at the masked positions, and accumulating the signed log-probability +terms into per-variant scores. Everything model-specific lives behind the +``MaskedLMAdapter`` interface in ``adapter.py``. + +A context may carry more than one mask, because ``mask-fill`` masks a variant's +whole mutated set at once, so masked positions are read from the task table +rather than inferred from the context string. +""" + +import numpy as np +import torch +from tqdm.auto import tqdm + +from .strategies import MASK_CHAR, validate_table + + +def window_contexts(table, budget: int) -> None: + """ + Trim contexts to a model's position limit, centred on their masked span. + + RNA-FM and RiNALMo accept a bounded number of positions. The non-coding + constructs used for the four-strategy comparison are 45 to 425 nucleotides + and never reach it, but the mRNA-coding constructs are kilobases, so a + window centred on the masked span is taken, mirroring the windowing in the + baselines' original ``compute_fitness.py``. Mutations outside the window are + lost from the context, which is inherent to windowing. + + Modifies ``table`` in place: contexts are rewritten and positions shifted. + + Raises: + ValueError: if a context's masked span itself does not fit in the + budget, since there is then no window that shows every scored position. + """ + if all(len(c) <= budget for c in table.contexts): + return + + # Masked span per context, from the task table rather than from the string. + # The table is sorted by context, so the spans come from one grouped reduce + # rather than a pass over every term, which matters on the mRNA-coding + # assays where this path is reached at all. + span_lo = {} + span_hi = {} + if table.n_terms(): + used, first = np.unique(table.ctx_id, return_index=True) + lows = np.minimum.reduceat(table.pos, first) + highs = np.maximum.reduceat(table.pos, first) + span_lo = dict(zip(used.tolist(), lows.tolist())) + span_hi = dict(zip(used.tolist(), highs.tolist())) + + shift = np.zeros(len(table.contexts), dtype=np.int64) + trimmed = list(table.contexts) + for k, context in enumerate(table.contexts): + if len(context) <= budget: + continue + if k not in span_lo: + trimmed[k] = context[:budget] + continue + low, high = span_lo[k], span_hi[k] + span = high - low + 1 + if span > budget: + raise ValueError( + f"Masked span of {span} positions does not fit in a window of " + f"{budget}; this context cannot be scored" + ) + centre = (low + high) // 2 + start = max(0, centre - budget // 2) + end = min(len(context), start + budget) + start = max(0, end - budget) + trimmed[k] = context[start:end] + shift[k] = start + + table.contexts = trimmed + if shift.any(): + table.pos = (table.pos - shift[table.ctx_id]).astype(np.int32) + validate_table(table) + + +def accumulate_scores( + adapter, + table, + n_rows: int, + batch_size: int, + max_batch_tokens: int, + progress: bool = True, +) -> np.ndarray: + """ + Run every context once and accumulate the task table into per-variant scores. + + Args: + adapter: a loaded MaskedLMAdapter. + table: the TaskTable built by ``strategies.build_tasks``. + n_rows: number of rows in the assay dataframe. + batch_size: maximum contexts per forward pass. + max_batch_tokens: cap on ``batch_size x sequence length``, so that long + assays automatically use a smaller batch. + progress: show a progress bar. + + Returns: + Array of shape ``(len(table.strategies), n_rows)``. Rows that could not + be scored are NaN under every strategy. + """ + contexts = table.contexts + if not contexts: + raise ValueError("No scorable variants found") + + n_special = len(adapter.prefix_ids) + len(adapter.suffix_ids) + # The gather is vectorized, so it needs the context-to-token map to be a + # constant shift. The adapter declares that, and the declaration is checked + # here: an off-by-one reads a neighbouring nucleotide's distribution and + # still produces plausible finite scores. + if not adapter.constant_token_offset: + raise ValueError( + f"{adapter.name} does not declare a constant context-to-token offset, " + "so the engine cannot gather its logits by a vectorized shift" + ) + offset = adapter.token_position(0) + longest = max(len(c) for c in contexts) + if any(adapter.token_position(p) != offset + p for p in (0, 1, longest - 1)): + raise ValueError( + f"{adapter.name} declares a constant context-to-token offset but does " + "not implement one" + ) + vocab_of_base = np.array( + [adapter.base_ids[b] for b in adapter.bases], dtype=np.int64 + ) + + seq_len = max(len(c) for c in contexts) + batch_size = max(1, min(batch_size, max_batch_tokens // (seq_len + n_special))) + print(f"Using batch size {batch_size} for sequence length {seq_len}") + + scores = np.full((len(table.strategies), n_rows), np.nan, dtype=float) + scores[:, table.scorable] = 0.0 + flat = scores.reshape(-1) + + device = adapter.device + pad_id = adapter.pad_id + + for start in tqdm( + range(0, len(contexts), batch_size), + desc="Scoring", + unit="batch", + disable=not progress, + ): + batch = contexts[start : start + batch_size] + max_len = max(len(c) for c in batch) + n_special + + if not adapter.allows_mixed_length_batches and len({len(c) for c in batch}) > 1: + raise ValueError( + f"{adapter.name} is not declared padding invariant, so contexts of " + "different lengths must not share a batch" + ) + + input_ids = torch.full((len(batch), max_len), pad_id, dtype=torch.long) + attention_mask = torch.zeros((len(batch), max_len), dtype=torch.long) + for b, context in enumerate(batch): + ids = adapter.encode_context(context, MASK_CHAR) + input_ids[b, : len(ids)] = torch.tensor(ids, dtype=torch.long) + attention_mask[b, : len(ids)] = 1 + + # The task table is sorted by context, so this batch's terms are one slice. + lo = int(np.searchsorted(table.ctx_id, start, "left")) + hi = int(np.searchsorted(table.ctx_id, start + len(batch), "left")) + if lo == hi: + continue + local_ctx = table.ctx_id[lo:hi].astype(np.int64) - start + key = local_ctx * max_len + table.pos[lo:hi] + uniq, inverse = np.unique(key, return_inverse=True) + # (context, position) is the identity of a gathered distribution: one + # context may be read at several positions, because mask-fill masks a + # variant's whole mutated set in a single context. + rows = torch.from_numpy((uniq // max_len).astype(np.int64)).to(device) + cols = torch.from_numpy((uniq % max_len).astype(np.int64) + offset).to(device) + + input_ids = input_ids.to(device) + attention_mask = attention_mask.to(device) + with torch.inference_mode(): + logits = adapter.logits_at(input_ids, attention_mask, rows, cols) + log_probs = torch.log_softmax(logits.float(), dim=-1) + log_probs = log_probs.cpu().numpy() + + vocab = vocab_of_base[table.base[lo:hi]] + values = log_probs[inverse, vocab] * table.sign[lo:hi] + np.add.at( + flat, + table.strategy[lo:hi].astype(np.int64) * n_rows + table.row[lo:hi], + values, + ) + + return scores diff --git a/fitness/baselines/masked_lm/runner.py b/fitness/baselines/masked_lm/runner.py new file mode 100644 index 0000000..fcb315a --- /dev/null +++ b/fitness/baselines/masked_lm/runner.py @@ -0,0 +1,430 @@ +""" +Command line entry point shared by the masked language model scorers. + +Handles the reference sheet and assay input, the wild-type cross-check, the +strategy selection, and the output file, so that a model's script is only an +adapter. +""" + +import argparse +import hashlib +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy.stats import spearmanr + +from .engine import accumulate_scores, window_contexts +from .strategies import STRATEGIES, build_tasks, normalize_strategy, recover_wild_type + + +def build_parser(adapter) -> argparse.ArgumentParser: + """Assemble the common arguments plus the adapter's own.""" + parser = argparse.ArgumentParser( + description=f"Run {adapter.name} masked-marginal inference on DMS assay sequences." + ) + parser.add_argument( + "--row_id", + type=int, + required=True, + help="Row ID in the reference sheet to process", + ) + parser.add_argument( + "--ref_sheet", + type=str, + required=True, + help="Path to reference sheet containing DMS_ID and RAW_CONSTRUCT_SEQ columns", + ) + parser.add_argument( + "--dms_dir_path", + type=str, + required=True, + help="Directory containing DMS CSV files", + ) + parser.add_argument( + "--output_dir_path", + type=str, + required=True, + help="Directory to save output files", + ) + parser.add_argument( + "--strategies", + type=str, + nargs="+", + default=list(STRATEGIES), + help="Masked-marginal fill strategies to compute. The fill named is what " + "the model sees at the variant's OTHER mutated positions while one " + "position is masked: wt-fill (wild-type bases, what the ESM and " + "ProteinGym code implement), mask-fill (masks, the ESM paper's " + "formula), mut-fill (mutant bases), match-fill (the allele being " + "scored). All four agree on single mutants. Default: all four, which " + "costs about 19%% more unique context examples than mut-fill alone, " + "because the contexts are shared", + ) + parser.add_argument( + "--legacy_column", + type=str, + default=None, + choices=[s.replace("_", "-") for s in STRATEGIES], + help="Also write the model's historical bare score column, holding this " + "strategy's scores. Defaults to the single requested strategy when " + "exactly one is requested, and to nothing otherwise", + ) + parser.add_argument( + "--device", + type=str, + default=None, + help="Device to run inference on (default: cuda:0 if available, else cpu)", + ) + parser.add_argument( + "--batch_size", + type=int, + default=adapter.default_batch_size, + help=f"Maximum masked contexts scored per forward pass (default: {adapter.default_batch_size})", + ) + parser.add_argument( + "--max_batch_tokens", + type=int, + default=adapter.default_max_batch_tokens, + help="Cap on batch_size x sequence length per forward pass, so that long " + f"assays automatically use a smaller batch (default: {adapter.default_max_batch_tokens})", + ) + if adapter.max_tokens is not None: + parser.add_argument( + "--max_tokens", + type=int, + default=adapter.max_tokens, + help=f"{adapter.name} position limit including special tokens. Contexts " + f"longer than this are windowed around the masked span (default: {adapter.max_tokens})", + ) + parser.add_argument( + "--quiet_skips", + action="store_true", + help="Summarize skipped variants instead of printing one line each", + ) + adapter.add_arguments(parser) + return parser + + +def load_reference_row(ref_sheet_path: str, row_id: int): + """ + Read one row of the reference sheet. + + Returns: + (dms_id, raw_construct_seq). The construct sequence is None when the + sheet does not carry one, in which case the wild type comes from the + assay alone. + + Raises: + ValueError: If row_id is not found or DMS_ID is missing + """ + try: + ref_df = pd.read_csv(ref_sheet_path, encoding="utf-8-sig") + except FileNotFoundError: + raise FileNotFoundError(f"Reference sheet not found: {ref_sheet_path}") + if "DMS_ID" not in ref_df.columns: + raise KeyError("Reference sheet must contain 'DMS_ID' column") + if row_id >= len(ref_df): + raise ValueError(f"Row ID {row_id} exceeds number of rows in reference sheet") + + dms_id = ref_df.loc[row_id, "DMS_ID"] + if pd.isna(dms_id): + raise ValueError(f"DMS_ID is missing for row {row_id}") + + construct = None + if "RAW_CONSTRUCT_SEQ" in ref_df.columns: + value = ref_df.loc[row_id, "RAW_CONSTRUCT_SEQ"] + if not pd.isna(value): + construct = str(value) + return str(dms_id), construct + + +def load_dms_data(dms_dir_path: str, dms_id: str) -> pd.DataFrame: + """ + Load DMS data for specified DMS_ID. + + Raises: + FileNotFoundError: If DMS file is not found + """ + dms_file = Path(dms_dir_path) / f"{dms_id}.csv" + if not dms_file.exists(): + raise FileNotFoundError(f"DMS file not found: {dms_file}") + + df = pd.read_csv(dms_file) + required_cols = ["mutant", "DMS_score", "sequence"] + missing_cols = [col for col in required_cols if col not in df.columns] + if missing_cols: + raise ValueError(f"Missing required columns in DMS file: {missing_cols}") + + return df + + +def resolve_wild_type(adapter, mutants, sequences, construct) -> str: + """ + Determine the assay's wild type and cross-check the two independent sources. + + The wild type is reconstructed from the assay itself, by reverting each + variant's own mutations, which uses no outside information and therefore + catches a coordinate mismatch between the reference sheet and the assay + file. It is then compared with the reference sheet's RAW_CONSTRUCT_SEQ. + + Raises: + ValueError: if the two disagree, since the wild-type-background + strategies would otherwise be scored against the wrong background. + """ + recovered = recover_wild_type(mutants, sequences, adapter.bases) + if construct is None: + print("Reference sheet has no RAW_CONSTRUCT_SEQ; using the recovered wild type") + return recovered + folded = adapter.canonicalize_sequence(construct) + if folded != recovered: + raise ValueError( + "The wild type recovered from the assay disagrees with the reference " + f"sheet's RAW_CONSTRUCT_SEQ (lengths {len(recovered)} and {len(folded)}). " + "The wild-type-background strategies need one agreed background" + ) + return recovered + + +def window_budget(adapter, max_tokens: int) -> int: + """ + How many nucleotides fit in one context, given the model's position limit. + + Uses the adapter's DECLARED special-token count rather than its loaded token + ids, so that an unsupported request is refused before a model is loaded. + ``check_alphabet`` verifies the declaration against the real tokenizer. + """ + return max_tokens - adapter.n_special_tokens + + +def needs_windowing(contexts, budget: int) -> bool: + """Whether any context is too long for the model's position limit.""" + return any(len(c) > budget for c in contexts) + + +def code_revision() -> dict: + """ + Identify the code that produced the scores. + + The git revision alone is misleading while the tree is dirty, which it is + during development, so the scoring source itself is hashed as well: the + shared package plus the model script that was invoked. That hash identifies + the scoring code exactly whether or not it has been committed. + """ + here = Path(__file__).resolve().parent + revision, dirty = "unknown", None + try: + revision = subprocess.run( + ["git", "-C", str(here), "rev-parse", "HEAD"], + capture_output=True, text=True, check=True, + ).stdout.strip() + status = subprocess.run( + ["git", "-C", str(here), "status", "--porcelain"], + capture_output=True, text=True, check=True, + ).stdout.strip() + dirty = bool(status) + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + digest = hashlib.sha256() + sources = sorted(here.glob("*.py")) + entry = Path(sys.argv[0]).resolve() + if entry.is_file(): + sources.append(entry) + for path in sources: + digest.update(path.name.encode()) + digest.update(path.read_bytes()) + return { + "git_head": revision, + "git_tree_dirty": dirty, + "scoring_source_sha256": digest.hexdigest(), + "scoring_source_files": [p.name for p in sources], + } + + +def runtime_environment() -> dict: + """ + Record what the scores were computed on. + + bfloat16 results depend on the GPU: AIDO.RNA-1.6B scored in bf16 on an L40S + and on an H100 agrees only to about 0.2 in score and 3e-4 in Spearman, which + is invisible unless the hardware is written down. + """ + import torch + + device_name = None + try: + if torch.cuda.is_available(): + device_name = torch.cuda.get_device_name() + except (AssertionError, RuntimeError): + pass + return { + "torch": torch.__version__, + "cuda": getattr(torch.version, "cuda", None), + "gpu": device_name, + } + + +def write_manifest(output_dir, dms_id, adapter, args, strategies, table, wild_type): + """ + Record how a prediction file was produced, next to the file itself. + + Which fill strategy a column holds, which alphabet the sequences were folded + to, and which checkpoint and dtype produced them are exactly the details + that were unrecoverable for the benchmark's earlier published predictions. + """ + manifest = { + "dms_id": dms_id, + "environment": runtime_environment(), + "model": adapter.name, + "score_column_stem": adapter.score_column, + "strategies": [s.replace("_", "-") for s in strategies], + "columns": { + s.replace("_", "-"): f"{adapter.score_column}_{s}" for s in strategies + }, + "alphabet": adapter.bases, + "wild_type_length": len(wild_type), + "contexts": len(table.contexts), + "terms": table.n_terms(), + "scorable_variants": int(table.scorable.sum()), + "code_revision": code_revision(), + "arguments": { + k: (str(v) if isinstance(v, Path) else v) for k, v in vars(args).items() + }, + } + path = Path(output_dir) / f"{dms_id}.manifest.json" + path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + + +def main(adapter): + """Run one assay with one model under the requested strategies.""" + parser = build_parser(adapter) + args = parser.parse_args() + + strategies = tuple(normalize_strategy(s) for s in args.strategies) + if len(set(strategies)) != len(strategies): + parser.error(f"Duplicate strategies requested: {args.strategies}") + legacy = args.legacy_column and normalize_strategy(args.legacy_column) + if legacy is None and len(strategies) == 1: + legacy = strategies[0] + if legacy is not None and legacy not in strategies: + parser.error(f"--legacy_column {args.legacy_column} was not requested") + + output_dir = Path(args.output_dir_path) + output_dir.mkdir(parents=True, exist_ok=True) + + try: + dms_id, construct = load_reference_row(args.ref_sheet, args.row_id) + print(f"Processing DMS ID: {dms_id}") + dms_df = load_dms_data(args.dms_dir_path, dms_id) + + print(f"Preprocessing sequences into the {adapter.bases} alphabet...") + sequences = [adapter.canonicalize_sequence(s) for s in dms_df["sequence"].tolist()] + mutants = dms_df["mutant"].tolist() + wild_type = resolve_wild_type(adapter, mutants, sequences, construct) + print(f"Wild type: {len(wild_type)} nt, cross-checked against the assay") + + print(f"Building contexts for: {', '.join(s.replace('_', '-') for s in strategies)}") + table = build_tasks( + mutants, + sequences, + wild_type, + adapter.bases, + strategies, + verbose=not args.quiet_skips, + ) + n_scorable = int(table.scorable.sum()) + print( + f"Prepared {table.n_terms()} log-probability terms across {n_scorable} " + f"scorable variants (of {len(sequences)} total), deduplicated to " + f"{len(table.contexts)} unique contexts" + ) + if not table.contexts: + raise ValueError("No scorable variants found") + + max_tokens = getattr(args, "max_tokens", None) + if max_tokens is not None: + # The declared count, not the loaded one: this guard runs before the + # model is loaded so that an unsupported request fails cheaply. + budget = window_budget(adapter, max_tokens) + if needs_windowing(table.contexts, budget): + if len(strategies) > 1: + raise ValueError( + f"Contexts exceed the {max_tokens} position limit and more " + "than one strategy was requested. A windowed wild-type " + "table and a windowed variant context are not in the same " + "coordinate frame, so the strategies would not be " + "comparable. Request one strategy at a time for this assay" + ) + print(f"Windowing contexts to {budget} positions around the masked span") + + device = args.device + if device is None: + import torch + + device = "cuda:0" if torch.cuda.is_available() else "cpu" + if adapter.requires_cuda and not device.startswith("cuda"): + raise ValueError( + f"{adapter.name} has no working CPU path; pass a CUDA device" + ) + args.device = device + print(f"Initializing {adapter.name} on {device}...") + adapter.load(args) + adapter.check_alphabet() + + if max_tokens is not None: + window_contexts(table, window_budget(adapter, max_tokens)) + + print("Running inference...") + scores = accumulate_scores( + adapter, + table, + n_rows=len(dms_df), + batch_size=args.batch_size, + max_batch_tokens=args.max_batch_tokens, + ) + + for i, strategy in enumerate(strategies): + dms_df[f"{adapter.score_column}_{strategy}"] = scores[i] + if legacy is not None: + dms_df[adapter.score_column] = scores[strategies.index(legacy)] + + output_file = output_dir / f"{dms_id}.csv" + dms_df.to_csv(output_file, index=False) + write_manifest(output_dir, dms_id, adapter, args, strategies, table, wild_type) + print(f"Saved results to: {output_file}") + + print("\nSummary:") + print(f"Number of sequences: {len(sequences)}") + for i, strategy in enumerate(strategies): + correlation, pvalue = spearmanr( + dms_df["DMS_score"], scores[i], nan_policy="omit" + ) + print( + f" {strategy.replace('_', '-'):10s} Spearman {correlation:+.4f} " + f"(p {pvalue:.2e})" + ) + if len(strategies) > 1: + finite = np.isfinite(scores).all(axis=0) + identical = np.allclose(scores[:, finite], scores[0, finite]) + n_multi = 0 + for mutant in np.asarray(mutants)[table.scorable]: + if str(mutant).count(",") > 0: + n_multi += 1 + if n_multi == 0 and not identical: + raise ValueError( + "Every variant is a single mutant, so the strategies must be " + "identical, but they are not. This is a bug" + ) + print( + f" strategies identical: {identical} " + f"({n_multi} multi-mutants among {n_scorable} scorable variants)" + ) + print(f"Output saved to: {output_file}") + + except Exception as err: + print(f"Error: {err}", file=sys.stderr) + sys.exit(1) diff --git a/fitness/baselines/masked_lm/strategies.py b/fitness/baselines/masked_lm/strategies.py new file mode 100644 index 0000000..4c0b4c9 --- /dev/null +++ b/fitness/baselines/masked_lm/strategies.py @@ -0,0 +1,469 @@ +""" +The four masked-marginal fill strategies, expressed as contexts and tasks. + +Every masked-marginal score computes a log-odds at each mutated position of a +variant and sums over the variant's mutations. The strategies differ in exactly +one thing: what fills the variant's OTHER mutated positions while the scored +position is masked. Following Meier et al. 2021 (ESM-1v) supplement, Appendix A, +with ``M`` the mutated positions, ``x^wt`` and ``x^mt`` the wild-type and variant +sequences, ``x_-i`` a mask at position ``i`` and ``x_-M`` masks at every position +in ``M``: + + wt-fill sum_i [ log p(mt_i | x^wt_-i) - log p(wt_i | x^wt_-i) ] + mask-fill sum_i [ log p(mt_i | x^wt_-M) - log p(wt_i | x^wt_-M) ] + mut-fill sum_i [ log p(mt_i | x^mt_-i) - log p(wt_i | x^mt_-i) ] + match-fill sum_i [ log p(mt_i | x^mt_-i) - log p(wt_i | x^wt_-i) ] + +``wt-fill`` is what the ESM authors' released code and ProteinGym's baseline +implement under the name ``masked-marginals``; ``mask-fill`` is the formula +written in the ESM paper; ``mut-fill`` and ``match-fill`` are strategies c and b +of the supplement. All four are identical on single mutants, because a variant +with one mutation has no other mutated positions and ``x^mt_-i`` equals +``x^wt_-i``. They diverge on multi-mutants, which are 99.4% of this benchmark's +non-coding variants. + +Only ``match-fill`` mixes two contexts: its mutant term is conditioned on the +variant's own sequence and its wild-type term on the wild type, so it is a +difference of two conditionals rather than a log-odds ratio. + +This module turns an assay into two things: + + contexts deduplicated strings over the model's alphabet plus a mask + placeholder, one per unique context example + tasks a flat table of (context, position, base, row, strategy, sign) + terms to accumulate, as parallel arrays + +Contexts are shared across strategies wherever they coincide, which is what +makes computing all four cost only about 19% more unique context examples than +computing ``mut-fill`` alone on the non-coding assays. These are context +examples, not model invocations: contexts are batched, and a context carrying +several masks is read at several positions from one pass. + +Coordinates +----------- + +Three coordinate systems appear and must not be confused: + + full position 0-based index into the assay's sequences, which is what + a mutation string names (1-based) and what this module + stores in ``TaskTable.pos`` + context position index into a context string. Equal to the full position + unless the context was windowed, in which case + ``engine.window_contexts`` subtracts the window origin + from ``TaskTable.pos`` so that the stored value is + always a context position + token position index into the model's input ids, which the adapter + computes from the context position, normally by adding + the number of leading special tokens + +Every task position must point at a mask in its context; ``validate_table`` +checks this, because a position that is off by one gathers a neighbouring +nucleotide's distribution and still produces plausible finite scores. +""" + +from array import array +from dataclasses import dataclass + +import numpy as np +import pandas as pd + +MASK_CHAR = "#" + +# Canonical order. The CLI accepts the hyphenated spelling used in the write-up; +# the underscored spelling is used for identifiers and output column suffixes. +STRATEGIES = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + + +def normalize_strategy(name: str) -> str: + """Accept either the hyphenated or the underscored spelling of a strategy.""" + key = name.strip().lower().replace("-", "_") + if key not in STRATEGIES: + raise ValueError(f"Unknown strategy {name!r}; expected one of {STRATEGIES}") + return key + + +@dataclass +class TaskTable: + """ + A deduplicated set of masked contexts plus the terms to read out of them. + + Attributes: + contexts: masked context strings, deduplicated. A context may + carry more than one mask (``mask-fill`` masks a variant's whole + mutated set at once), so a context does not identify the position + being scored and ``pos`` is carried explicitly. + ctx_id, pos, base, row, strategy, sign: parallel arrays, one entry per + term to accumulate. Term ``k`` adds + ``sign[k] * log p(base[k] | contexts[ctx_id[k]])`` evaluated at + position ``pos[k]`` into the score of variant ``row[k]`` under + strategy ``strategy[k]``. Sorted by ``ctx_id`` so that each + context's terms are a contiguous slice. + scorable: per-variant mask; False rows are reported as NaN under every + strategy. + strategies: the strategy names indexed by ``strategy``. + """ + + contexts: list + ctx_id: np.ndarray + pos: np.ndarray + base: np.ndarray + row: np.ndarray + strategy: np.ndarray + sign: np.ndarray + scorable: np.ndarray + strategies: tuple + + def n_terms(self) -> int: + return int(self.ctx_id.size) + + +def parse_mutations(mutant_str: str, bases: str) -> list: + """ + Parse a mutation string such as ``"A4U,A5G"`` into a list of + ``(pos0, wt_base, mut_base)`` tuples with 0-based positions, in the alphabet + the model uses. + + Raises: + ValueError: for non-substitution edits (indels) or bases outside the + model's alphabet, so the caller can score the affected variant as NaN. + """ + fold_from, fold_to = ("U", "T") if "T" in bases else ("T", "U") + mutations = [] + for token in str(mutant_str).replace(" ", "").split(","): + if not token: + continue + wt_base = token[0].upper().replace(fold_from, fold_to) + mut_base = token[-1].upper().replace(fold_from, fold_to) + pos = int(token[1:-1]) - 1 # 1-based -> 0-based + if wt_base not in bases or mut_base not in bases: + raise ValueError(f"Unsupported mutation token: {token}") + mutations.append((pos, wt_base, mut_base)) + return mutations + + +def recover_wild_type(mutants, sequences, bases: str) -> str: + """ + Recover the assay's wild-type sequence by reverting each variant's own + mutations, and require that every variant agrees. + + The wild-type-background strategies need a wild type, and taking it from the + reference sheet alone would not catch a coordinate mismatch between the + sheet and the assay file. Reverting the mutations uses only the assay, so + the two are independent and can be cross-checked by the caller. + + Returns: + The single wild-type sequence implied by the assay. + + Raises: + ValueError: if the variants do not all imply the same wild type, which + would mean the assay mixes backgrounds or the mutation coordinates do + not line up with the sequences. + """ + candidates = {} + for mutant_str, seq in zip(mutants, sequences): + if pd.isna(mutant_str): + continue + try: + mutations = parse_mutations(mutant_str, bases) + except ValueError: + continue + if not mutations: + continue + reverted = list(seq) + ok = True + for pos, wt_base, mut_base in mutations: + if pos < 0 or pos >= len(seq) or seq[pos] != mut_base: + ok = False + break + reverted[pos] = wt_base + if ok: + candidates.setdefault("".join(reverted), 0) + candidates["".join(reverted)] += 1 + if not candidates: + raise ValueError("No variant could be reverted to a wild-type sequence") + if len(candidates) > 1: + top = sorted(candidates.items(), key=lambda kv: -kv[1])[:3] + raise ValueError( + "Variants imply more than one wild-type sequence " + f"({len(candidates)} distinct; top counts {[c for _, c in top]})" + ) + return next(iter(candidates)) + + +def difference_counts(sequences, wild_type: str) -> np.ndarray: + """ + Count, for every variant, how many positions differ from the wild type. + + Used to verify that a variant differs from the wild type at exactly its + declared mutated positions and nowhere else. Done as one array comparison + rather than per row, since the assays run to hundreds of thousands of + variants. + + Returns: + int array with one entry per sequence, or -1 where the sequence length + differs from the wild type's and the comparison is undefined. + """ + length = len(wild_type) + counts = np.full(len(sequences), -1, dtype=np.int64) + same_length = np.array([len(s) == length for s in sequences], dtype=bool) + if not same_length.any(): + return counts + packed = np.frombuffer( + "".join(s for s, keep in zip(sequences, same_length) if keep).encode("ascii"), + dtype=np.uint8, + ).reshape(-1, length) + reference = np.frombuffer(wild_type.encode("ascii"), dtype=np.uint8) + counts[same_length] = (packed != reference).sum(axis=1) + return counts + + +def validate_table(table) -> None: + """ + Check the task table's invariants: every scored position is masked in the + context it is read from. + + A position that is off by one, or that survived windowing incorrectly, still + yields finite and plausible scores, so this is checked rather than assumed. + + Raises: + ValueError: if any task points at a position that is not masked. + """ + if table.n_terms() == 0: + return + lengths = {len(c) for c in table.contexts} + if len(lengths) == 1: + # The usual case: one construct length per assay, so the whole bank + # packs into an array and the check is one comparison. + length = lengths.pop() + packed = np.frombuffer( + "".join(table.contexts).encode("ascii"), dtype=np.uint8 + ).reshape(len(table.contexts), length) + if table.pos.min() < 0 or table.pos.max() >= length: + raise ValueError("A task position falls outside its context") + bad = packed[table.ctx_id, table.pos] != ord(MASK_CHAR) + if bad.any(): + k = int(np.flatnonzero(bad)[0]) + raise ValueError( + f"Task {k} reads position {table.pos[k]} of context " + f"{table.ctx_id[k]}, which is not a mask" + ) + return + stride = max(lengths) + 1 + for packed_key in np.unique(table.ctx_id.astype(np.int64) * stride + table.pos): + ctx, pos = int(packed_key // stride), int(packed_key % stride) + context = table.contexts[ctx] + if pos >= len(context) or context[pos] != MASK_CHAR: + raise ValueError( + f"Task reads position {pos} of context {ctx}, which is not a mask" + ) + + +class _ContextBank: + """Deduplicates context strings and hands out their indices.""" + + def __init__(self): + self._index = {} + self.contexts = [] + + def add(self, context: str) -> int: + key = self._index.get(context) + if key is None: + key = len(self.contexts) + self._index[context] = key + self.contexts.append(context) + return key + + +def _single_mask(seq: str, pos: int) -> str: + return seq[:pos] + MASK_CHAR + seq[pos + 1 :] + + +def _multi_mask(seq: str, positions) -> str: + chars = list(seq) + for pos in positions: + chars[pos] = MASK_CHAR + return "".join(chars) + + +def build_tasks( + mutants, + sequences, + wild_type: str, + bases: str, + strategies=STRATEGIES, + verbose: bool = True, +) -> TaskTable: + """ + Expand an assay's variants into deduplicated masked contexts and the terms + read out of them, for every requested strategy at once. + + A variant is scorable only if it is scorable under EVERY requested strategy, + so that the strategies are compared on exactly the same set of variants. + + Every scorable variant is checked to satisfy all of: + + the mutation string parses into substitutions over the alphabet; + its positions are inside the sequence and are distinct; + no mutation is a no-op (a mutant base equal to its wild-type base); + the variant sequence carries the mutant base at each mutated position; + the sequence contains no mask placeholder; + + and additionally, when a wild-type-background strategy is requested: + + the sequence has the same length as the wild type; + the wild type carries the declared wild-type base at each position; + the sequence equals the wild type at every position outside the mutated + set, so that the declared mutations describe the variant completely. + + Anything else is reported as NaN under every strategy, matching the existing + scorers' all-or-nothing policy: a variant with one bad mutation contributes + none of its mutations. On the 31 non-coding assays (856,628 variants) none + of these checks currently rejects anything, so they cost no coverage. + + Args: + mutants: the assay's ``mutant`` column. + sequences: the assay's ``sequence`` column, already folded to the + model's alphabet. + wild_type: the assay's wild-type sequence, in the same alphabet. + bases: the model's alphabet, ``"ACGU"`` or ``"ACGT"``. + strategies: which strategies to build terms for. + verbose: print the reason each skipped variant was skipped. + + Returns: + A TaskTable. + """ + strategies = tuple(normalize_strategy(s) for s in strategies) + if not strategies: + raise ValueError("At least one strategy is required") + strat_id = {name: i for i, name in enumerate(strategies)} + base_id = {b: i for i, b in enumerate(bases)} + needs_wt = any(s in ("wt_fill", "mask_fill", "match_fill") for s in strategies) + + bank = _ContextBank() + ctx_id, pos_a, row_a = array("i"), array("i"), array("i") + base_a, strat_a, sign_a = array("b"), array("b"), array("b") + scorable = np.zeros(len(sequences), dtype=bool) + + # The wild-type single-mask contexts depend only on the position, so they are + # shared by every variant and worth caching. + wt_ctx_cache = {} + + def wt_masked(pos: int) -> int: + key = wt_ctx_cache.get(pos) + if key is None: + key = bank.add(_single_mask(wild_type, pos)) + wt_ctx_cache[pos] = key + return key + + def emit(context_key: int, pos: int, base: str, row: int, strategy: str, sign: int): + ctx_id.append(context_key) + pos_a.append(pos) + base_a.append(base_id[base]) + row_a.append(row) + strat_a.append(strat_id[strategy]) + sign_a.append(sign) + + # One array comparison for the whole assay, so that the per-variant check + # that nothing outside the mutated set differs from the wild type is cheap. + n_differences = difference_counts(sequences, wild_type) if needs_wt else None + + n_skipped = 0 + for i, (mutant_str, seq) in enumerate(zip(mutants, sequences)): + if pd.isna(mutant_str): + continue # wild-type row: leave as NaN + try: + mutations = parse_mutations(mutant_str, bases) + if not mutations: + continue + if MASK_CHAR in seq: + raise ValueError( + f"Sequence contains the mask placeholder {MASK_CHAR!r}, which " + "would be read as a mask and would corrupt context dedup" + ) + positions = [pos for pos, _, _ in mutations] + if len(set(positions)) != len(positions): + raise ValueError(f"Mutation string names a position twice: {mutant_str}") + for pos, wt_base, mut_base in mutations: + if pos < 0 or pos >= len(seq): + raise ValueError(f"Mutation position {pos + 1} outside sequence") + if wt_base == mut_base: + raise ValueError(f"Mutation {wt_base}{pos + 1}{mut_base} is a no-op") + if seq[pos] != mut_base: + raise ValueError( + f"Sequence has {seq[pos]} at position {pos + 1}, " + f"expected the mutant base {mut_base}" + ) + if needs_wt: + if len(seq) != len(wild_type): + raise ValueError( + f"Sequence length {len(seq)} differs from the wild type " + f"({len(wild_type)}), so wild-type-background strategies " + "have no common coordinate frame" + ) + if wild_type[pos] != wt_base: + raise ValueError( + f"Wild type has {wild_type[pos]} at position {pos + 1}, " + f"expected {wt_base}" + ) + if needs_wt and n_differences[i] != len(positions): + raise ValueError( + f"Sequence differs from the wild type at {n_differences[i]} " + f"positions but declares {len(positions)} mutations, so the " + "mutation string does not describe the variant completely" + ) + except (ValueError, IndexError) as err: + n_skipped += 1 + if verbose: + print(f"Skipping variant {mutant_str}: {err}") + continue # unsupported edit: leave as NaN under every strategy + + mut_ctx = {} + if "mut_fill" in strat_id or "match_fill" in strat_id: + for pos in positions: + mut_ctx[pos] = bank.add(_single_mask(seq, pos)) + joint_ctx = None + if "mask_fill" in strat_id: + joint_ctx = bank.add(_multi_mask(wild_type, positions)) + + for pos, wt_base, mut_base in mutations: + if "wt_fill" in strat_id: + key = wt_masked(pos) + emit(key, pos, mut_base, i, "wt_fill", 1) + emit(key, pos, wt_base, i, "wt_fill", -1) + if "mask_fill" in strat_id: + emit(joint_ctx, pos, mut_base, i, "mask_fill", 1) + emit(joint_ctx, pos, wt_base, i, "mask_fill", -1) + if "mut_fill" in strat_id: + emit(mut_ctx[pos], pos, mut_base, i, "mut_fill", 1) + emit(mut_ctx[pos], pos, wt_base, i, "mut_fill", -1) + if "match_fill" in strat_id: + # The mutant term shares mut-fill's context and the wild-type + # term shares wt-fill's, which is why match-fill is free once + # both of those are being computed. + emit(mut_ctx[pos], pos, mut_base, i, "match_fill", 1) + emit(wt_masked(pos), pos, wt_base, i, "match_fill", -1) + scorable[i] = True + + if n_skipped and not verbose: + print(f"Skipped {n_skipped} variants that could not be scored") + + table = TaskTable( + contexts=bank.contexts, + ctx_id=np.frombuffer(ctx_id, dtype=np.int32).copy(), + pos=np.frombuffer(pos_a, dtype=np.int32).copy(), + base=np.frombuffer(base_a, dtype=np.int8).copy(), + row=np.frombuffer(row_a, dtype=np.int32).copy(), + strategy=np.frombuffer(strat_a, dtype=np.int8).copy(), + sign=np.frombuffer(sign_a, dtype=np.int8).copy(), + scorable=scorable, + strategies=strategies, + ) + order = np.argsort(table.ctx_id, kind="stable") + table.ctx_id = table.ctx_id[order] + table.pos = table.pos[order] + table.base = table.base[order] + table.row = table.row[order] + table.strategy = table.strategy[order] + table.sign = table.sign[order] + validate_table(table) + return table diff --git a/fitness/merge_scoring_files.py b/fitness/merge_scoring_files.py index a544dee..7da33c3 100755 --- a/fitness/merge_scoring_files.py +++ b/fitness/merge_scoring_files.py @@ -48,14 +48,20 @@ def combine_csv_data( save = True for model_name in model_list: - model_path = os.path.join(model_predictions_folder, model_name, csv_file) + folder = resolve_source(score_cols_dict, model_name)[0] + model_path = os.path.join(model_predictions_folder, folder, csv_file) if not os.path.exists(model_path): logger.warning(f"Model file {csv_file} not found in {model_name}") continue model_df = pd.read_csv(model_path) model_mutation_col = get_mutation_column(model_df) - score_col = score_cols_dict[model_name] + score_col = resolve_source(score_cols_dict, model_name)[1] + if score_col not in model_df.columns: + raise KeyError( + f"{model_path} has no column {score_col!r}; it holds " + f"{list(model_df.columns)}" + ) model_df = model_df[[model_mutation_col, score_col]] model_df.columns = [mutation_col, f"{model_name}_score"] model_df[mutation_col] = model_df[mutation_col].apply(standardize_mutation) @@ -77,6 +83,45 @@ def combine_csv_data( logger.info(f"Saved combined data to {output_file}") +def resolve_source(score_cols_dict, model_name): + """ + Return the (folder, column) a model's scores are read from. + + An entry is either a bare column name, meaning the predictions live in a + folder named after the model, or a dict giving the folder and column + explicitly. The explicit form is what lets several entries read different + columns of the same prediction files, which is how the four masked-marginal + fill strategies are exposed: one run writes all four columns into one folder. + """ + try: + spec = score_cols_dict[model_name] + except KeyError: + raise KeyError(f"No score column configured for model {model_name!r}") + if isinstance(spec, str): + return model_name, spec + return spec["folder"], spec["column"] + + +def four_fill_entries(name, folder, column_stem): + """ + Register a masked language model's four fill strategies. + + A single scoring run writes one file per assay holding all four columns, so + the four entries share a folder and differ only in the column they read. + The fill named is what the model sees at a variant's OTHER mutated positions + while one position is masked: see fitness/baselines/masked_lm. + """ + return { + f"{name}_{strategy}": {"folder": folder, "column": f"{column_stem}_{strategy}"} + for strategy in ("wt_fill", "mask_fill", "mut_fill", "match_fill") + } + + +# The masked language models resolve to the wild-type fill, which is the +# convention the leaderboard publishes. Their prediction folders hold all four +# fill strategies as separate columns, so the entries below name the folder and +# the column explicitly, and the {name}_{strategy} entries further down read the +# other three from those same files. SCORE_COLS = { "evo1": "evo_1_131k_base_score", "evo1.5": "evo_1.5_8k_base_score", @@ -84,21 +129,40 @@ def combine_csv_data( "evo2_40b": "evo2_40b_score", "GenSLM": "logit_scores", "NT": "kmer_pseudo_LL", - "RNA-FM": "RNA_FM_scores", - "rinalmo": "logit_scores", + "RNA-FM": {"folder": "rna_fm_4fill", "column": "RNA_FM_scores_wt_fill"}, + "rinalmo": {"folder": "rinalmo_4fill", "column": "logit_scores_wt_fill"}, "RNAErnie": "Mutation_Scores", "orthrus": "orthrus_score", - "aido_rna": "aido_rna_score", - "rnagenesis": "rnagenesis_score", - # AIDO.RNA size series, for the scaling comparison. Every checkpoint writes the - # same aido_rna_score column, so they differ only by prediction folder. - "aido_rna_1m": "aido_rna_score", - "aido_rna_25m": "aido_rna_score", - "aido_rna_300m": "aido_rna_score", - "aido_rna_650m": "aido_rna_score", + "aido_rna": {"folder": "aido_rna_4fill", "column": "aido_rna_score_wt_fill"}, + "rnagenesis": {"folder": "rnagenesis_4fill", "column": "rnagenesis_score_wt_fill"}, + # AIDO.RNA size series. Every checkpoint writes the same column, so they + # differ by prediction folder. + "aido_rna_1m": {"folder": "aido_rna_1m_4fill", "column": "aido_rna_score_wt_fill"}, + "aido_rna_25m": {"folder": "aido_rna_25m_4fill", "column": "aido_rna_score_wt_fill"}, + "aido_rna_300m": {"folder": "aido_rna_300m_4fill", "column": "aido_rna_score_wt_fill"}, + "aido_rna_650m": {"folder": "aido_rna_650m_4fill", "column": "aido_rna_score_wt_fill"}, "EVmutation": "prediction_epistatic", } +# The four fill strategies, for every masked language model that runs them. Each +# model has one prediction folder holding all four columns. These are not in +# ALL_MODELS: pass them to --models explicitly, since one model appears four +# times and a default merge should not multiply the released leaderboard. +FOUR_FILL_MODELS = [] +for _name, _folder, _stem in [ + ("rna_fm", "rna_fm_4fill", "RNA_FM_scores"), + ("rinalmo", "rinalmo_4fill", "logit_scores"), + ("rnagenesis", "rnagenesis_4fill", "rnagenesis_score"), + ("aido_rna", "aido_rna_4fill", "aido_rna_score"), + ("aido_rna_1m", "aido_rna_1m_4fill", "aido_rna_score"), + ("aido_rna_25m", "aido_rna_25m_4fill", "aido_rna_score"), + ("aido_rna_300m", "aido_rna_300m_4fill", "aido_rna_score"), + ("aido_rna_650m", "aido_rna_650m_4fill", "aido_rna_score"), +]: + _entries = four_fill_entries(_name, _folder, _stem) + SCORE_COLS.update(_entries) + FOUR_FILL_MODELS.extend(_entries) + ALL_MODELS = [ "evo1", "evo1.5", @@ -139,6 +203,14 @@ def main(): required=True, help="Path to the folder where combined CSV files will be saved.", ) + parser.add_argument( + "--models", + nargs="+", + default=None, + help="Model entries to merge (default: every entry in ALL_MODELS). Use " + "this to merge a subset, such as one masked model's four fill " + "strategies: rna_fm_wt_fill rna_fm_mask_fill rna_fm_mut_fill rna_fm_match_fill", + ) parser.add_argument( "--assays_with_MSAs_only", action="store_true", @@ -146,7 +218,13 @@ def main(): ) args = parser.parse_args() - model_list = ["EVmutation"] if args.assays_with_MSAs_only else ALL_MODELS + if args.assays_with_MSAs_only: + model_list = ["EVmutation"] + else: + model_list = args.models if args.models else ALL_MODELS + unknown = [m for m in model_list if m not in SCORE_COLS] + if unknown: + parser.error(f"No score column configured for: {unknown}") combine_csv_data( args.processed_folder, diff --git a/fitness/performance_fitness.py b/fitness/performance_fitness.py index b8f05b3..5e256ae 100755 --- a/fitness/performance_fitness.py +++ b/fitness/performance_fitness.py @@ -27,13 +27,27 @@ def calculate_metrics( assay_scores: np.ndarray, model_scores: np.ndarray ) -> Dict[str, float]: - """Calculate Spearman correlation, AUC, and MCC.""" + """ + Calculate Spearman correlation, AUC, and MCC. + + All three are DIRECTED. They used to be reported as an absolute Spearman, an + AUC folded onto its better side with max(auc, 1 - auc), and an absolute MCC, + which together credited a model whose scores anti-correlate with fitness + exactly as much as one that correlates. That is the distinction the signed + metric adopted in v0.1.1 exists to make, and folding two of the three + metrics let a single row report a model as badly wrong by Spearman and + moderately good by AUC at the same time. + + Read them as: Spearman in [-1, 1] with 0 unrelated; AUC in [0, 1] with 0.5 + random, so below 0.5 means the model ranks variants the wrong way round; MCC + in [-1, 1] with 0 random. Higher remains better for all three. + """ spearman_corr = stats.spearmanr(assay_scores, model_scores).correlation binary_true = (assay_scores > np.median(assay_scores)).astype(int) binary_pred = (model_scores > np.median(model_scores)).astype(int) auc = roc_auc_score(y_true=binary_true, y_score=model_scores) mcc = matthews_corrcoef(y_true=binary_true, y_pred=binary_pred) - return {"Spearman": abs(spearman_corr), "AUC": max(auc, 1 - auc), "MCC": abs(mcc)} + return {"Spearman": spearman_corr, "AUC": auc, "MCC": mcc} def get_performance_dataset( diff --git a/leaderboard/fitness/README.md b/leaderboard/fitness/README.md index da6005d..483296b 100644 --- a/leaderboard/fitness/README.md +++ b/leaderboard/fitness/README.md @@ -13,53 +13,114 @@ Two changes from before: | Rank | Model | Ribozyme (n=26) | tRNA (n=3) | Aptamer (n=2) | Macro (3 ncRNA) | |---:|:--|--:|--:|--:|--:| -| 1 | Evo 2 (40B) | 0.1081 | 0.4310 | 0.0970 | 0.2120 | -| 2 | RNA-ERNIE | 0.1343 | 0.4161 | 0.0306 | 0.1937 | -| 3 | Evo 2 (7B) | 0.0651 | 0.3867 | 0.1192 | 0.1904 | -| 4 | AIDO.RNA (1.6B) | 0.0740 | 0.3885 | 0.0847 | 0.1824 | -| 5 | RNAGenesis | 0.0465 | 0.3735 | 0.0292 | 0.1498 | -| 6 | RiNALMo | -0.0092 | 0.3910 | 0.0356 | 0.1391 | -| 7 | Evo 1.5 | 0.0278 | 0.3850 | 0.0007 | 0.1378 | -| 8 | Nucleotide Transformer | 0.1329 | 0.3166 | -0.0886 | 0.1203 | -| 9 | RNA-FM | -0.0303 | 0.3342 | 0.0004 | 0.1014 | -| 10 | Orthrus | -0.0276 | -0.0068 | 0.1495 | 0.0384 | -| 11 | Evo 1 | -0.0216 | 0.0948 | 0.0058 | 0.0263 | -| 12 | GenSLM | -0.0045 | -0.0934 | -0.0036 | -0.0338 | +| 1 | AIDO.RNA (650M) | 0.0660 | 0.4894 | 0.0934 | 0.2163 | +| 2 | Evo 2 (40B) | 0.1081 | 0.4310 | 0.0970 | 0.2120 | +| 3 | AIDO.RNA (1.6B) | 0.0609 | 0.4884 | 0.0718 | 0.2070 | +| 4 | RNA-ERNIE | 0.1343 | 0.4161 | 0.0306 | 0.1937 | +| 5 | Evo 2 (7B) | 0.0651 | 0.3867 | 0.1192 | 0.1904 | +| 6 | RNAGenesis | 0.0750 | 0.4381 | 0.0343 | 0.1825 | +| 7 | RiNALMo | -0.0243 | 0.4856 | 0.0459 | 0.1690 | +| 8 | AIDO.RNA (300M) | -0.0256 | 0.4551 | 0.0372 | 0.1556 | +| 9 | AIDO.RNA (25M) | -0.0299 | 0.4569 | 0.0348 | 0.1540 | +| 10 | Evo 1.5 | 0.0278 | 0.3850 | 0.0007 | 0.1378 | +| 11 | RNA-FM | -0.0225 | 0.4147 | -0.0043 | 0.1293 | +| 12 | Nucleotide Transformer | 0.1329 | 0.3166 | -0.0886 | 0.1203 | +| 13 | AIDO.RNA (1M) | 0.0014 | 0.1777 | 0.0626 | 0.0806 | +| 14 | Orthrus | -0.0276 | -0.0068 | 0.1495 | 0.0384 | +| 15 | Evo 1 | -0.0216 | 0.0948 | 0.0058 | 0.0263 | +| 16 | GenSLM | -0.0045 | -0.0934 | -0.0036 | -0.0338 | + +**The top three are a tie, not a ranking.** AIDO.RNA (650M), Evo 2 (40B) and AIDO.RNA (1.6B) span +0.0093, inside the roughly 0.01 band that 31 assays cannot resolve (see Notes). What the top of the +table does show is that a 650M-parameter masked RNA model is level with a 40B autoregressive one. + +All five released AIDO.RNA checkpoints are listed as separate entries rather than in a side table, +since they are separate models scored the same way. Their scores rise steeply to 650M and then stop: +the 650M scores above the 1.6B under `wt-fill` and `match-fill`, though not under `mask-fill` or +`mut-fill`, so the shape of that curve depends on the scoring convention and should be read with that +in mind. Per category columns are the mean signed Spearman over the assays in that category; Macro is the unweighted average of the three category means. Underlying values: [`leaderboard_signed_3ncRNA.csv`](leaderboard_signed_3ncRNA.csv). -AIDO.RNA and RNAGenesis are new entries. RiNALMo and RNA-FM were rescored, see below. +The masked language models (AIDO.RNA at five sizes, RNAGenesis, RiNALMo, RNA-FM) are scored with `wt-fill`, +the convention the ESM and ProteinGym reference implementations use; see below. Their previous +numbers on this table used a different fill and were lower. **RNA-ERNIE is the one masked model not +on this convention**: it needs a paddlepaddle environment we do not have, so its score is carried +over from its original script, whose convention was never verified. Read its rank with that caveat. ## Scoring convention -All masked language models are now scored the same way: the masked-marginal log-likelihood ratio, -with each mutated position masked one at a time in the variant's own sequence, summed over the -variant's mutated positions. This matters because 99.4% of the ncRNA variants are multi-mutants, -which is exactly where scoring conventions differ. - -RiNALMo and RNA-FM were rescored to reach that convention, and their earlier scoring scripts were -replaced. The earlier RiNALMo scores looked the mutant base up in an RNA-form sequence, and its -alphabet is DNA, so mutations to U scored against the unknown-token logit. The earlier RNA-FM scores -match a single unmasked wild-type pass rather than the masked strategy its run script requested. - -## AIDO.RNA across model sizes - -All five released AIDO.RNA checkpoints on the same 31 assays. Values: -[`aido_rna_scaling.csv`](aido_rna_scaling.csv). - -| Checkpoint | Params | Ribozyme | tRNA | Aptamer | Macro | -|:--|--:|--:|--:|--:|--:| -| AIDO.RNA-1M-MARS | 1M | 0.0015 | 0.1763 | 0.0624 | 0.0801 | -| AIDO.RNA-25M-MARS | 25M | -0.0009 | 0.2718 | 0.0367 | 0.1025 | -| AIDO.RNA-300M-MARS | 299M | 0.0058 | 0.3684 | 0.0365 | 0.1369 | -| AIDO.RNA-650M | 646M | 0.0460 | 0.3928 | 0.0939 | 0.1775 | -| AIDO.RNA-1.6B | 1606M | 0.0740 | 0.3885 | 0.0847 | 0.1824 | - -Performance increases with size at every step. The macro gain from 650M to 1.6B is small (+0.0049), -but that is a property of the weighting rather than of the model: under an equal-assay average over -all 31 assays the same step is +0.0225, the second largest. tRNA and aptamer hold two thirds of the -macro weight and both stop improving at 650M, while ribozyme holds one third and keeps improving. +The masked language models are scored with the masked-marginal log-likelihood ratio: mask a mutated +position, take `log p(mutant) - log p(wild type)` there, and sum over the variant's mutated +positions. Four such conventions exist in the literature, all from Meier et al. 2021 (ESM-1v), +supplement Appendix A, and they differ in exactly one thing: **what fills the variant's OTHER mutated +positions while one position is masked.** With `M` the mutated positions, `x_-i` a mask at `i` and +`x_-M` masks at every position in `M`: + +| name | fill at the other mutated sites | score | +|:--|:--|:--| +| **`wt-fill`** | wild-type bases | `sum_i log p(mt_i \| x^wt_-i) - log p(wt_i \| x^wt_-i)` | +| `mask-fill` | masks | `sum_i log p(mt_i \| x^wt_-M) - log p(wt_i \| x^wt_-M)` | +| `mut-fill` | mutant bases | `sum_i log p(mt_i \| x^mt_-i) - log p(wt_i \| x^mt_-i)` | +| `match-fill` | the allele being scored | `sum_i log p(mt_i \| x^mt_-i) - log p(wt_i \| x^wt_-i)` | + +**The leaderboard uses `wt-fill`.** It is what the ESM authors' own +`examples/variant-prediction/predict.py` and ProteinGym's `proteingym/baselines/esm/compute_fitness.py` +implement under the option name `masked-marginals`, so it is the convention the surrounding +zero-shot literature is calibrated on. The choice is on that provenance, not on which scores best. + +All four are computed and published, see the sensitivity section below. `fitness/baselines/masked_lm` +computes them in a single pass per checkpoint: the four share most of their masked contexts, so all +four together cost about 19% more unique context examples than `mut-fill` alone (2,929,196 against +2,458,521 over the 31 assays; these are batched inputs, not model invocations). Every scoring script +accepts `--strategies`. + +Two properties worth knowing. **On single mutants all four are identical**, because a variant with +one mutation has no other mutated positions; that is used as a fixture test. And **99.4% of the ncRNA +variants are multi-mutants**, so the conventions diverge on essentially everything here. + +The name is overloaded and the overloading has caused real errors. The ESM code option called +`masked-marginals` is `wt-fill`, while the formula written in the ESM paper is `mask-fill`, and +`wt-marginals` is a different method again: one unmasked forward pass with no masking at all. + +RiNALMo and RNA-FM were also rescored to fix outright bugs, not just the convention. The earlier +RiNALMo scores looked the mutant base up in an RNA-form sequence while its alphabet is DNA, so +mutations to U scored against the unknown-token logit. The earlier RNA-FM scores match a single +unmasked wild-type pass (`wt-marginals`) rather than the masked strategy its run script requested. + +## Sensitivity to the fill strategy + +Every masked model is scored under all four conventions, so the effect of the choice is visible +rather than assumed. Macro over the 3 ncRNA categories; regenerate this table, the per-category +values and the bootstrap intervals with `fitness/analyze_fill_strategies.py`. + +| Checkpoint | `wt-fill` | `mask-fill` | `mut-fill` | `match-fill` | +|:--|--:|--:|--:|--:| +| AIDO.RNA-650M | 0.2163 | 0.1993 | 0.1767 | 0.2017 | +| AIDO.RNA-1.6B | 0.2070 | 0.2010 | 0.1824 | 0.1989 | +| RNAGenesis | 0.1825 | 0.1802 | 0.1497 | 0.1760 | +| RiNALMo | 0.1690 | 0.1643 | 0.1391 | 0.1611 | +| AIDO.RNA-300M-MARS | 0.1556 | 0.1514 | 0.1369 | 0.1539 | +| AIDO.RNA-25M-MARS | 0.1540 | 0.1380 | 0.1026 | 0.1393 | +| RNA-FM | 0.1293 | 0.1155 | 0.1014 | 0.1267 | +| AIDO.RNA-1M-MARS | 0.0806 | 0.0904 | 0.0800 | 0.0855 | + +What this does and does not support: + +- **`mut-fill` is the weakest of the four on every checkpoint.** A paired assay bootstrap, resampled + within each category, puts the `wt-fill` minus `mut-fill` interval above zero on 5 of the 8. +- **`wt-fill` is not separable from `mask-fill` or `match-fill`** at this sample size: those + intervals exclude zero on only 2 and 1 of the 8 checkpoints. `wt-fill` has the highest observed + macro on 7 of 8, but the convention is chosen on provenance and should not be defended on these + numbers. +- **The differences are concentrated in tRNA.** Mean spread across the four strategies is 0.094 in + tRNA against 0.021 in ribozyme and 0.007 in aptamer, because the conventions are identical on + single mutants and the tRNA assays are the deepest (4.07 mutations per variant against 2.91 and + 1.68). Since the macro weights 3 tRNA assays as heavily as 26 ribozyme assays, that effect is + amplified: under an unweighted mean over all 31 assays `wt-fill` leads on only 3 of 8 checkpoints. +- **Model ranking is far more stable than the absolute numbers.** Across the eight checkpoints the + only ordering change between strategies is that AIDO.RNA-650M and AIDO.RNA-1.6B trade places. ## Why we now exclude the mRNA assays @@ -82,7 +143,7 @@ less than about 0.01 should not be read as meaningful separations. ## Scoring scripts -Scoring scripts are on the [`v0.2`](https://github.com/MarksLab-DasLab/RNAGym/tree/v0.2) branch: +Scoring scripts, paths relative to the repository root: - Evo 2 40B: `fitness/baselines/Evo/score_evo2_single_dms.py` and `score_evo2.sh` - Orthrus: `fitness/baselines/Orthrus/score_orthrus_single_dms.py` and `score_orthrus.sh` - AIDO.RNA: `fitness/baselines/AIDO_RNA/score_aido_rna_single_dms.py` and `score_aido_rna.sh` @@ -90,5 +151,24 @@ Scoring scripts are on the [`v0.2`](https://github.com/MarksLab-DasLab/RNAGym/tr - RNA-FM: `fitness/baselines/RNA_FM/score_rna_fm_single_dms.py` and `score_rna_fm.sh` - RiNALMo: `fitness/baselines/RiNALMo/score_rinalmo_single_dms.py` and `score_rinalmo.sh` -All are registered in `fitness/merge_scoring_files.py` and `fitness/performance_fitness.py` on `v0.2`. -Reproduce the aggregate with `performance_fitness.py --type ncRNA`. +The four masked models share one scoring engine, `fitness/baselines/masked_lm`, which implements the +four fill strategies, the context deduplication and the batching; each model's script is a thin +adapter supplying its alphabet, tokenization and forward pass. `tests/test_masked_lm.py` covers the +strategies against a per-variant reference implementation and needs no checkpoint or GPU. + +All are registered in `fitness/merge_scoring_files.py` and `fitness/performance_fitness.py`. + +Reproduce the aggregate with `performance_fitness.py --type ncRNA`, whose Spearman is now signed. +It previously reported the absolute value, which credited a model whose scores anti-correlate with +fitness exactly as much as one that correlates, so it could not produce the numbers this page +publishes. Its per-category means now match the columns above, and the macro is their unweighted +mean. AUC and MCC are directed too: an AUC below 0.5 or an MCC below 0 means the model ranks variants +the wrong way round, where both were previously folded onto their better side. + +`fitness/analyze_fill_strategies.py` regenerates the sensitivity table, the category spreads and the +bootstrap intervals from the prediction files. + +Scores computed in bfloat16 depend on the GPU: the same code and checkpoint on an L40S and an H100 +differ by up to 0.2 in score and about 0.001 in per-assay Spearman. Every prediction file is written +with a manifest recording the strategies, alphabet, dtype, checkpoint and GPU. All numbers on this +page were produced on H100s. diff --git a/leaderboard/fitness/aido_rna_scaling.csv b/leaderboard/fitness/aido_rna_scaling.csv deleted file mode 100644 index c6a7d3b..0000000 --- a/leaderboard/fitness/aido_rna_scaling.csv +++ /dev/null @@ -1,6 +0,0 @@ -checkpoint,params_M,Ribozyme,tRNA,Aptamer,macro_3ncRNA -AIDO.RNA-1M-MARS,1,0.0015201861885637,0.1763240437720008,0.0623867974103732,0.0800770091236459 -AIDO.RNA-25M-MARS,25,-0.0009421067064766,0.271753681365211,0.0367220267208074,0.1025112004598473 -AIDO.RNA-300M-MARS,299,0.005809027714429,0.3683549928409946,0.0365494535558483,0.136904491370424 -AIDO.RNA-650M,646,0.0459675954454011,0.392759745777773,0.0938717001723726,0.1775330137985156 -AIDO.RNA-1.6B,1606,0.0740258014833371,0.3884879086583498,0.0846631863871998,0.1823922988429622 diff --git a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv index ad60f68..cec9cea 100644 --- a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv +++ b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv @@ -1,13 +1,17 @@ model,Ribozyme,tRNA,Aptamer,macro_3ncRNA +aido_rna_650m,0.0660332868987325,0.4893565654410639,0.0934325517551516,0.216274134698316 evo2_40b,0.1081020224381099,0.4309848333293184,0.09695482733614616,0.21201389436785814 +aido_rna,0.0608547145460644,0.4883787106096389,0.0718036455462622,0.2070123569006552 RNAErnie,0.13427241327809655,0.416115138737617,0.030641733377440288,0.1936764284643846 evo2,0.06511905079221095,0.3867345488695757,0.11920502912796216,0.1903528762632496 -aido_rna,0.0740258014833371,0.3884879086583498,0.0846631863871998,0.1823922988429622 -rnagenesis,0.0465473164565716,0.3735138133776106,0.0292223987359924,0.1497611761900582 -rinalmo,-0.0092385423512286,0.3909694048182631,0.0355731404591736,0.139101334308736 +rnagenesis,0.0750044159987944,0.4380555746849195,0.0343114577905016,0.1824571494914051 +rinalmo,-0.0243482728829925,0.4855667318003715,0.0458576152000779,0.1690253580391523 +aido_rna_300m,-0.0255811990086168,0.4550919102078445,0.0371789292494538,0.1555632134828938 +aido_rna_25m,-0.0298765175467791,0.456947919154572,0.0347915567069773,0.1539543194382567 evo1.5,0.02777228452283577,0.384997183239441,0.0007072149253433031,0.13782556089587336 +RNA-FM,-0.0224760788032156,0.4146672587305333,-0.0042742403435094,0.129305646527936 NT,0.13293178039954628,0.31660718643288793,-0.088591177571467,0.12031592975365574 -RNA-FM,-0.0303011446898677,0.3341866911496434,0.0004181460382951,0.1014345641660236 +aido_rna_1m,0.0013833725774681,0.1777039407912451,0.0625990558827026,0.0805621230838053 orthrus,-0.027566921396824637,-0.0068396133810284115,0.1494642014057122,0.038352555542619716 evo1,-0.021640697016188013,0.09484040594853649,0.0057677958977196755,0.026322501610022715 GenSLM,-0.004542860367691017,-0.09337311344747723,-0.0036277512402647302,-0.033847908351810986 diff --git a/tests/test_masked_lm.py b/tests/test_masked_lm.py new file mode 100644 index 0000000..86f67d7 --- /dev/null +++ b/tests/test_masked_lm.py @@ -0,0 +1,641 @@ +""" +Unit tests for the shared masked-marginal scoring engine. + +These run on CPU against a small deterministic stand-in model, so they check the +four fill strategies, the context bank and the accumulation without needing a +checkpoint or a GPU. The stand-in's logits at a position depend on the whole +context, which is what makes the four strategies disagree on multi-mutants, as +they must. +""" + +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "fitness" / "baselines")) + +from masked_lm import ( # noqa: E402 + STRATEGIES, + MaskedLMAdapter, + accumulate_scores, + build_tasks, +) +from masked_lm.strategies import MASK_CHAR, recover_wild_type # noqa: E402 + +VOCAB = {"": 0, "": 1, "": 2, "": 3, "A": 4, "C": 5, "G": 6, "U": 7, "N": 8} + + +class FakeAdapter(MaskedLMAdapter): + """A deterministic stand-in whose logits depend on the entire context.""" + + name = "fake" + bases = "ACGU" + score_column = "fake_score" + n_special_tokens = 2 # and + + def __init__(self): + self.device = "cpu" + self.prefix_ids = [VOCAB[""]] + self.suffix_ids = [VOCAB[""]] + self.pad_id = VOCAB[""] + self.mask_id = VOCAB[""] + self.unk_id = VOCAB["N"] + self.base_ids = {b: VOCAB[b] for b in self.bases} + generator = torch.Generator().manual_seed(0) + self.embedding = torch.randn(len(VOCAB), 8, generator=generator) + self.projection = torch.randn(8, len(VOCAB), generator=generator) + + @staticmethod + def add_arguments(parser): + pass + + def load(self, args): + pass + + def _logits(self, input_ids): + embedded = self.embedding[input_ids] + # Every position sees the whole context, so filling the other mutated + # positions differently changes the prediction, as in a real masked LM. + pooled = embedded.mean(dim=1, keepdim=True) + return (embedded + pooled) @ self.projection + + def logits_at(self, input_ids, attention_mask, rows, cols): + return self._logits(input_ids)[rows, cols] + + def naive_log_probs(self, context: str, position: int): + """Reference path: encode one context on its own and read one position.""" + ids = list(self.prefix_ids) + for char in context: + ids.append(self.mask_id if char == MASK_CHAR else self.base_ids[char]) + ids.extend(self.suffix_ids) + logits = self._logits(torch.tensor([ids], dtype=torch.long)) + return torch.log_softmax(logits[0, position + len(self.prefix_ids)].float(), -1).numpy() + + +def naive_score(adapter, wild_type, sequence, mutations, strategy): + """ + Score one variant straight from the formulas, one context at a time. + + Deliberately independent of the context bank and the batching engine. + """ + total = 0.0 + positions = [pos for pos, _, _ in mutations] + for pos, wt_base, mut_base in mutations: + if strategy == "wt_fill": + ctx = wild_type[:pos] + MASK_CHAR + wild_type[pos + 1 :] + lp = adapter.naive_log_probs(ctx, pos) + total += lp[adapter.base_ids[mut_base]] - lp[adapter.base_ids[wt_base]] + elif strategy == "mask_fill": + chars = list(wild_type) + for other in positions: + chars[other] = MASK_CHAR + lp = adapter.naive_log_probs("".join(chars), pos) + total += lp[adapter.base_ids[mut_base]] - lp[adapter.base_ids[wt_base]] + elif strategy == "mut_fill": + ctx = sequence[:pos] + MASK_CHAR + sequence[pos + 1 :] + lp = adapter.naive_log_probs(ctx, pos) + total += lp[adapter.base_ids[mut_base]] - lp[adapter.base_ids[wt_base]] + elif strategy == "match_fill": + mut_ctx = sequence[:pos] + MASK_CHAR + sequence[pos + 1 :] + wt_ctx = wild_type[:pos] + MASK_CHAR + wild_type[pos + 1 :] + total += adapter.naive_log_probs(mut_ctx, pos)[adapter.base_ids[mut_base]] + total -= adapter.naive_log_probs(wt_ctx, pos)[adapter.base_ids[wt_base]] + else: + raise ValueError(strategy) + return total + + +def make_assay(n_variants, max_mutations, length=24, seed=0): + """Build a synthetic assay: a wild type plus random substitution variants.""" + rng = np.random.default_rng(seed) + bases = list("ACGU") + wild_type = "".join(rng.choice(bases, size=length)) + mutants, sequences, parsed = [], [], [] + for _ in range(n_variants): + k = int(rng.integers(1, max_mutations + 1)) + positions = sorted(rng.choice(length, size=k, replace=False).tolist()) + tokens, seq = [], list(wild_type) + muts = [] + for pos in positions: + wt_base = wild_type[pos] + mut_base = str(rng.choice([b for b in bases if b != wt_base])) + seq[pos] = mut_base + tokens.append(f"{wt_base}{pos + 1}{mut_base}") + muts.append((pos, wt_base, mut_base)) + mutants.append(",".join(tokens)) + sequences.append("".join(seq)) + parsed.append(muts) + return wild_type, mutants, sequences, parsed + + +def score_with_engine(adapter, wild_type, mutants, sequences, strategies): + table = build_tasks(mutants, sequences, wild_type, adapter.bases, strategies, verbose=False) + scores = accumulate_scores( + adapter, table, n_rows=len(sequences), batch_size=7, max_batch_tokens=10**6, + progress=False, + ) + return table, scores + + +@pytest.fixture(scope="module") +def adapter(): + return FakeAdapter() + + +def test_matches_the_formulas(adapter): + """Every strategy reproduces a per-variant, one-context-at-a-time reference.""" + wild_type, mutants, sequences, parsed = make_assay(40, 4, seed=1) + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + _, scores = score_with_engine(adapter, wild_type, mutants, sequences, strategies) + for s, strategy in enumerate(strategies): + expected = [ + naive_score(adapter, wild_type, seq, muts, strategy) + for seq, muts in zip(sequences, parsed) + ] + assert np.allclose(scores[s], expected, atol=1e-6), strategy + + +def test_single_mutants_are_identical(adapter): + """ + With one mutation there are no other mutated positions, so the four + strategies are the same computation. This is the Pitt_2010_ribozyme fixture + in miniature: disagreement here means a bug. + """ + wild_type, mutants, sequences, _ = make_assay(30, 1, seed=2) + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + _, scores = score_with_engine(adapter, wild_type, mutants, sequences, strategies) + assert np.allclose(scores, scores[0], atol=1e-9) + + +def test_multi_mutants_diverge(adapter): + """The strategies must not silently collapse onto one another.""" + wild_type, mutants, sequences, _ = make_assay(30, 4, seed=3) + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + _, scores = score_with_engine(adapter, wild_type, mutants, sequences, strategies) + for i in range(1, 4): + assert not np.allclose(scores[i], scores[0]) + + +def test_contexts_are_shared_across_strategies(adapter): + """ + Asking for all four costs far less than asking for each separately, because + match-fill reuses mut-fill's and wt-fill's contexts and single mutants make + mask-fill's context equal to wt-fill's. + """ + wild_type, mutants, sequences, _ = make_assay(30, 4, seed=4) + sizes = {} + for strategy in ("wt_fill", "mask_fill", "mut_fill", "match_fill"): + table, _ = score_with_engine(adapter, wild_type, mutants, sequences, (strategy,)) + sizes[strategy] = len(table.contexts) + joint, _ = score_with_engine( + adapter, wild_type, mutants, sequences, ("wt_fill", "mask_fill", "mut_fill", "match_fill") + ) + assert len(joint.contexts) < sum(sizes.values()) + # match-fill introduces no context that mut-fill and wt-fill do not already need + assert sizes["match_fill"] <= sizes["mut_fill"] + sizes["wt_fill"] + combined, _ = score_with_engine( + adapter, wild_type, mutants, sequences, ("wt_fill", "mut_fill", "match_fill") + ) + pair, _ = score_with_engine(adapter, wild_type, mutants, sequences, ("wt_fill", "mut_fill")) + assert len(combined.contexts) == len(pair.contexts) + + +def test_strategy_subsets_agree_with_the_full_run(adapter): + """Requesting one strategy gives the same numbers as requesting all four.""" + wild_type, mutants, sequences, _ = make_assay(25, 3, seed=5) + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + _, full = score_with_engine(adapter, wild_type, mutants, sequences, strategies) + for i, strategy in enumerate(strategies): + _, alone = score_with_engine(adapter, wild_type, mutants, sequences, (strategy,)) + assert np.allclose(full[i], alone[0], atol=1e-9), strategy + + +def test_mask_fill_masks_the_whole_mutated_set(adapter): + """mask-fill uses one context per variant, carrying |M| masks.""" + wild_type = "ACGUACGUACGU" + mutants = ["A1C,G3U"] + sequences = ["CCUUACGUACGU"] + table = build_tasks(mutants, sequences, wild_type, "ACGU", ("mask_fill",), verbose=False) + assert len(table.contexts) == 1 + assert table.contexts[0].count(MASK_CHAR) == 2 + assert sorted(table.pos.tolist()) == [0, 0, 2, 2] + + +def test_wt_and_mut_fill_mask_one_position(adapter): + """wt-fill masks one position of the wild type; mut-fill masks one position + of the variant, so the variant's other mutation stays visible.""" + wild_type = "ACGUACGUACGU" + mutants = ["A1C,G3U"] + sequences = ["CCUUACGUACGU"] + + wt_table = build_tasks(mutants, sequences, wild_type, "ACGU", ("wt_fill",), verbose=False) + assert sorted(wt_table.contexts) == sorted(["#CGUACGUACGU", "AC#UACGUACGU"]) + + mut_table = build_tasks(mutants, sequences, wild_type, "ACGU", ("mut_fill",), verbose=False) + assert sorted(mut_table.contexts) == sorted(["#CUUACGUACGU", "CC#UACGUACGU"]) + + +def test_task_records_carry_the_right_base_sign_and_strategy(adapter): + """ + The term table is where a silent mix-up would be invisible, so check every + field of a two-mutation variant under all four strategies by hand. + """ + wild_type = "ACGUACGUACGU" + mutants = ["A1C,G3U"] + sequences = ["CCUUACGUACGU"] + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + table = build_tasks(mutants, sequences, wild_type, "ACGU", strategies, verbose=False) + + bases = "ACGU" + seen = set() + for k in range(table.n_terms()): + seen.add( + ( + table.contexts[table.ctx_id[k]], + int(table.pos[k]), + bases[table.base[k]], + strategies[table.strategy[k]], + int(table.sign[k]), + ) + ) + assert seen == { + # wt-fill: both alleles from the wild-type context, one mask + ("#CGUACGUACGU", 0, "C", "wt_fill", 1), + ("#CGUACGUACGU", 0, "A", "wt_fill", -1), + ("AC#UACGUACGU", 2, "U", "wt_fill", 1), + ("AC#UACGUACGU", 2, "G", "wt_fill", -1), + # mask-fill: both alleles from one wild-type context masked at both sites + ("#C#UACGUACGU", 0, "C", "mask_fill", 1), + ("#C#UACGUACGU", 0, "A", "mask_fill", -1), + ("#C#UACGUACGU", 2, "U", "mask_fill", 1), + ("#C#UACGUACGU", 2, "G", "mask_fill", -1), + # mut-fill: both alleles from the variant's own context + ("#CUUACGUACGU", 0, "C", "mut_fill", 1), + ("#CUUACGUACGU", 0, "A", "mut_fill", -1), + ("CC#UACGUACGU", 2, "U", "mut_fill", 1), + ("CC#UACGUACGU", 2, "G", "mut_fill", -1), + # match-fill: mutant allele from the variant context, wild-type allele + # from the wild-type context + ("#CUUACGUACGU", 0, "C", "match_fill", 1), + ("#CGUACGUACGU", 0, "A", "match_fill", -1), + ("CC#UACGUACGU", 2, "U", "match_fill", 1), + ("AC#UACGUACGU", 2, "G", "match_fill", -1), + } + + +def test_every_scored_position_is_masked(adapter): + """The table's own invariant check must reject a position that is not a mask.""" + from masked_lm.strategies import validate_table + + wild_type, mutants, sequences, _ = make_assay(10, 3, seed=8) + table = build_tasks(mutants, sequences, wild_type, "ACGU", ("mask_fill",), verbose=False) + validate_table(table) + table.pos = table.pos + 1 # shift every read off its mask + with pytest.raises(ValueError): + validate_table(table) + + +def test_multi_mask_context_is_gathered_at_every_masked_position(adapter): + """ + mask-fill reads one context at several positions. The engine must gather + each of them, not just the first. + """ + wild_type = "ACGUACGUACGU" + mutants = ["A1C,G3U"] + sequences = ["CCUUACGUACGU"] + table = build_tasks(mutants, sequences, wild_type, "ACGU", ("mask_fill",), verbose=False) + assert len(table.contexts) == 1 + scores = accumulate_scores( + adapter, table, n_rows=1, batch_size=4, max_batch_tokens=10**6, progress=False + ) + context = table.contexts[0] + expected = 0.0 + for pos, wt_base, mut_base in ((0, "A", "C"), (2, "G", "U")): + lp = adapter.naive_log_probs(context, pos) + expected += lp[adapter.base_ids[mut_base]] - lp[adapter.base_ids[wt_base]] + assert np.isclose(scores[0, 0], expected, atol=1e-6) + + +def test_prefix_offset_is_applied(adapter): + """ + A model with leading special tokens must read the token that carries the + nucleotide, not the special token. Scoring the same context with and without + a prefix must agree. + """ + wild_type, mutants, sequences, _ = make_assay(12, 2, seed=9) + table = build_tasks(mutants, sequences, wild_type, "ACGU", ("mut_fill",), verbose=False) + with_prefix = accumulate_scores( + adapter, table, len(sequences), batch_size=8, max_batch_tokens=10**6, progress=False + ) + + bare = FakeAdapter() + bare.prefix_ids = [] + bare.suffix_ids = [] + without_prefix = accumulate_scores( + bare, table, len(sequences), batch_size=8, max_batch_tokens=10**6, progress=False + ) + # Different inputs, so different numbers, but both must be finite and the + # offset must have moved with the prefix rather than staying at zero. + assert np.isfinite(with_prefix).all() and np.isfinite(without_prefix).all() + assert not np.allclose(with_prefix, without_prefix) + + shifted = FakeAdapter() + shifted.token_position = lambda pos: pos # forget the prefix offset + wrong = accumulate_scores( + shifted, table, len(sequences), batch_size=8, max_batch_tokens=10**6, progress=False + ) + assert not np.allclose(with_prefix, wrong) + + +def test_unknown_context_bases_are_encoded_as_unknown(adapter): + """ + A construct may contain an N. It is encoded as the unknown token in the + context, while mutation alleles themselves must be canonical bases. + """ + wild_type = "ACGUNCGUACGU" + mutants = ["A1C", "N5C"] + sequences = ["CCGUNCGUACGU", "ACGUCCGUACGU"] + table = build_tasks(mutants, sequences, wild_type, "ACGU", ("mut_fill",), verbose=False) + assert table.scorable.tolist() == [True, False] # N is not a scorable allele + ids = adapter.encode_context("ACGUNCGU", MASK_CHAR) + assert ids[5] == adapter.unk_id # position 4 plus the one leading token + + +def test_mask_placeholder_in_a_sequence_is_rejected(adapter): + """A sequence carrying the placeholder would silently become a mask.""" + wild_type = "ACGUACGUACGU" + table = build_tasks( + ["A1C"], ["CCGUACGUAC" + MASK_CHAR + "U"], wild_type, "ACGU", ("mut_fill",), verbose=False + ) + assert not table.scorable.any() + + +def test_malformed_variants_are_rejected(adapter): + """Duplicate positions, no-ops, and undeclared changes are all rejected.""" + wild_type = "ACGUACGUACGU" + cases = [ + ("A1C,A1G", "GCGUACGUACGU"), # same position named twice + ("A1A", "ACGUACGUACGU"), # no-op mutation + ("A1C", "CCGUACGUACGA"), # a change the mutation string omits + ] + for mutant, sequence in cases: + table = build_tasks([mutant], [sequence], wild_type, "ACGU", STRATEGIES, verbose=False) + assert not table.scorable.any(), mutant + + +def test_unscorable_variants_are_nan_under_every_strategy(adapter): + """A row that any strategy cannot score is NaN for all of them, so the + strategies are compared on exactly the same variants.""" + wild_type = "ACGUACGUACGU" + mutants = [ + "A1C", # fine + "A99C", # position outside the sequence + "A1-", # not a substitution + "G1C", # wild-type base disagrees with the wild type + None, # the wild-type row itself + ] + sequences = ["CCGUACGUACGU", "ACGUACGUACGU", "ACGUACGUACGU", "CCGUACGUACGU", wild_type] + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + table, scores = score_with_engine(adapter, wild_type, mutants, sequences, strategies) + assert table.scorable.tolist() == [True, False, False, False, False] + assert np.isfinite(scores[:, 0]).all() + assert np.isnan(scores[:, 1:]).all() + + +def test_wild_type_recovery(): + """The wild type is recoverable from the variants alone, and disagreement is + an error rather than a silently chosen majority.""" + wild_type, mutants, sequences, _ = make_assay(20, 3, seed=6) + assert recover_wild_type(mutants, sequences, "ACGU") == wild_type + broken = list(sequences) + broken[0] = "A" * len(wild_type) + mutants = list(mutants) + mutants[0] = "A1A" + with pytest.raises(ValueError): + recover_wild_type(mutants, broken, "ACGU") + + +def test_batching_does_not_change_scores(adapter): + """Scores must not depend on how contexts are packed into batches.""" + wild_type, mutants, sequences, _ = make_assay(25, 3, seed=7) + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + table = build_tasks(mutants, sequences, wild_type, "ACGU", strategies, verbose=False) + reference = accumulate_scores( + adapter, table, len(sequences), batch_size=1, max_batch_tokens=10**6, progress=False + ) + for batch_size in (3, 16, 1000): + other = accumulate_scores( + adapter, table, len(sequences), batch_size=batch_size, + max_batch_tokens=10**6, progress=False, + ) + assert np.allclose(reference, other, atol=1e-10), batch_size + + +def test_shared_distribution_feeds_several_rows_and_strategies(adapter): + """ + One gathered distribution is read by several variants, strategies and signs. + A flat-index or accumulation error would show up as one row stealing + another's contribution, so check the accumulation against a hand sum. + """ + wild_type = "ACGUACGUACGU" + # All three variants mutate position 0, so wt-fill reads one shared context + # there; the first two also share their mut-fill context at position 0. + mutants = ["A1C", "A1C,G3U", "A1G"] + sequences = ["CCGUACGUACGU", "CCUUACGUACGU", "GCGUACGUACGU"] + strategies = ("wt_fill", "mut_fill", "match_fill") + table, scores = score_with_engine(adapter, wild_type, mutants, sequences, strategies) + + shared = "#CGUACGUACGU" + assert shared in table.contexts + lp = adapter.naive_log_probs(shared, 0) + # every wt-fill term at position 0 comes from that one distribution + assert np.isclose(scores[0, 0], lp[adapter.base_ids["C"]] - lp[adapter.base_ids["A"]]) + assert np.isclose(scores[0, 2], lp[adapter.base_ids["G"]] - lp[adapter.base_ids["A"]]) + for s, strategy in enumerate(strategies): + expected = [ + naive_score(adapter, wild_type, seq, muts, strategy) + for seq, muts in zip(sequences, [ + [(0, "A", "C")], + [(0, "A", "C"), (2, "G", "U")], + [(0, "A", "G")], + ]) + ] + assert np.allclose(scores[s], expected, atol=1e-6), strategy + + +def test_gather_key_packing_at_the_boundaries(adapter): + """ + The engine packs (context, position) into one integer to find unique + distributions. Check the corners: the last position of one context and the + first position of the next must not collide. + """ + wild_type = "ACGUACGUACGU" + length = len(wild_type) + # one variant mutating the final position, one mutating the first + mutants = [f"U{length}A", "A1C"] + sequences = ["ACGUACGUACGA", "CCGUACGUACGU"] + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + _, scores = score_with_engine(adapter, wild_type, mutants, sequences, strategies) + expected_last = naive_score(adapter, wild_type, sequences[0], [(length - 1, "U", "A")], "mut_fill") + expected_first = naive_score(adapter, wild_type, sequences[1], [(0, "A", "C")], "mut_fill") + assert np.isclose(scores[2, 0], expected_last, atol=1e-6) + assert np.isclose(scores[2, 1], expected_first, atol=1e-6) + assert not np.isclose(expected_last, expected_first) + + +def test_batch_boundary_falls_around_a_heavily_used_context(adapter): + """ + Terms are sliced per batch with searchsorted, so a context carrying many + terms must be handled whichever side of a batch boundary it lands on. + """ + wild_type = "ACGUACGUACGU" + # 20 variants all mutating position 0, so one wt-fill context carries 40 terms + mutants, sequences = [], [] + for i in range(20): + other = 2 + (i % 8) + wt_other = wild_type[other] + mut_other = "A" if wt_other != "A" else "C" + mutants.append(f"A1C,{wt_other}{other + 1}{mut_other}") + seq = list(wild_type) + seq[0], seq[other] = "C", mut_other + sequences.append("".join(seq)) + strategies = ("wt_fill", "mask_fill", "mut_fill", "match_fill") + table = build_tasks(mutants, sequences, wild_type, "ACGU", strategies, verbose=False) + reference = accumulate_scores( + adapter, table, len(sequences), batch_size=len(table.contexts), + max_batch_tokens=10**6, progress=False, + ) + for batch_size in range(1, 6): + chunked = accumulate_scores( + adapter, table, len(sequences), batch_size=batch_size, + max_batch_tokens=10**6, progress=False, + ) + assert np.allclose(reference, chunked, atol=1e-10), batch_size + + +def test_windowing_keeps_masks_and_positions_consistent(adapter): + """ + Windowing shifts positions. Masks near both ends of a long construct must + survive it, and the table's invariant must still hold afterwards. + """ + from masked_lm.engine import window_contexts + from masked_lm.strategies import validate_table + + length = 200 + rng = np.random.default_rng(11) + wild_type = "".join(rng.choice(list("ACGU"), size=length)) + mutants, sequences = [], [] + for pos in (1, length - 2): + wt_base = wild_type[pos] + mut_base = "A" if wt_base != "A" else "C" + seq = list(wild_type) + seq[pos] = mut_base + mutants.append(f"{wt_base}{pos + 1}{mut_base}") + sequences.append("".join(seq)) + table = build_tasks(mutants, sequences, wild_type, "ACGU", ("mut_fill",), verbose=False) + window_contexts(table, 64) + validate_table(table) + assert all(len(c) <= 64 for c in table.contexts) + scores = accumulate_scores( + adapter, table, len(sequences), batch_size=2, max_batch_tokens=10**6, progress=False + ) + assert np.isfinite(scores).all() + + +def test_window_guard_uses_the_declared_special_token_count(): + """ + The guard that refuses multi-strategy scoring on over-long constructs runs + before the model is loaded, so it must use the adapter's declared special + token count. Reading the unloaded prefix_ids instead would give a budget two + positions too large and let a 1023 nt construct through. + """ + from masked_lm.runner import needs_windowing, window_budget + + unloaded = FakeAdapter.__new__(FakeAdapter) # not loaded: no prefix_ids yet + unloaded.max_tokens = 1024 + assert window_budget(unloaded, 1024) == 1022 + assert needs_windowing(["A" * 1023], window_budget(unloaded, 1024)) + assert not needs_windowing(["A" * 1022], window_budget(unloaded, 1024)) + + +@pytest.mark.parametrize( + "module_dir,module_name,expect", + [ + ("RNA_FM", "score_rna_fm_single_dms", {"cls": "RNAFMAdapter", "bases": "ACGU", + "column": "RNA_FM_scores", "special": 2, "batch": 512, "tokens": 65536, "dtype": None}), + ("RiNALMo", "score_rinalmo_single_dms", {"cls": "RiNALMoAdapter", "bases": "ACGT", + "column": "logit_scores", "special": 2, "batch": 512, "tokens": 65536, "dtype": None}), + ("AIDO_RNA", "score_aido_rna_single_dms", {"cls": "AIDORNAAdapter", "bases": "ACGT", + "column": "aido_rna_score", "special": 2, "batch": 512, "tokens": 49152, + "dtype": "bfloat16"}), + ("RNAGenesis", "score_rnagenesis_single_dms", {"cls": "RNAGenesisAdapter", "bases": "ACGU", + "column": "rnagenesis_score", "special": 0, "batch": 256, "tokens": 32768, + "dtype": "bfloat16"}), + ], +) +def test_adapters_keep_their_historical_defaults(module_dir, module_name, expect): + """ + Alphabet, output column, batching and dtype defaults are load bearing: they + are what the shipped predictions were produced with. A refactor that changes + one silently changes the benchmark's numbers. + """ + import argparse + import importlib.util + + path = ( + Path(__file__).resolve().parent.parent + / "fitness" / "baselines" / module_dir / f"{module_name}.py" + ) + spec = importlib.util.spec_from_file_location(module_name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + adapter = getattr(module, expect["cls"])() + + assert adapter.bases == expect["bases"] + assert adapter.score_column == expect["column"] + assert adapter.n_special_tokens == expect["special"] + assert adapter.default_batch_size == expect["batch"] + assert adapter.default_max_batch_tokens == expect["tokens"] + + parser = argparse.ArgumentParser() + adapter.add_arguments(parser) + defaults = {a.dest: a.default for a in parser._actions} + assert defaults.get("dtype") == expect["dtype"] + + +def test_performance_fitness_metrics_are_directed(): + """ + All three benchmark metrics report direction: a model that ranks variants + the wrong way round must score worse than random, not the same as a model + that ranks them correctly. Spearman was once an absolute value, the AUC was + folded with max(auc, 1 - auc) and the MCC was absolute, which let one row + call a model badly wrong and moderately good at the same time. Pinned here + rather than left to a reviewer to notice again. + """ + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + from fitness.performance_fitness import calculate_metrics + + # 40 distinct values, so the median splits both classes evenly and a fully + # reversed prediction reaches MCC -1 exactly. That is not general: with an + # odd count the median observation falls in the same class under both + # splits, and reversing arange(5) gives -2/3 rather than -1. + truth = np.arange(40, dtype=float) + perfect = calculate_metrics(truth, truth) + backwards = calculate_metrics(truth, -truth) + assert perfect["Spearman"] == pytest.approx(1.0) + assert backwards["Spearman"] == pytest.approx(-1.0) + assert perfect["AUC"] == pytest.approx(1.0) + assert backwards["AUC"] == pytest.approx(0.0) + assert perfect["MCC"] == pytest.approx(1.0) + assert backwards["MCC"] == pytest.approx(-1.0) + + rng = np.random.default_rng(0) + noisy = truth + rng.normal(0, 5, truth.size) + forward = calculate_metrics(truth, noisy) + reversed_ = calculate_metrics(truth, -noisy) + assert forward["Spearman"] == pytest.approx(-reversed_["Spearman"]) + assert forward["AUC"] > 0.5 > reversed_["AUC"] + assert forward["MCC"] > 0 > reversed_["MCC"]