Permutation score analysis - Pre-frontal cortex#

Library imports#

import os
import sys

import numpy as np
import pandas as pd
import torch
from scipy import stats
from sklearn.preprocessing import MinMaxScaler

import matplotlib.pyplot as plt
import mplscience
import seaborn as sns

import scanpy as sc
import scvelo as scv
import scvi
from scvelo.plotting.simulation import compute_dynamics
from velovi import VELOVI

sys.path.append("../..")
from paths import DATA_DIR, FIG_DIR
Global seed set to 0

General settings#

scvi.settings.dl_pin_memory_gpu_training = False
sns.reset_defaults()
sns.reset_orig()
scv.settings.set_figure_params('scvelo', dpi_save=400, dpi=80, transparent=True, fontsize=20, color_map='viridis')
SAVE_FIGURES = True
if SAVE_FIGURES:
    os.makedirs(FIG_DIR / 'permutation' / 'pfc', exist_ok=True)

Function definitions#

def fit_velovi(bdata):
    VELOVI.setup_anndata(bdata, spliced_layer="Ms", unspliced_layer="Mu")

    vae = VELOVI(bdata)
    vae.train()

    df = vae.history["elbo_train"].iloc[20:].reset_index().rename(columns={'elbo_train': 'elbo'})
    df['set'] = 'train'

    _df = vae.history["elbo_validation"].iloc[20:].reset_index().rename(columns={'elbo_validation': 'elbo'})
    _df['set'] = 'validation'

    df = pd.concat([df, _df], axis=0).reset_index(drop=True)

    with mplscience.style_context():
        sns.set_style(style="whitegrid")
        fig, ax = plt.subplots(figsize=(6, 4))
        sns.lineplot(data=df, x='epoch', y='elbo', hue='set', palette=['#0173B2', '#DE8F05'], ax=ax)

    latent_time = vae.get_latent_time(n_samples=25)
    velocities = vae.get_velocity(n_samples=25, velo_statistic="mean")

    t = latent_time
    scaling = 20 / t.max(0)

    bdata.layers["velocities_velovi"] = velocities / scaling
    bdata.layers["latent_time_velovi"] = latent_time

    bdata.var["fit_alpha"] = vae.get_rates()["alpha"] / scaling
    bdata.var["fit_beta"] = vae.get_rates()["beta"] / scaling
    bdata.var["fit_gamma"] = vae.get_rates()["gamma"] / scaling
    bdata.var["fit_t_"] = (
        torch.nn.functional.softplus(vae.module.switch_time_unconstr)
        .detach()
        .cpu()
        .numpy()
    ) * scaling
    bdata.layers["fit_t"] = latent_time.values * scaling[np.newaxis, :]
    bdata.var['fit_scaling'] = 1.0

    return vae
def plot_phase_portrait(adata, gene, color, permuted=False, figsize=(6, 6)):
    fig, ax = plt.subplots(figsize=figsize)

    df = pd.DataFrame(
        {
            'unspliced': adata[:, gene].layers['Mu'].squeeze().copy(),
            'spliced': adata[:, gene].layers['Ms'].squeeze().copy(),
            'color': color
        }
    )

    with mplscience.style_context():
        sns.scatterplot(data=df, x='spliced', y='unspliced', c=color, s=25, ax=ax);

        _, unspliced, spliced = compute_dynamics(adata, basis=gene, extrapolate=True, sort=True)
        df = pd.DataFrame(
            {
                'unspliced': unspliced.squeeze(),
                'spliced': spliced.squeeze(),
            }
        )

        ax.plot(spliced, unspliced, color="purple", linewidth=2)

        spliced_steady_state = np.linspace(np.min(spliced), np.max(spliced))
        unspliced_steady_state = adata.var.loc[gene, 'fit_gamma'] / adata.var.loc[gene, 'fit_beta'] * (spliced_steady_state - np.min(spliced_steady_state)) + np.min(unspliced)
        ax.plot(spliced_steady_state, unspliced_steady_state, color='purple', linestyle="--", linewidth=2);

    ax.axis('off')
    
    if SAVE_FIGURES:
        if permuted:
            fname = f'phase_portrait_{gene}_permuted'
        else:
            fname = f'phase_portrait_{gene}'
        fig.savefig(
            FIG_DIR / 'permutation' / 'pfc' / f'{fname}.svg',
            format="svg",
            transparent=True,
            bbox_inches='tight'
        )
def plot_perm_scores(adata, perm_scores, gene, color_label, figsize=(6, 4)):
    df = pd.DataFrame(perm_scores.loc[gene])
    df["Cell type"] = df.index
    order = adata.obs[color_label].cat.categories.tolist()
    
    with mplscience.style_context():
        sns.set_style(style="whitegrid")
        fig, ax = plt.subplots(figsize=figsize)
        sns.barplot(
            data=df,
            y=gene,
            x="Cell type",
            palette=adata.uns[f"{color_label}_colors"],
            order=order,
            ax=ax,
        )
        ax.set_ylim([-4.5, 30])
        ax.tick_params(axis='x', rotation=45)

    if SAVE_FIGURES:
        fig.savefig(
            FIG_DIR / 'permutation' / 'pfc' / f'permutation_score_{gene}.svg',
            format="svg",
            transparent=True,
            bbox_inches='tight'
        )

Data loading#

adata = sc.read(DATA_DIR / "pfc" / "pfc.h5ad")
adata
AnnData object with n_obs × n_vars = 1267 × 54144
    obs: 'cell_index', 'clusters_coarse', 'clusters', 'X', 'nGene', 'nUMI', 'percent.mito', 'Sample', 'treatment', 'Period', 'stage', 'DevStage'
    uns: 'clusters_colors'
    layers: 'spliced', 'unspliced'

Data preprocessing#

scv.pp.filter_and_normalize(adata, min_shared_counts=20, n_top_genes=2000)
scv.pp.moments(adata, n_pcs=30, n_neighbors=30)

scaler = MinMaxScaler()
adata.layers["Mu"] = scaler.fit_transform(adata.layers["Mu"])

scaler = MinMaxScaler()
adata.layers["Ms"] = scaler.fit_transform(adata.layers["Ms"])
Filtered out 48098 genes that are detected 20 counts (shared).
Normalized count data: X, spliced, unspliced.
Extracted 2000 highly variable genes.
Logarithmized X.
computing neighbors
    finished (0:00:05) --> added 
    'distances' and 'connectivities', weighted adjacency matrices (adata.obsp)
computing moments based on connectivities
    finished (0:00:00) --> added 
    'Ms' and 'Mu', moments of un/spliced abundances (adata.layers)

Model fitting#

velovi_vae = fit_velovi(adata)
/home/icb/philipp.weiler/miniconda3/envs/velovi-py39/lib/python3.9/site-packages/torch/distributed/_sharded_tensor/__init__.py:8: DeprecationWarning: torch.distributed._sharded_tensor will be deprecated, use torch.distributed._shard.sharded_tensor instead
  warnings.warn(
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
Set SLURM handle signals.
Epoch 500/500: 100%|██████████| 500/500 [01:05<00:00,  7.60it/s, loss=-5.68e+03, v_num=1]
../_images/57645270d367f0e550f97fe1e17ceb39a4b8b3617e3fcacd3aac2c8d82f17f3f.png

Permutation score evaluation#

perm_scores, permuted_adata = velovi_vae.get_permutation_scores(labels_key='clusters')
INFO     Input AnnData not setup with scvi-tools. attempting to transfer AnnData setup       
INFO     Input AnnData not setup with scvi-tools. attempting to transfer AnnData setup       
full_perm_df = pd.DataFrame(columns=["Score", "Dataset"])

max_ratio = np.nanmax(perm_scores.values, axis=1)
scores = max_ratio.tolist()
dataset = ['PFC'] * len(max_ratio)

full_perm_df["Score"] = scores
full_perm_df["Dataset"] = dataset
color = adata.obs['clusters'].astype(str).replace(
    dict(zip(adata.obs['clusters'].cat.categories, adata.uns['clusters_colors']))
).tolist()

Lowest ranked genes#

for gene in adata.var_names[np.argsort(scores)[:2]]:
    plot_phase_portrait(adata, gene, color)
    plot_phase_portrait(permuted_adata, gene, color, permuted=True)
    plot_perm_scores(adata, perm_scores, gene, 'clusters')
../_images/565bf8b1bf04766b1b0652e85c41b60cab60442af6cbc3e98cf009f89159bf0c.png ../_images/871353862da1f40ddbbd7ffa4d99b888b195db209ca18e7aaa3f02e0f8c0a29a.png ../_images/3ddab26c79af0d3b3a146d48c7c7e35004cbb40dd57844df856378f79d73725b.png ../_images/d21e0ae51fa813574e6789784b332befa92545b5c25f8102d38b09c7c2f62dfb.png ../_images/3e82726c10b1b5507ccabef9eeccd87d4386e54ce3a0216213c0bf2560dc1db8.png ../_images/ec64b09457a10ec1f1fae6cc19dc1c81b97d4d0f75f21e8e56be409ed325819b.png

Highest ranked genes#

for gene in adata.var_names[np.argsort(scores)[-2:]]:
    plot_phase_portrait(adata, gene, color)
    plot_phase_portrait(permuted_adata, gene, color, permuted=True)
    plot_perm_scores(adata, perm_scores, gene, 'clusters')
../_images/44dc1540c01f8c583a86732dd4c5089a7a79e30f8efae59785b70239cbe590c4.png ../_images/c3c06cb057fe09961d90e113bf9f661ab0f5cf2f51b66643c0414ff785981ea0.png ../_images/1548e499b7d1c7a19704eebd5d01e5fa278351f13da8e6cf3533799bfe3bc75a.png ../_images/f9fe2c742cdce94c2cb6dddf849c5f2bfadec613db9b342c20732d70f9cfcaaf.png ../_images/3d15b66e76ec518d20e8836477b32b9cca9f241efb49e260d1e7fca8d6498836.png ../_images/4dc03eee5b78c9ea09ec37f4cbdce19ef8c32a6df2a17fd231898322cb330557.png