Ribosnitch tutorial (variant_aware.py — ribosnitch mode)

Ribosnitch tutorial (variant_aware.py — ribosnitch mode)#

This notebook is a tutorialized, step-by-step walkthrough of the unified variant_aware.py entry point. It decomposes the run into annotated, inspectable steps (explicit I/O expectations + lightweight sanity checks such as printing tensor shapes), while keeping the underlying execution logic identical to the script (we call the same functions and helpers).

CLI arguments#

Common arguments#

Argument	Type / choices	Required	Meaning
`--variation_mode`	`before / after`	yes	Score reference window (‘before’) or ALT-substituted window (‘after’).
`--fasta_sequence_path`	`PATH`	yes	Input FASTA containing window sequences.
`--variant_out_file`	`PATH`	yes	Output file (appended).
`--Transformer_path`	`PATH`	yes	Transformer path used by build_Transformer_embeddings.
`--model_save_path`	`PATH`	yes	Directory with BRIDGE checkpoints (*.pth).
`--device`	`cuda / cuda:N / cpu`	no	Torch device (default: cuda if available else cpu).

Pipeline selection flags#

Argument	Type / choices	Required	Meaning
`--gwas`	`-`	no	Force GWAS pipeline (default if no pipeline flag provided).
`--ribosnitch/--ribosnitches`	`-`	no	Enable ribosnitch pipeline.
`--catalog_variants/--genomic_variants`	`-`	no	Enable catalog-variants pipeline (ClinVar/TCGA/1000G).

Ribosnitch-only arguments#

Argument	Type / choices	Required	Meaning
`--ribosnitch_after_variation/--ribosnitches_after_variation`	`-`	no	Force ALT substitution in ribosnitch mode regardless of –variation_mode.
`--ribosnitch_out_dir/--ribosnitches_out_dir`	`PATH`	no	Root output directory for ribosnitch results (default ./results/ribosnitches).

Input FASTA requirements (Ribosnitch)#

Header must end with two cell-line tokens (e.g., HEK293 HepG2).
If ALT substitution is used, header must also contain tokens parsable by parse_variant_block_flexible.

Output format (Ribosnitch)#

<header_without_>\t<checkpoint_stem>\t<float>

0. Imports + repo bootstrap#

from __future__ import annotations

import sys
from pathlib import Path

def find_repo_root(start: Path | None = None) -> Path:
    # Heuristically locate the BRIDGE repo root so we can import `variant_aware.py` and `utils/`.
    # Works when this notebook lives under docs/tutorials/notebooks/.
    start = (start or Path().resolve())
    for p in [start, *start.parents]:
        if (p / "variant_aware.py").exists() and (p / "utils").exists():
            return p
    raise FileNotFoundError(
        "Cannot locate repo root (expected to find variant_aware.py and utils/). "
        "Run this notebook from within the BRIDGE repository."
    )

REPO_ROOT = find_repo_root()
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))
print("Repo root:", REPO_ROOT)

import variant_aware as va
from utils import variant as var_utils

print("Imported:", va.__file__)

Repo root: /fs1/private/user/wangyubo/code/BRIDGE
Imported: /fs1/private/user/wangyubo/code/BRIDGE/variant_aware.py

1. Minimal runnable demo data (FASTA + dummy checkpoint names)#

We create a toy FASTA header that ends with two cell-line tokens. This makes _extract_cell_lines runnable. We also create an empty model directory with dummy .pth filenames purely to demonstrate checkpoint selection logic. (Actual scoring requires real checkpoints.)

from pathlib import Path

TMP = Path("./_tmp_tutorial")
TMP.mkdir(exist_ok=True)

from pathlib import Path

# Ribosnitch header requirements in `variant_aware.py`:
# - last two tokens are cell-line names (used to select *_<cell>.pth checkpoints)
# - header also contains a region token + SNV token so we can run the mutation demo.
header = ">variant_1 chr19:11120155-11120255(+)[Familial_hypercholesterolemia|not_specified] 11120205:T>C HEK293 HepG2"
seq = "TCCGATGTCAACTTGTTGGCTGAAAACCTACTGTCCCCAGAGGATATGGTTCTCTTCCACAACCTCACCCAGCCAAGAGGTAAGGGTGGGTCAGCCCCACC"

fasta_path = Path("ribosnitch.fa")
with fasta_path.open("w") as f:
    f.write(header + "\n")
    f.write(seq + "\n")

# Dummy model dir for selection demo
model_dir = Path("../../../results/model")
# model_dir.mkdir(exist_ok=True)
# for fn in ["RBP1_HepG2.pth", "RBP2_K562.pth", "OtherCell_GM12878.pth"]:
#     (model_dir / fn).touch()

print("FASTA:", fasta_path)
# print("Model files:", sorted([p.name for p in model_dir.iterdir()]))

FASTA: ribosnitch.fa

2. Load FASTA#

headers, seqs = va.read_fasta(fasta_path)
header0 = headers[0].lstrip(">")
seq0 = seqs[0]
print("header0:", header0)
print("seq0 len:", len(seq0))

header0: variant_1 chr19:11120155-11120255(+)[Familial_hypercholesterolemia|not_specified] 11120205:T>C HEK293 HepG2
seq0 len: 101

3. Step 1 in ribosnitch: extract cell lines#

This matches _extract_cell_lines(...).

cell1, cell2 = va._extract_cell_lines(header0)
print("cell lines:", cell1, cell2)

cell lines: HEK293 HepG2

4. Step 2: select checkpoints by suffix#

This matches the model_files = [...] selection logic.

import os
from pathlib import Path

model_files = [
    fn for fn in os.listdir(model_dir)
    if fn.endswith(f"_{cell1}.pth") or fn.endswith(f"_{cell2}.pth")
]
print("selected checkpoints:", model_files)

selected checkpoints: ['FIP1L1_HEK293.pth', 'HNRNPD_HEK293.pth', 'CSTF2T_HepG2.pth', 'UCHL5_HepG2.pth', 'HNRNPC_HepG2.pth', 'FXR2_HEK293.pth', 'PTBP1_HepG2.pth', 'BCLAF1_HepG2.pth', 'CDC40_HepG2.pth', 'IGF2BP1_HEK293.pth', 'PRPF8_HepG2.pth', 'BCCIP_HepG2.pth', 'EIF3H_HepG2.pth', 'AUH_HepG2.pth', 'IGF2BP1_HepG2.pth', 'SRSF1_HepG2.pth', 'SRSF9_HepG2.pth', 'CSTF2_HEK293.pth', 'RBPMS_HEK293.pth', 'FKBP4_HepG2.pth', 'TNRC6A_HEK293.pth', 'RTCB_HEK293.pth', 'PCBP1_HepG2.pth', 'SUPV3L1_HepG2.pth', 'NUDT21_HEK293.pth', 'TAF15_HEK293.pth', 'FAM120A_HepG2.pth', 'HNRNPUL1_HepG2.pth', 'SUB1_HepG2.pth', 'SMNDC1_HepG2.pth', 'YBX3_HepG2.pth', 'FXR2_HepG2.pth', 'EWSR1_HEK293.pth', 'HNRNPM_HepG2.pth', 'IGF2BP3_HepG2.pth', 'TRA2A_HepG2.pth', 'U2AF1_HepG2.pth', 'QKI_HEK293.pth', 'FBL_HEK293.pth', 'FASTKD2_HepG2.pth', 'DDX55_HepG2.pth', 'PCBP2_HepG2.pth', 'DDX3X_HepG2.pth', 'DDX52_HepG2.pth', 'IGF2BP2_HEK293.pth', 'HLTF_HepG2.pth', 'AGO_HEK293.pth', 'WDR33_HEK293.pth', 'TIA1_HepG2.pth', 'FUS_HEK293.pth', 'ELAVL1_HEK293.pth', 'RBM22_HepG2.pth', 'NOP56_HEK293.pth', 'AKAP1_HepG2.pth', 'ZNF800_HepG2.pth', 'TAF15_HepG2.pth', 'UTP18_HepG2.pth', 'EXOSC5_HepG2.pth', 'DKC1_HepG2.pth', 'NCBP2_HepG2.pth', 'G3BP1_HepG2.pth', 'EIF3D_HepG2.pth', 'NOLC1_HepG2.pth', 'LSM11_HepG2.pth', 'PABPN1_HepG2.pth', 'NKRF_HepG2.pth', 'DDX59_HepG2.pth', 'XRN2_HepG2.pth', 'HNRNPA1_HepG2.pth', 'RPS3_HepG2.pth', 'UPF1_HepG2.pth', 'SRSF7_HepG2.pth', 'NOL12_HepG2.pth', 'LIN28B_HepG2.pth', 'DHX30_HepG2.pth', 'HNRNPU_HepG2.pth', 'LIN28B_HEK293.pth', 'AQR_HepG2.pth', 'PUM2_HEK293.pth', 'EFTUD2_HepG2.pth', 'CPSF6_HEK293.pth', 'AGGF1_HepG2.pth', 'C22ORF28_HEK293.pth', 'ZC3H11A_HepG2.pth', 'NOP58_HEK293.pth', 'SDAD1_HepG2.pth', 'FMR1_HEK293.pth', 'WDR43_HepG2.pth', 'ZC3H7B_HEK293.pth', 'GTF2F1_HepG2.pth', 'DDX6_HepG2.pth', 'CAPRIN1_HEK293.pth', 'RBFOX2_HepG2.pth', 'SF3B4_HepG2.pth', 'ALKBH5_HEK293.pth', 'FXR1_HEK293.pth', 'BUD13_HepG2.pth', 'CPSF4_HEK293.pth', 'DGCR8_HepG2.pth', 'SLTM_HepG2.pth', 'SND1_HepG2.pth', 'CPSF2_HEK293.pth', 'CPSF7_HEK293.pth', 'RBM15_HepG2.pth', 'CPSF3_HEK293.pth', 'MOV10_HEK293.pth', 'CPSF1_HEK293.pth', 'LIN28A_HEK293.pth', 'CSTF2T_HEK293.pth', 'GRWD1_HepG2.pth', 'SF3A3_HepG2.pth', 'IGF2BP3_HEK293.pth', 'DROSHA_HepG2.pth', 'HNRNPK_HepG2.pth', 'U2AF2_HepG2.pth', 'TBRG4_HepG2.pth', 'LARP4_HepG2.pth', 'PPIG_HepG2.pth', 'PRPF4_HepG2.pth', 'LARP7_HepG2.pth', 'C17ORF85_HEK293.pth', 'FTO_HepG2.pth', 'NIP7_HepG2.pth']

5. Step 3: mutation demo (`--variation_mode=after`)#

This uses the flexible parser and shows strand-aware substitution.

# This is exactly the helper used by the ribosnitch-after branch.
mut_seq = va._maybe_mutate_sequence_from_header(headers[0], seq0)

# Print context around the mutated site
# We recompute idx0 to show the same sanity-check math
var_pos, ref, alt, strand, seq_start = va.parse_variant_block_flexible(headers[0])
if strand == "-":
    ref, alt = va.apply_complement(ref), va.apply_complement(alt)
idx0 = var_pos - seq_start

L = 8
ctx_before = seq0[max(0, idx0-L): idx0] + "[" + seq0[idx0] + "]" + seq0[idx0+1: idx0+1+L]
ctx_after  = mut_seq[max(0, idx0-L): idx0] + "[" + mut_seq[idx0] + "]" + mut_seq[idx0+1: idx0+1+L]
print("before:", ctx_before)
print("after :", ctx_after)

before: GATATGGT[T]CTCTTCCA
after : GATATGGT[C]CTCTTCCA

6. Run the full ribosnitch pipeline function#

This calls run_ribosnitches(...) directly (same execution logic as the CLI). Requires real .pth checkpoints in model_save_path.

import argparse
from pathlib import Path
import torch

Transformer_path = Path(os.environ.get("BRIDGE_TRANSFORMER_PATH", "../../../RBPformer"))
model_save_path = Path(os.environ.get("BRIDGE_MODEL_SAVE_PATH", "../../../results/model"))
variant_out_file = Path("ribosnitch_scores.txt")

args = argparse.Namespace(
    variation_mode="after",
    fasta_sequence_path=fasta_path,
    variant_out_file=variant_out_file,
    Transformer_path=Transformer_path,
    model_save_path=model_save_path,
    ribosnitch=True,
    ribosnitch_after_variation=False,
    ribosnitch_out_dir=Path("ribosnitch_results"),
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

if Transformer_path.exists() and model_save_path.exists():
    va.run_ribosnitches(headers, seqs, args, device)
    print("Wrote scores to:", variant_out_file)
    print(variant_out_file.read_text().splitlines()[:3])
else:
    print("Skip: set BRIDGE_TRANSFORMER_PATH and BRIDGE_MODEL_SAVE_PATH.")

Wrote scores to: ribosnitch_scores.txt
['variant_1 chr19:11120155-11120255(+)[Familial_hypercholesterolemia|not_specified] 11120205:T>C HEK293 HepG2\tFIP1L1_HEK293\t-7.150293827056885', 'variant_1 chr19:11120155-11120255(+)[Familial_hypercholesterolemia|not_specified] 11120205:T>C HEK293 HepG2\tHNRNPD_HEK293\t0.31320297718048096', 'variant_1 chr19:11120155-11120255(+)[Familial_hypercholesterolemia|not_specified] 11120205:T>C HEK293 HepG2\tCSTF2T_HepG2\t-2.3324456214904785']

The Kernel crashed while executing code in the current cell or a previous cell. 

Please review the code in the cell(s) to identify a possible cause of the failure. 

Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. 

View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details.