GWAS tutorial (variant_aware.py — GWAS mode)

GWAS tutorial (variant_aware.py — GWAS 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).

Input FASTA requirements#

  • Sequence: fixed-length window (commonly ~101 nt).

  • Header must encode (variant_pos, ref, alt, strand, seq_start) for parse_variant_block. The 0-based in-window index is idx0 = variant_pos - seq_start.

Output format#

<header_without_>\tPrediction_score:<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)#

We generate a toy 101-nt window with a centered SNV so that the header parsing and allele substitution are runnable without any external files.

from pathlib import Path

def find_repo_root(start: Path) -> Path:
    for p in [start] + list(start.parents):
        if (p / ".git").exists():
            return p
    return start

repo_root = find_repo_root(Path.cwd())
fasta_path = repo_root / "dataset_variant" / "AUH_HepG2.fa"

n_records = 3
count = 0

with fasta_path.open("r") as f:
    while count < n_records:
        header = f.readline()
        if not header:
            break
        seq = f.readline()
        if not seq:
            break
        header = header.strip()
        seq = seq.strip()
        print(f"[{count+1}] {header}")
        print(f"    len={len(seq)} seq[:60]={seq[:60]}{'...' if len(seq)>60 else ''}")
        count += 1

[1] >variant_1 chr1:1014401-1014501(+)[synonymous_variant]{benign} 1014451:C>T ENSG00000187608[0.001]{rs116002608}
    len=101 seq[:60]=GGCCTCAAGCCCCTGAGCACCGTGTTCATGAATCTGCGCCTGCGGGGAGGCGGCACAGAG...
[2] >variant_2 chr1:6599395-6599495(-)[synonymous_variant]{NA} 6599445:G>A ENSG00000162413[0.339]{rs2232460}
    len=101 seq[:60]=CACCTGAGTGGGCTGATCTTCTGCCTGTCATGCGTCCTGGCAGGTGGGTCCGATGGCTCC...
[3] >variant_3 chr1:6599395-6599495(+)[non_coding_transcript_exon_variant]{NA} 6599445:G>A ENSG00000295286[0.339]{rs2232460}
    len=101 seq[:60]=TTCACGCTTGAGTTGTACCTCCACACGCAGTCATAGAGCCGGGAGCCATCGGACCCACCT...

2. Load FASTA (matches variant_aware.py main)#

We reuse the exact FASTA reader used by the unified entry point.

headers, seqs = va.read_fasta(fasta_path)
print("n_records:", len(headers))
print("first header:", headers[0])
print("first seq len:", len(seqs[0]))

n_records: 3
first header: >variant_1 chr1:1014401-1014501(+)[synonymous_variant]{benign} 1014451:C>T ENSG00000187608[0.001]{rs116002608}
first seq len: 101

3. Parse variant metadata from header#

This corresponds to the parse_variant_block(header) call inside process_sequences_gwas.

header0, seq0 = headers[0], seqs[0]
var_pos, ref, alt, strand, seq_start = va.parse_variant_block(header0)
print("var_pos:", var_pos)
print("ref/alt:", ref, "->", alt)
print("strand:", strand)
print("seq_start:", seq_start)

# Strand-aware complement (same logic as in variant_aware.py)
if strand == "-":
    ref, alt = va.apply_complement(ref), va.apply_complement(alt)
print("ref/alt after strand handling:", ref, "->", alt)

var_pos: 1014451
ref/alt: C -> T
strand: +
seq_start: 1014401
ref/alt after strand handling: C -> T

4. Locate the SNV inside the window + sanity checks#

We explicitly show the index math and print the local sequence context.

idx0 = var_pos - seq_start
print("0-based idx0:", idx0)

assert 0 <= idx0 < len(seq0), "Variant index out of bounds."
print("Base at idx0:", seq0[idx0], "| expected REF:", ref)

# Show a small context window around the variant
L = 8
context = seq0[max(0, idx0-L): idx0] + "[" + seq0[idx0] + "]" + seq0[idx0+1: idx0+1+L]
print("Context:", context)

0-based idx0: 50
Base at idx0: C | expected REF: C
Context: CGGGGAGG[C]GGCACAGA

5. Apply allele substitution (REF→ALT) demo#

This is the same substitute_base(...) call used when --variation_mode=after.

mut_seq = va.substitute_base(seq0, idx0, alt)
mut_context = mut_seq[max(0, idx0-L): idx0] + "[" + mut_seq[idx0] + "]" + mut_seq[idx0+1: idx0+1+L]
print("Mutated base at idx0:", mut_seq[idx0])
print("Mutated context:", mut_context)

Mutated base at idx0: T
Mutated context: CGGGGAGG[T]GGCACAGA

6. Run the full GWAS pipeline function#

This calls process_sequences_gwas(...) directly, i.e., the same execution logic as the CLI entry point. Requires real checkpoints; model-id strategy is explained in the CLI help.

If you already have a working BRIDGE environment (dependencies installed and valid Transformer/checkpoint paths), you can run an end-to-end example as follows:

  1. Make sure these three paths exist:

    • FASTA_PATH: your FASTA file (ClinVar/TCGA/1000G catalog variants)

    • TRANSFORMER_PATH: the Transformer directory (e.g. RBPformer)

    • MODEL_SAVE_PATH: the BRIDGE checkpoint directory (containing <model_id>.pth or equivalent naming)

  2. Start with a small file and run --device cpu as a smoke test, then switch to cuda:0 for large-scale scoring.

The output file is appended to. If a file with the same name already exists, new results will be appended at the end.

import argparse
from pathlib import Path
import torch
import os

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("gwas_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,
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
hub = va.ModelHub(Transformer_path, device)

if Transformer_path.exists() and model_save_path.exists():
    va.process_sequences_gwas(headers, seqs, args, hub)
    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: gwas_scores.txt
['variant_1 chr1:1014401-1014501(+)[synonymous_variant]{benign} 1014451:C>T ENSG00000187608[0.001]{rs116002608}\tPrediction_score:-4.160311', 'variant_2 chr1:6599395-6599495(-)[synonymous_variant]{NA} 6599445:G>A ENSG00000162413[0.339]{rs2232460}\tPrediction_score:-3.795815', 'variant_3 chr1:6599395-6599495(+)[non_coding_transcript_exon_variant]{NA} 6599445:G>A ENSG00000295286[0.339]{rs2232460}\tPrediction_score:-5.907292']
Contents