Deriving PSI values from splice junctions

Note

This tutorial uses AlphaGenome 1 (alphagenome.models.v1), but the same logic applies to AlphaGenome 2 (alphagenome.models.v2): pass condition= instead of ontology_terms= to predict_variant, and the returned JunctionData object has the same structure.

This tutorial derives percent spliced in (PSI) values from predicted splice junction counts. For a junction from donor D to acceptor A, PSI5 is the fraction of reads leaving donor D that splice to acceptor A (choice among acceptors for donor D), and PSI3 is the fraction of reads entering acceptor A that come from donor D (choice among donors for acceptor A).

Imports

# @title Install AlphaGenome
# @markdown Run this cell to install AlphaGenome.
from IPython.display import clear_output
! pip install alphagenome
clear_output()

from alphagenome import colab_utils
from alphagenome.data import gene_annotation
from alphagenome.data import genome
from alphagenome.data import transcript as transcript_utils
from alphagenome.models.v1 import dna_client
from alphagenome.visualization import plot_components
import matplotlib.pyplot as plt
import pandas as pd

Load the model

dna_model = dna_client.create(colab_utils.get_api_key())

Predict splice junctions for a variant

Predict splice junctions for a variant near a known splice site:

# Define a variant near a splice site in the APOL4 gene.
variant = genome.Variant(
    chromosome='chr22',
    position=36201698,
    reference_bases='A',
    alternate_bases='C',
)

# Create a 1 Mb interval centered on the variant.
interval = variant.reference_interval.resize(dna_client.SEQUENCE_LENGTH_1MB)

# Predict junction counts for both REF and ALT alleles
output = dna_model.predict_variant(
    interval=interval,
    variant=variant,
    requested_outputs=[dna_client.OutputType.SPLICE_JUNCTIONS],
    ontology_terms=['CL:0000084'],
)

The junction data holds one row per predicted junction and one column per RNA-seq track for T cells (CL:0000084):

print(f"Junctions shape: {output.reference.splice_junctions.junctions.shape}")
print(f"Value shape: {output.reference.splice_junctions.values.shape}")

Junctions shape: (107,)
Value shape: (107, 2)

Junction counts come from RNA-seq, and T cells have two RNA-seq assays: poly(A)+ RNA-seq, which enriches for mature spliced mRNA, and total RNA-seq, which also captures pre-mRNA and can show incomplete splicing.

output.reference.splice_junctions.metadata

	name	ontology_curie	biosample_name	biosample_type	biosample_life_stage	gtex_tissue	data_source	Assay title
0	junction_CL:0000084 polyA plus RNA-seq	CL:0000084	T-cell	primary_cell	adult		encode	polyA plus RNA-seq
1	junction_CL:0000084 total RNA-seq	CL:0000084	T-cell	primary_cell	adult		encode	total RNA-seq

Derive the PSI values

Create a function to calculate PSI values from splice junction counts for either RNA track:

def compute_psi_from_raw(junction_data, track_idx):
  # Create the dataframe from known attributes
  data = []
  for i, j in enumerate(junction_data.junctions):
    # Determine Donor and Acceptor based on strand
    if j.strand == '+':
      donor, acceptor = j.start, j.end
    else:
      donor, acceptor = j.end, j.start

    data.append({
        'donor': donor,
        'acceptor': acceptor,
        'count': junction_data.values[i, track_idx],
    })

  df = pd.DataFrame(data)

  # Pivot into the matrix: Rows = Donors, Cols = Acceptors
  matrix = df.pivot_table(
      index='donor', columns='acceptor', values='count', fill_value=0
  )

  # PSI5: Normalize row-wise (Percent of donor usage for each acceptor)
  # PSI5 = Count / (Total counts for that Donor)
  psi5 = matrix.div(matrix.sum(axis=1), axis=0).fillna(0)

  # PSI3: Normalize column-wise (Percent of acceptor usage for each donor)
  # PSI3 = Count / (Total counts for that Acceptor)
  psi3 = matrix.div(matrix.sum(axis=0), axis=1).fillna(0)

  return psi5, psi3

Select RNA track and calculate:

# Select RNA track index and compute PSI5 and PSI3 for the polyA plus RNA-seq track
titles = (
    output.reference.splice_junctions.metadata['Assay title']
    .astype(str)
    .tolist()
)
polya_idx = next(
    (i for i, t in enumerate(titles) if 'polya' in t.lower()), None
)
total_idx = next(
    (i for i, t in enumerate(titles) if 'total' in t.lower()), None
)

print(f'Indices found -> PolyA: {polya_idx}, Total: {total_idx}')

psi5_polya, psi3_polya = compute_psi_from_raw(
    output.reference.splice_junctions, polya_idx
)

Indices found -> PolyA: 0, Total: 1

Inspect results

print(
    'The proportion of total splicing activity at a specific acceptor site that'
    ' is contributed by a given donor'
)
display(psi3_polya.head())

print(
    'The proportion of total splicing activity at a specific donor site that is'
    ' directed to a given acceptor'
)
display(psi5_polya.head())

The proportion of total splicing activity at a specific acceptor site that is contributed by a given donor
The proportion of total splicing activity at a specific donor site that is directed to a given acceptor

acceptor	36191912	36195432	36195437	36197831	36198896	36199369	36199376	36199608	36200291	36201796	36201851	36201888	36202055	36202094	36202933
donor
36192925	0.018771	0.000000	0.000000	0.000000	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
36195094	0.016950	0.000000	0.000000	0.000000	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
36195310	0.909040	0.000000	0.000000	0.000000	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
36195912	0.004808	0.093865	0.016693	0.000000	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
36197944	0.005055	0.087138	0.017416	0.100716	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0

acceptor	36191912	36195432	36195437	36197831	36198896	36199369	36199376	36199608	36200291	36201796	36201851	36201888	36202055	36202094	36202933
donor
36192925	1.000000	0.000000	0.000000	0.00000	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
36195094	1.000000	0.000000	0.000000	0.00000	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
36195310	1.000000	0.000000	0.000000	0.00000	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
36195912	0.153819	0.215299	0.630882	0.00000	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0
36197944	0.130356	0.161100	0.530524	0.17802	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0	0.0

Visualize junctions with a Sashimi plot

For splicing, it helps to show alternative isoforms, so we annotate all protein-coding transcripts rather than MANE Select only:

# The GTF file contains information on the location of all transcripts.
# Note that we use genome assembly hg38 for human.
gtf = pd.read_feather(
    'https://storage.googleapis.com/alphagenome/reference/gencode/'
    'hg38/gencode.v46.annotation.gtf.gz.feather'
)

transcript_extractor = transcript_utils.TranscriptExtractor(
    gene_annotation.filter_protein_coding(gtf)
)

Build the Sashimi plot. Numbers above the arcs are predicted normalized junction read counts, not PSI values:

# We will plot for both polyA and total RNA-seq
ref_output = output.reference
alt_output = output.alternate

# Plot
_ = plot_components.plot(
    [
        # Transcript annotations
        plot_components.TranscriptAnnotation(
            transcript_extractor.extract(interval)
        ),
        # Sashimi ref and alt plots
        plot_components.Sashimi(
            ref_output.splice_junctions,
            ylabel_template='Reference {biosample_name} ({strand})\n{name}',
            color='skyblue',
        ),
        plot_components.Sashimi(
            alt_output.splice_junctions,
            ylabel_template='Alternate {biosample_name} ({strand})\n{name}',
            color='red',
        ),
    ],
    # Annotate the variant position as a vertical line across all panels
    annotations=[plot_components.VariantAnnotation([variant])],
    interval=interval.resize(40000),
)

plt.show()