From bd2df3db42ccd0be82ff00bdabf822b211244eae Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:12:30 -0400 Subject: [PATCH 01/11] Add a shared masked-marginal engine with four fill strategies Every masked-marginal score masks a mutated position and sums a log-odds over the variant's mutations; the published conventions differ in what fills the variant's other mutated positions while one is masked, which Meier et al. 2021 (ESM-1v) supplement Appendix A defines as wild-type, mask, mutant and allele-matched fills. Add fitness/baselines/masked_lm, which builds a deduplicated context bank for all four at once and accumulates them in a single pass, so the four together cost about 19% more unique context examples than the mutant fill alone (2,929,196 against 2,458,521 on the 31 ncRNA assays). Compute the four together because match-fill cannot be recovered from the final mut-fill and wt-fill scores: it needs their per-position halves. Reduce each model script to an adapter of about 100 lines supplying an alphabet, a tokenization and one forward pass; the four scorers lose 1,559 lines of duplicated masking, batching and IO. Take the wild type from RAW_CONSTRUCT_SEQ and independently recover it by reverting each variant's mutations, and stop if they disagree, since three of the four strategies condition on it. Reject a tokenizer that maps a base to the unknown token, which is the bug the released RiNALMo scores carried, and assert that a scored position is masked, that context and token positions differ by a constant shift, and that all four strategies cover the same variants. Write a manifest beside each prediction file recording the strategies, columns, alphabet, dtype, checkpoint, counts, a hash of the scoring source and the GPU, because bfloat16 scores differ between GPU models. Add tests/test_masked_lm.py, which checks each strategy against a per-variant reference implementation on a stand-in model, needs no checkpoint or GPU, and pins each adapter's alphabet, column, batching and dtype defaults. --- fitness/baselines/AIDO_RNA/score_aido_rna.sh | 18 +- .../AIDO_RNA/score_aido_rna_single_dms.py | 429 ++----------- .../baselines/RNAGenesis/score_rnagenesis.sh | 18 +- .../RNAGenesis/score_rnagenesis_single_dms.py | 435 ++----------- fitness/baselines/RNA_FM/score_rna_fm.sh | 18 +- .../RNA_FM/score_rna_fm_single_dms.py | 439 ++----------- fitness/baselines/RiNALMo/score_rinalmo.sh | 18 +- .../RiNALMo/score_rinalmo_single_dms.py | 467 ++------------ fitness/baselines/masked_lm/README.md | 73 +++ fitness/baselines/masked_lm/__init__.py | 35 + fitness/baselines/masked_lm/adapter.py | 145 +++++ fitness/baselines/masked_lm/engine.py | 194 ++++++ fitness/baselines/masked_lm/runner.py | 430 +++++++++++++ fitness/baselines/masked_lm/strategies.py | 469 ++++++++++++++ tests/test_masked_lm.py | 606 ++++++++++++++++++ 15 files changed, 2219 insertions(+), 1575 deletions(-) create mode 100644 fitness/baselines/masked_lm/README.md create mode 100644 fitness/baselines/masked_lm/__init__.py create mode 100644 fitness/baselines/masked_lm/adapter.py create mode 100644 fitness/baselines/masked_lm/engine.py create mode 100644 fitness/baselines/masked_lm/runner.py create mode 100644 fitness/baselines/masked_lm/strategies.py create mode 100644 tests/test_masked_lm.py 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/tests/test_masked_lm.py b/tests/test_masked_lm.py new file mode 100644 index 0000000..01be165 --- /dev/null +++ b/tests/test_masked_lm.py @@ -0,0 +1,606 @@ +""" +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"] From 5e9928b3e2886ba857e35da000bd4232a5334b2f Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:12:30 -0400 Subject: [PATCH 02/11] Let a merge entry name its prediction folder and column One scoring run now writes all four fill strategies into one folder as separate columns, but combine_csv_data used the model name as folder, lookup key and output column at once, so a folder could hold only one score. Resolve each entry through resolve_source: a bare column name keeps the previous meaning, and a {folder, column} entry lets several entries read different columns of the same prediction files. Register the eight masked checkpoints' four strategies with four_fill_entries, kept out of ALL_MODELS so a default merge still produces one row per model. Add --models to merge a subset, and fail on an unknown entry or on a prediction file that lacks its configured column. --- fitness/merge_scoring_files.py | 79 ++++++++++++++++++++++++++++++++-- 1 file changed, 76 insertions(+), 3 deletions(-) diff --git a/fitness/merge_scoring_files.py b/fitness/merge_scoring_files.py index a544dee..0fad656 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,40 @@ 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") + } + + SCORE_COLS = { "evo1": "evo_1_131k_base_score", "evo1.5": "evo_1.5_8k_base_score", @@ -99,6 +139,25 @@ def combine_csv_data( "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 +198,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 +213,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, From 75ddaadaece7542c1a0ead4dd493ec29d19a02c6 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:12:45 -0400 Subject: [PATCH 03/11] Score the masked models with the wild-type fill on the leaderboard Adopt wt-fill, the convention the ESM and ProteinGym reference implementations use under the name masked-marginals, in place of the mutant fill the previous table used. Choose it on that provenance rather than on the scores: wt-fill has the highest observed macro on 7 of 8 checkpoints, but a paired assay bootstrap separates it from mask-fill on only 2 and from match-fill on only 1, while mut-fill is the weakest of the four on every checkpoint. Rescore AIDO.RNA, RNAGenesis, RiNALMo and RNA-FM, which all rise: AIDO.RNA takes second place from RNA-ERNIE and Evo 2 (7B), and RNA-FM passes Nucleotide Transformer. Note that RNA-ERNIE is the one masked model not on this convention, since it needs a paddlepaddle environment, so its score is carried over unverified. Publish all four strategies in four_fill_sensitivity.csv, and record that their differences sit almost entirely in the tRNA assays, which the macro weighting amplifies: under an unweighted assay mean wt-fill leads on only 3 of 8. Replace the claim that AIDO.RNA improves at every size step, since under wt-fill the 650M checkpoint scores above the 1.6B. Add fitness/analyze_fill_strategies.py, which reproduces the sensitivity table, the category spreads and the bootstrap intervals from the prediction files. --- fitness/analyze_fill_strategies.py | 211 ++++++++++++++++++ leaderboard/fitness/README.md | 140 +++++++++--- leaderboard/fitness/aido_rna_scaling.csv | 10 +- leaderboard/fitness/four_fill_sensitivity.csv | 33 +++ .../fitness/leaderboard_signed_3ncRNA.csv | 8 +- 5 files changed, 362 insertions(+), 40 deletions(-) create mode 100644 fitness/analyze_fill_strategies.py create mode 100644 leaderboard/fitness/four_fill_sensitivity.csv 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/leaderboard/fitness/README.md b/leaderboard/fitness/README.md index da6005d..8fe0fbd 100644 --- a/leaderboard/fitness/README.md +++ b/leaderboard/fitness/README.md @@ -14,14 +14,14 @@ 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 | +| 2 | AIDO.RNA (1.6B) | 0.0609 | 0.4884 | 0.0718 | 0.2070 | +| 3 | RNA-ERNIE | 0.1343 | 0.4161 | 0.0306 | 0.1937 | +| 4 | Evo 2 (7B) | 0.0651 | 0.3867 | 0.1192 | 0.1904 | +| 5 | RNAGenesis | 0.0750 | 0.4381 | 0.0343 | 0.1825 | +| 6 | RiNALMo | -0.0243 | 0.4856 | 0.0459 | 0.1690 | | 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 | +| 8 | RNA-FM | -0.0225 | 0.4147 | -0.0043 | 0.1293 | +| 9 | Nucleotide Transformer | 0.1329 | 0.3166 | -0.0886 | 0.1203 | | 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 | @@ -29,37 +29,104 @@ Two changes from before: 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 four masked language models (AIDO.RNA, 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. +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. ## AIDO.RNA across model sizes -All five released AIDO.RNA checkpoints on the same 31 assays. Values: +All five released AIDO.RNA checkpoints on the same 31 assays, under `wt-fill`. 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. +| AIDO.RNA-1M-MARS | 1M | 0.0014 | 0.1777 | 0.0626 | 0.0806 | +| AIDO.RNA-25M-MARS | 25M | -0.0299 | 0.4569 | 0.0348 | 0.1540 | +| AIDO.RNA-300M-MARS | 299M | -0.0256 | 0.4551 | 0.0372 | 0.1556 | +| AIDO.RNA-650M | 646M | 0.0660 | 0.4894 | 0.0934 | 0.2163 | +| AIDO.RNA-1.6B | 1606M | 0.0609 | 0.4884 | 0.0718 | 0.2070 | + +Performance rises steeply to 650M and then does not improve: **the 650M checkpoint scores above the +1.6B** (0.2163 against 0.2070). Whether the curve is monotone depends on the fill strategy, which is +a reason to read it cautiously rather than a result about the models: under `mask-fill` and +`mut-fill` the series is monotone in size, while under `wt-fill` and `match-fill` the 650M leads. The +650M against 1.6B gap is also not resolved by 31 assays; a paired assay bootstrap puts it at +0.0093 +with a 95% interval of [-0.0057, +0.0245]. + +## 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; full per-category values in +[`four_fill_sensitivity.csv`](four_fill_sensitivity.csv). + +| 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 +149,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 +157,16 @@ 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`, and the sensitivity table, +category spreads and bootstrap intervals above with `fitness/analyze_fill_strategies.py`. + +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 index c6a7d3b..30ea467 100644 --- a/leaderboard/fitness/aido_rna_scaling.csv +++ b/leaderboard/fitness/aido_rna_scaling.csv @@ -1,6 +1,6 @@ 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 +AIDO.RNA-1M-MARS,1,0.0013833725774681,0.1777039407912451,0.0625990558827026,0.0805621230838053 +AIDO.RNA-25M-MARS,25,-0.0298765175467791,0.456947919154572,0.0347915567069773,0.1539543194382567 +AIDO.RNA-300M-MARS,299,-0.0255811990086168,0.4550919102078445,0.0371789292494538,0.1555632134828938 +AIDO.RNA-650M,646,0.0660332868987325,0.4893565654410639,0.0934325517551516,0.216274134698316 +AIDO.RNA-1.6B,1606,0.0608547145460644,0.4883787106096389,0.0718036455462622,0.2070123569006552 diff --git a/leaderboard/fitness/four_fill_sensitivity.csv b/leaderboard/fitness/four_fill_sensitivity.csv new file mode 100644 index 0000000..14eedaf --- /dev/null +++ b/leaderboard/fitness/four_fill_sensitivity.csv @@ -0,0 +1,33 @@ +model,strategy,Ribozyme,tRNA,Aptamer,macro_3ncRNA +AIDO.RNA-1.6B,mask-fill,0.0711558718183486,0.4528447167435985,0.0790099175279454,0.2010035020299642 +AIDO.RNA-1.6B,match-fill,0.0594302533307336,0.4597095238591757,0.0775229431656438,0.198887573451851 +AIDO.RNA-1.6B,mut-fill,0.073937583955693,0.3885241187083097,0.0847234840806227,0.1823950622482084 +AIDO.RNA-1.6B,wt-fill,0.0608547145460644,0.4883787106096389,0.0718036455462622,0.2070123569006552 +AIDO.RNA-1M-MARS,mask-fill,-0.0003858066042423,0.2099289676844771,0.061706392753086,0.0904165179444402 +AIDO.RNA-1M-MARS,match-fill,-0.0036154309099485,0.2001046465356077,0.0599300356988054,0.0854730837748215 +AIDO.RNA-1M-MARS,mut-fill,0.001744967493573,0.1760305856902303,0.062215601235242,0.0799970514730151 +AIDO.RNA-1M-MARS,wt-fill,0.0013833725774681,0.1777039407912451,0.0625990558827026,0.0805621230838053 +AIDO.RNA-25M-MARS,mask-fill,-0.0080679765252972,0.3846516118384668,0.0375659162215725,0.1380498505115807 +AIDO.RNA-25M-MARS,match-fill,-0.0223139990013149,0.4054395286704635,0.0347959986502595,0.1393071761064694 +AIDO.RNA-25M-MARS,mut-fill,-0.0007045190339156,0.2715814746903859,0.0370100959874292,0.1026290172146331 +AIDO.RNA-25M-MARS,wt-fill,-0.0298765175467791,0.456947919154572,0.0347915567069773,0.1539543194382567 +AIDO.RNA-300M-MARS,mask-fill,-0.002318707778476,0.4156447069635412,0.040882006181148,0.1514026684554044 +AIDO.RNA-300M-MARS,match-fill,-0.0202025458235484,0.4459552615131028,0.0360900276971144,0.1539475811288896 +AIDO.RNA-300M-MARS,mut-fill,0.0056655763225726,0.3683046102759297,0.0367067135240057,0.136892300040836 +AIDO.RNA-300M-MARS,wt-fill,-0.0255811990086168,0.4550919102078445,0.0371789292494538,0.1555632134828938 +AIDO.RNA-650M,mask-fill,0.0561792246625304,0.447101236218364,0.0947219340048093,0.1993341316285679 +AIDO.RNA-650M,match-fill,0.0495166659489845,0.4639243972041769,0.0917207407628353,0.2017206013053322 +AIDO.RNA-650M,mut-fill,0.046158489045562,0.3927081253182438,0.0913799489696062,0.1767488544444707 +AIDO.RNA-650M,wt-fill,0.0660332868987325,0.4893565654410639,0.0934325517551516,0.216274134698316 +RNA-FM,mask-fill,-0.0409644821312328,0.3866832315666649,0.0009191892371025,0.1155459795575115 +RNA-FM,match-fill,-0.023997554480859,0.403098385867694,0.0010741984500069,0.1267250099456139 +RNA-FM,mut-fill,-0.0303011460608878,0.3341866961122435,0.0004181460382951,0.1014345653632169 +RNA-FM,wt-fill,-0.0224760788032156,0.4146672587305333,-0.0042742403435094,0.129305646527936 +RNAGenesis,mask-fill,0.0646649413419138,0.4430794477960186,0.03294921045597,0.1802311998646341 +RNAGenesis,match-fill,0.0602226295416455,0.444001298043382,0.0237764923043939,0.1760001399631404 +RNAGenesis,mut-fill,0.0465260030681693,0.3735092434830764,0.0289580251896279,0.1496644239136245 +RNAGenesis,wt-fill,0.0750044159987944,0.4380555746849195,0.0343114577905016,0.1824571494914051 +RiNALMo,mask-fill,-0.0143484959918101,0.459569427178537,0.0476616875034881,0.1642942062300717 +RiNALMo,match-fill,-0.0294320277613483,0.4741679689727105,0.0384591474849403,0.1610650295654341 +RiNALMo,mut-fill,-0.0092389501238413,0.3909676352806882,0.0356096135810886,0.1391127662459785 +RiNALMo,wt-fill,-0.0243482728829925,0.4855667318003715,0.0458576152000779,0.1690253580391523 diff --git a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv index ad60f68..f556df4 100644 --- a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv +++ b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv @@ -1,13 +1,13 @@ model,Ribozyme,tRNA,Aptamer,macro_3ncRNA 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 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 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 From ca3508af5094c90878cc4511535e1a5317daa775 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:07:17 -0400 Subject: [PATCH 04/11] Correct what performance_fitness.py reproduces The page said to reproduce the aggregate with performance_fitness.py, but that script returns the absolute Spearman, a direction-folded AUC and an absolute MCC, so it cannot produce the signed values the leaderboard publishes. Say so, and point at analyze_fill_strategies.py, which computes the signed per-assay Spearman and the category macro from the prediction files. Leave performance_fitness.py itself alone: adding a signed option changes a shared aggregation used by the other tables and belongs in its own change. --- leaderboard/fitness/README.md | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/leaderboard/fitness/README.md b/leaderboard/fitness/README.md index 8fe0fbd..198333a 100644 --- a/leaderboard/fitness/README.md +++ b/leaderboard/fitness/README.md @@ -163,8 +163,15 @@ adapter supplying its alphabet, tokenization and forward pass. `tests/test_maske 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`, and the sensitivity table, -category spreads and bootstrap intervals above with `fitness/analyze_fill_strategies.py`. + +**`performance_fitness.py` does not reproduce the signed values on this page.** It reports the +absolute Spearman, along with a direction-folded AUC and an absolute MCC, so a model whose scores +anti-correlate with fitness is credited the same as one that correlates, which is exactly the +distinction the v0.1.1 metric was introduced to make. The four rescored rows here, the scaling series +and the sensitivity table are reproduced by `fitness/analyze_fill_strategies.py`, which computes the +signed per-assay Spearman and the category macro directly from the prediction files; the eight +unaffected rows are carried over unchanged from the previous release. Giving +`performance_fitness.py` a signed option is worth doing separately. 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 From b28e786840d9745fcfdfdef1e5c1e8f510a09365 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:05:09 -0400 Subject: [PATCH 05/11] List every AIDO.RNA checkpoint on the leaderboard itself The five AIDO.RNA sizes were split between the leaderboard, which carried only the 1.6B, and a side table, which is an odd division for five models scored on the same assays with the same convention. Fold all five into leaderboard_signed_3ncRNA.csv and drop aido_rna_scaling.csv. Say that the top three entries are a tie: 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, which the Notes section already states. Keep the observation that the size series stops improving after 650M, and that whether it is monotone depends on the fill strategy, as prose rather than as a second table. --- leaderboard/fitness/README.md | 60 +++++++++---------- leaderboard/fitness/aido_rna_scaling.csv | 6 -- .../fitness/leaderboard_signed_3ncRNA.csv | 4 ++ 3 files changed, 31 insertions(+), 39 deletions(-) delete mode 100644 leaderboard/fitness/aido_rna_scaling.csv diff --git a/leaderboard/fitness/README.md b/leaderboard/fitness/README.md index 198333a..ab427a6 100644 --- a/leaderboard/fitness/README.md +++ b/leaderboard/fitness/README.md @@ -13,23 +13,37 @@ 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 | AIDO.RNA (1.6B) | 0.0609 | 0.4884 | 0.0718 | 0.2070 | -| 3 | RNA-ERNIE | 0.1343 | 0.4161 | 0.0306 | 0.1937 | -| 4 | Evo 2 (7B) | 0.0651 | 0.3867 | 0.1192 | 0.1904 | -| 5 | RNAGenesis | 0.0750 | 0.4381 | 0.0343 | 0.1825 | -| 6 | RiNALMo | -0.0243 | 0.4856 | 0.0459 | 0.1690 | -| 7 | Evo 1.5 | 0.0278 | 0.3850 | 0.0007 | 0.1378 | -| 8 | RNA-FM | -0.0225 | 0.4147 | -0.0043 | 0.1293 | -| 9 | Nucleotide Transformer | 0.1329 | 0.3166 | -0.0886 | 0.1203 | -| 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). -The four masked language models (AIDO.RNA, RNAGenesis, RiNALMo, RNA-FM) are scored with `wt-fill`, +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 @@ -75,26 +89,6 @@ RiNALMo scores looked the mutant base up in an RNA-form sequence while its alpha 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. -## AIDO.RNA across model sizes - -All five released AIDO.RNA checkpoints on the same 31 assays, under `wt-fill`. Values: -[`aido_rna_scaling.csv`](aido_rna_scaling.csv). - -| Checkpoint | Params | Ribozyme | tRNA | Aptamer | Macro | -|:--|--:|--:|--:|--:|--:| -| AIDO.RNA-1M-MARS | 1M | 0.0014 | 0.1777 | 0.0626 | 0.0806 | -| AIDO.RNA-25M-MARS | 25M | -0.0299 | 0.4569 | 0.0348 | 0.1540 | -| AIDO.RNA-300M-MARS | 299M | -0.0256 | 0.4551 | 0.0372 | 0.1556 | -| AIDO.RNA-650M | 646M | 0.0660 | 0.4894 | 0.0934 | 0.2163 | -| AIDO.RNA-1.6B | 1606M | 0.0609 | 0.4884 | 0.0718 | 0.2070 | - -Performance rises steeply to 650M and then does not improve: **the 650M checkpoint scores above the -1.6B** (0.2163 against 0.2070). Whether the curve is monotone depends on the fill strategy, which is -a reason to read it cautiously rather than a result about the models: under `mask-fill` and -`mut-fill` the series is monotone in size, while under `wt-fill` and `match-fill` the 650M leads. The -650M against 1.6B gap is also not resolved by 31 assays; a paired assay bootstrap puts it at +0.0093 -with a 95% interval of [-0.0057, +0.0245]. - ## Sensitivity to the fill strategy Every masked model is scored under all four conventions, so the effect of the choice is visible diff --git a/leaderboard/fitness/aido_rna_scaling.csv b/leaderboard/fitness/aido_rna_scaling.csv deleted file mode 100644 index 30ea467..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.0013833725774681,0.1777039407912451,0.0625990558827026,0.0805621230838053 -AIDO.RNA-25M-MARS,25,-0.0298765175467791,0.456947919154572,0.0347915567069773,0.1539543194382567 -AIDO.RNA-300M-MARS,299,-0.0255811990086168,0.4550919102078445,0.0371789292494538,0.1555632134828938 -AIDO.RNA-650M,646,0.0660332868987325,0.4893565654410639,0.0934325517551516,0.216274134698316 -AIDO.RNA-1.6B,1606,0.0608547145460644,0.4883787106096389,0.0718036455462622,0.2070123569006552 diff --git a/leaderboard/fitness/leaderboard_signed_3ncRNA.csv b/leaderboard/fitness/leaderboard_signed_3ncRNA.csv index f556df4..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 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 +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 From b2b4be3e9709e06192f7b51b8cc519de746b224f Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:11:11 -0400 Subject: [PATCH 06/11] Say which rows analyze_fill_strategies.py reproduces Now that all five AIDO.RNA sizes are leaderboard entries, eight of the sixteen rows are masked checkpoints, so calling them the four rescored rows was ambiguous. --- leaderboard/fitness/README.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/leaderboard/fitness/README.md b/leaderboard/fitness/README.md index ab427a6..7dd25c5 100644 --- a/leaderboard/fitness/README.md +++ b/leaderboard/fitness/README.md @@ -161,10 +161,11 @@ All are registered in `fitness/merge_scoring_files.py` and `fitness/performance_ **`performance_fitness.py` does not reproduce the signed values on this page.** It reports the absolute Spearman, along with a direction-folded AUC and an absolute MCC, so a model whose scores anti-correlate with fitness is credited the same as one that correlates, which is exactly the -distinction the v0.1.1 metric was introduced to make. The four rescored rows here, the scaling series -and the sensitivity table are reproduced by `fitness/analyze_fill_strategies.py`, which computes the -signed per-assay Spearman and the category macro directly from the prediction files; the eight -unaffected rows are carried over unchanged from the previous release. Giving +distinction the v0.1.1 metric was introduced to make. All eight masked checkpoint rows here, meaning the +four model families and the five AIDO.RNA sizes, together with the sensitivity table, are reproduced +by `fitness/analyze_fill_strategies.py`, which computes the signed per-assay Spearman and the +category macro directly from the prediction files; the eight remaining rows are carried over +unchanged from the previous release and it does not recompute them. Giving `performance_fitness.py` a signed option is worth doing separately. Scores computed in bfloat16 depend on the GPU: the same code and checkpoint on an L40S and an H100 From 7a6250dd193b7cb3bb2b8adfbfd7c24c3841e48d Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:15:40 -0400 Subject: [PATCH 07/11] Report the signed Spearman in performance_fitness.py The script returned the absolute Spearman, so it credited a model whose scores anti-correlate with fitness exactly as much as one that correlates, which is the distinction the signed metric adopted in v0.1.1 exists to make, and it therefore could not produce the numbers the leaderboard publishes. Return the signed value. Its per-category means now reproduce this page's category columns for all eight masked checkpoints to floating-point precision, with the macro their unweighted mean. Leave AUC and MCC folded onto their better direction for now, and say so in the docstring: it is the same conflation, but unfolding them changes published numbers this change is not otherwise touching. Drop four_fill_sensitivity.csv, whose contents are the table on the leaderboard page and are regenerated by analyze_fill_strategies.py. Note that tests/test_fitness.py compares against result files published before this fix, so it will fail on any assay where a model anti-correlates until those artifacts are regenerated. --- fitness/performance_fitness.py | 17 ++++++++-- leaderboard/fitness/README.md | 22 ++++++------- leaderboard/fitness/four_fill_sensitivity.csv | 33 ------------------- 3 files changed, 26 insertions(+), 46 deletions(-) delete mode 100644 leaderboard/fitness/four_fill_sensitivity.csv diff --git a/fitness/performance_fitness.py b/fitness/performance_fitness.py index b8f05b3..e5df536 100755 --- a/fitness/performance_fitness.py +++ b/fitness/performance_fitness.py @@ -27,13 +27,26 @@ 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. + + The Spearman correlation is SIGNED. It used to be reported as an absolute + value, which credited a model whose scores anti-correlate with fitness + exactly as much as one that correlates, and that is the distinction the + signed metric adopted in v0.1.1 exists to make. The leaderboard has + published signed values since then, so the absolute form no longer + corresponds to anything the benchmark reports. + + AUC and MCC are still folded onto their better direction below. That is the + same conflation and should be revisited, but it is left alone here because + it would change published numbers this change is not otherwise touching. + """ 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": max(auc, 1 - auc), "MCC": abs(mcc)} def get_performance_dataset( diff --git a/leaderboard/fitness/README.md b/leaderboard/fitness/README.md index 7dd25c5..efdc2d0 100644 --- a/leaderboard/fitness/README.md +++ b/leaderboard/fitness/README.md @@ -92,8 +92,8 @@ unmasked wild-type pass (`wt-marginals`) rather than the masked strategy its run ## 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; full per-category values in -[`four_fill_sensitivity.csv`](four_fill_sensitivity.csv). +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` | |:--|--:|--:|--:|--:| @@ -158,15 +158,15 @@ strategies against a per-variant reference implementation and needs no checkpoin All are registered in `fitness/merge_scoring_files.py` and `fitness/performance_fitness.py`. -**`performance_fitness.py` does not reproduce the signed values on this page.** It reports the -absolute Spearman, along with a direction-folded AUC and an absolute MCC, so a model whose scores -anti-correlate with fitness is credited the same as one that correlates, which is exactly the -distinction the v0.1.1 metric was introduced to make. All eight masked checkpoint rows here, meaning the -four model families and the five AIDO.RNA sizes, together with the sensitivity table, are reproduced -by `fitness/analyze_fill_strategies.py`, which computes the signed per-assay Spearman and the -category macro directly from the prediction files; the eight remaining rows are carried over -unchanged from the previous release and it does not recompute them. Giving -`performance_fitness.py` a signed option is worth doing separately. +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. Note that AUC and MCC in that script are still folded onto their better direction, which is the +same conflation and is worth revisiting separately. + +`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 diff --git a/leaderboard/fitness/four_fill_sensitivity.csv b/leaderboard/fitness/four_fill_sensitivity.csv deleted file mode 100644 index 14eedaf..0000000 --- a/leaderboard/fitness/four_fill_sensitivity.csv +++ /dev/null @@ -1,33 +0,0 @@ -model,strategy,Ribozyme,tRNA,Aptamer,macro_3ncRNA -AIDO.RNA-1.6B,mask-fill,0.0711558718183486,0.4528447167435985,0.0790099175279454,0.2010035020299642 -AIDO.RNA-1.6B,match-fill,0.0594302533307336,0.4597095238591757,0.0775229431656438,0.198887573451851 -AIDO.RNA-1.6B,mut-fill,0.073937583955693,0.3885241187083097,0.0847234840806227,0.1823950622482084 -AIDO.RNA-1.6B,wt-fill,0.0608547145460644,0.4883787106096389,0.0718036455462622,0.2070123569006552 -AIDO.RNA-1M-MARS,mask-fill,-0.0003858066042423,0.2099289676844771,0.061706392753086,0.0904165179444402 -AIDO.RNA-1M-MARS,match-fill,-0.0036154309099485,0.2001046465356077,0.0599300356988054,0.0854730837748215 -AIDO.RNA-1M-MARS,mut-fill,0.001744967493573,0.1760305856902303,0.062215601235242,0.0799970514730151 -AIDO.RNA-1M-MARS,wt-fill,0.0013833725774681,0.1777039407912451,0.0625990558827026,0.0805621230838053 -AIDO.RNA-25M-MARS,mask-fill,-0.0080679765252972,0.3846516118384668,0.0375659162215725,0.1380498505115807 -AIDO.RNA-25M-MARS,match-fill,-0.0223139990013149,0.4054395286704635,0.0347959986502595,0.1393071761064694 -AIDO.RNA-25M-MARS,mut-fill,-0.0007045190339156,0.2715814746903859,0.0370100959874292,0.1026290172146331 -AIDO.RNA-25M-MARS,wt-fill,-0.0298765175467791,0.456947919154572,0.0347915567069773,0.1539543194382567 -AIDO.RNA-300M-MARS,mask-fill,-0.002318707778476,0.4156447069635412,0.040882006181148,0.1514026684554044 -AIDO.RNA-300M-MARS,match-fill,-0.0202025458235484,0.4459552615131028,0.0360900276971144,0.1539475811288896 -AIDO.RNA-300M-MARS,mut-fill,0.0056655763225726,0.3683046102759297,0.0367067135240057,0.136892300040836 -AIDO.RNA-300M-MARS,wt-fill,-0.0255811990086168,0.4550919102078445,0.0371789292494538,0.1555632134828938 -AIDO.RNA-650M,mask-fill,0.0561792246625304,0.447101236218364,0.0947219340048093,0.1993341316285679 -AIDO.RNA-650M,match-fill,0.0495166659489845,0.4639243972041769,0.0917207407628353,0.2017206013053322 -AIDO.RNA-650M,mut-fill,0.046158489045562,0.3927081253182438,0.0913799489696062,0.1767488544444707 -AIDO.RNA-650M,wt-fill,0.0660332868987325,0.4893565654410639,0.0934325517551516,0.216274134698316 -RNA-FM,mask-fill,-0.0409644821312328,0.3866832315666649,0.0009191892371025,0.1155459795575115 -RNA-FM,match-fill,-0.023997554480859,0.403098385867694,0.0010741984500069,0.1267250099456139 -RNA-FM,mut-fill,-0.0303011460608878,0.3341866961122435,0.0004181460382951,0.1014345653632169 -RNA-FM,wt-fill,-0.0224760788032156,0.4146672587305333,-0.0042742403435094,0.129305646527936 -RNAGenesis,mask-fill,0.0646649413419138,0.4430794477960186,0.03294921045597,0.1802311998646341 -RNAGenesis,match-fill,0.0602226295416455,0.444001298043382,0.0237764923043939,0.1760001399631404 -RNAGenesis,mut-fill,0.0465260030681693,0.3735092434830764,0.0289580251896279,0.1496644239136245 -RNAGenesis,wt-fill,0.0750044159987944,0.4380555746849195,0.0343114577905016,0.1824571494914051 -RiNALMo,mask-fill,-0.0143484959918101,0.459569427178537,0.0476616875034881,0.1642942062300717 -RiNALMo,match-fill,-0.0294320277613483,0.4741679689727105,0.0384591474849403,0.1610650295654341 -RiNALMo,mut-fill,-0.0092389501238413,0.3909676352806882,0.0356096135810886,0.1391127662459785 -RiNALMo,wt-fill,-0.0243482728829925,0.4855667318003715,0.0458576152000779,0.1690253580391523 From fe86f36dc49c98254b8e148358b80889aab38d92 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:27:42 -0400 Subject: [PATCH 08/11] Pin the signed Spearman and correct the metric description fitness/README.md still described the Spearman as absolute, which was the last place in the repository saying the active metric ignores direction. Add a test that a reversed prediction returns a negative Spearman equal in magnitude to the forward one, so that reverting to an absolute value fails rather than passing silently. --- fitness/README.md | 2 +- tests/test_masked_lm.py | 22 ++++++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/fitness/README.md b/fitness/README.md index 4acdd8a..4609c5d 100644 --- a/fitness/README.md +++ b/fitness/README.md @@ -34,7 +34,7 @@ The following directories are gitignored. Download them from ## Metrics -- **Spearman** — absolute Spearman correlation between predicted and experimental scores +- **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 using median-threshold binarisation (max of AUC and 1−AUC) - **MCC** — absolute Matthews correlation coefficient (median-threshold binarisation) diff --git a/tests/test_masked_lm.py b/tests/test_masked_lm.py index 01be165..2722bb5 100644 --- a/tests/test_masked_lm.py +++ b/tests/test_masked_lm.py @@ -604,3 +604,25 @@ def test_adapters_keep_their_historical_defaults(module_dir, module_name, expect adapter.add_arguments(parser) defaults = {a.dest: a.default for a in parser._actions} assert defaults.get("dtype") == expect["dtype"] + + +def test_performance_fitness_reports_a_signed_spearman(): + """ + The benchmark metric is signed: a model that ranks variants the wrong way + round must score negative, not positive. This was reported as an absolute + value once, which credited anti-correlation as strongly as correlation, so + it is 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 + + truth = np.arange(40, dtype=float) + assert calculate_metrics(truth, truth)["Spearman"] == pytest.approx(1.0) + assert calculate_metrics(truth, -truth)["Spearman"] == pytest.approx(-1.0) + + rng = np.random.default_rng(0) + noisy = truth + rng.normal(0, 5, truth.size) + forward = calculate_metrics(truth, noisy)["Spearman"] + reversed_ = calculate_metrics(truth, -noisy)["Spearman"] + assert forward > 0 and reversed_ < 0 + assert forward == pytest.approx(-reversed_) From 435ded8149ab9cc9c0c7f428a6b717381b57cd9c Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:53:33 -0400 Subject: [PATCH 09/11] Resolve the masked models to the fill the leaderboard publishes The canonical entries still pointed at the superseded prediction folders and columns, so a default merge followed by performance_fitness.py reproduced the old mutant-fill numbers rather than the leaderboard, which is the kind of quiet disagreement this change set exists to remove. Point RNA-FM, RiNALMo, AIDO.RNA at all five sizes and RNAGenesis at their {name}_4fill folder and their wt_fill column. The {name}_{strategy} entries still read the other three fills from those same files. Verified end to end: merging the released predictions with these folders and running performance_fitness.py --type all reproduces every one of the sixteen leaderboard rows to at most 8.3e-17. Note that the masked models were rescored on the 31 non-coding assays only, so they now report NaN on the mRNA assays rather than a score from a superseded convention. --- fitness/merge_scoring_files.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/fitness/merge_scoring_files.py b/fitness/merge_scoring_files.py index 0fad656..7da33c3 100755 --- a/fitness/merge_scoring_files.py +++ b/fitness/merge_scoring_files.py @@ -117,6 +117,11 @@ def four_fill_entries(name, folder, column_stem): } +# 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", @@ -124,18 +129,18 @@ def four_fill_entries(name, folder, column_stem): "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", } From 0e25e3edd47cbdb909e899a43ea89ee90654f02f Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Thu, 20 Aug 2026 22:54:04 -0400 Subject: [PATCH 10/11] Report AUC and MCC with their direction as well The Spearman became signed in this branch, but the AUC was still folded with max(auc, 1 - auc) and the MCC was still absolute, so one row could report a model as badly wrong by Spearman and moderately good by the other two at the same time. On the non-coding assays 190 of 496 model and assay pairs rank variants backwards, and every one of them was credited with an AUC above 0.5 and a positive MCC; the worst reads Spearman -0.528 alongside AUC 0.747. Return the AUC and the MCC as computed. An AUC below 0.5 or an MCC below 0 now means the model ranks variants the wrong way round, higher stays better for all three, and the leaderboard is unaffected because it reports the Spearman. Extend the metric test to pin all three directions, so a reversed prediction has to score worse than random rather than the same as a correct one. --- fitness/README.md | 4 ++-- fitness/performance_fitness.py | 23 ++++++++++++----------- leaderboard/fitness/README.md | 4 ++-- tests/test_masked_lm.py | 31 ++++++++++++++++++++----------- 4 files changed, 36 insertions(+), 26 deletions(-) diff --git a/fitness/README.md b/fitness/README.md index 4609c5d..2660156 100644 --- a/fitness/README.md +++ b/fitness/README.md @@ -35,8 +35,8 @@ The following directories are gitignored. Download them from ## Metrics - **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 using median-threshold binarisation (max of AUC and 1−AUC) -- **MCC** — absolute Matthews correlation coefficient (median-threshold binarisation) +- **AUC** — AUROC using median-threshold binarisation. Below 0.5 means the model ranks variants the wrong way round +- **MCC** — Matthews correlation coefficient, median-threshold binarisation, signed ## Usage diff --git a/fitness/performance_fitness.py b/fitness/performance_fitness.py index e5df536..5e256ae 100755 --- a/fitness/performance_fitness.py +++ b/fitness/performance_fitness.py @@ -30,23 +30,24 @@ def calculate_metrics( """ Calculate Spearman correlation, AUC, and MCC. - The Spearman correlation is SIGNED. It used to be reported as an absolute - value, which credited a model whose scores anti-correlate with fitness - exactly as much as one that correlates, and that is the distinction the - signed metric adopted in v0.1.1 exists to make. The leaderboard has - published signed values since then, so the absolute form no longer - corresponds to anything the benchmark reports. - - AUC and MCC are still folded onto their better direction below. That is the - same conflation and should be revisited, but it is left alone here because - it would change published numbers this change is not otherwise touching. + 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": 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 efdc2d0..483296b 100644 --- a/leaderboard/fitness/README.md +++ b/leaderboard/fitness/README.md @@ -162,8 +162,8 @@ Reproduce the aggregate with `performance_fitness.py --type ncRNA`, whose Spearm 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. Note that AUC and MCC in that script are still folded onto their better direction, which is the -same conflation and is worth revisiting separately. +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. diff --git a/tests/test_masked_lm.py b/tests/test_masked_lm.py index 2722bb5..d7668e8 100644 --- a/tests/test_masked_lm.py +++ b/tests/test_masked_lm.py @@ -606,23 +606,32 @@ def test_adapters_keep_their_historical_defaults(module_dir, module_name, expect assert defaults.get("dtype") == expect["dtype"] -def test_performance_fitness_reports_a_signed_spearman(): +def test_performance_fitness_metrics_are_directed(): """ - The benchmark metric is signed: a model that ranks variants the wrong way - round must score negative, not positive. This was reported as an absolute - value once, which credited anti-correlation as strongly as correlation, so - it is pinned here rather than left to a reviewer to notice again. + 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 truth = np.arange(40, dtype=float) - assert calculate_metrics(truth, truth)["Spearman"] == pytest.approx(1.0) - assert calculate_metrics(truth, -truth)["Spearman"] == pytest.approx(-1.0) + 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)["Spearman"] - reversed_ = calculate_metrics(truth, -noisy)["Spearman"] - assert forward > 0 and reversed_ < 0 - assert forward == pytest.approx(-reversed_) + 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"] From db9964da31793c2c9dd3417988406e9366c9f404 Mon Sep 17 00:00:00 2001 From: Leo-T-Zang <85547488+Leo-T-Zang@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:00:07 -0400 Subject: [PATCH 11/11] Describe the AUC and MCC binarisation precisely The AUC scores continuous predictions against median-binarised assay labels, while the MCC binarises both sides, which the shared phrase obscured. Note in the metric test why it uses an even, tie-free sample: an MCC of exactly -1 needs the median to split both classes evenly. --- fitness/README.md | 4 ++-- tests/test_masked_lm.py | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/fitness/README.md b/fitness/README.md index 2660156..5eca058 100644 --- a/fitness/README.md +++ b/fitness/README.md @@ -35,8 +35,8 @@ The following directories are gitignored. Download them from ## Metrics - **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 using median-threshold binarisation. Below 0.5 means the model ranks variants the wrong way round -- **MCC** — Matthews correlation coefficient, median-threshold binarisation, signed +- **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/tests/test_masked_lm.py b/tests/test_masked_lm.py index d7668e8..86f67d7 100644 --- a/tests/test_masked_lm.py +++ b/tests/test_masked_lm.py @@ -618,6 +618,10 @@ def test_performance_fitness_metrics_are_directed(): 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)