veloVI benchmark on dyngen data

Notebook benchmarks velocity and latent time inference using veloVI on dyngen-generated data.

Library imports

import numpy as np
import pandas as pd

import anndata as ad
from velovi import VELOVI

from rgv_tools import DATA_DIR
from rgv_tools.benchmarking import (
    get_time_correlation,
    get_velocity_correlation,
    set_output,
)

/home/icb/yifan.chen/miniconda3/envs/regvelo-py310/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing read_csv from `anndata` is deprecated. Import anndata.io.read_csv instead.
  warnings.warn(msg, FutureWarning)
/home/icb/yifan.chen/miniconda3/envs/regvelo-py310/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing read_loom from `anndata` is deprecated. Import anndata.io.read_loom instead.
  warnings.warn(msg, FutureWarning)
/home/icb/yifan.chen/miniconda3/envs/regvelo-py310/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing read_text from `anndata` is deprecated. Import anndata.io.read_text instead.
  warnings.warn(msg, FutureWarning)
/home/icb/yifan.chen/miniconda3/envs/regvelo-py310/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing CSCDataset from `anndata.experimental` is deprecated. Import anndata.abc.CSCDataset instead.
  warnings.warn(msg, FutureWarning)
/home/icb/yifan.chen/miniconda3/envs/regvelo-py310/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing CSRDataset from `anndata.experimental` is deprecated. Import anndata.abc.CSRDataset instead.
  warnings.warn(msg, FutureWarning)
/home/icb/yifan.chen/miniconda3/envs/regvelo-py310/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing read_elem from `anndata.experimental` is deprecated. Import anndata.io.read_elem instead.
  warnings.warn(msg, FutureWarning)

Constants

DATASET = "dyngen"

COMPLEXITY = "complexity_1"

SAVE_DATA = True
if SAVE_DATA:
    (DATA_DIR / DATASET / COMPLEXITY / "results").mkdir(parents=True, exist_ok=True)

SAVE_DATASETS = True
if SAVE_DATASETS:
    (DATA_DIR / DATASET / COMPLEXITY / "trained_velovi").mkdir(parents=True, exist_ok=True)

Velocity pipeline

velocity_correlation = []
time_correlation = []

cnt = 0
for filename in (DATA_DIR / DATASET / COMPLEXITY / "processed").iterdir():
    if filename.suffix != ".zarr":
        continue

    simulation_id = int(filename.stem.removeprefix("simulation_"))
    print(f"Run {cnt}, dataset {simulation_id}.")

    adata = ad.io.read_zarr(filename)

    VELOVI.setup_anndata(adata, spliced_layer="Ms", unspliced_layer="Mu")
    vae = VELOVI(adata)
    vae.train(max_epochs=1500)

    set_output(adata, vae, n_samples=30)

    # save data
    adata.write_zarr(DATA_DIR / DATASET / COMPLEXITY / "trained_velovi" / f"trained_{simulation_id}.zarr")

    velocity_correlation.append(
        get_velocity_correlation(
            ground_truth=adata.layers["true_velocity"], estimated=adata.layers["velocity"], aggregation=np.mean
        )
    )

    ## calculate per gene latent time correlation
    time_corr = [
        get_time_correlation(ground_truth=adata.obs["true_time"], estimated=adata.layers["fit_t"][:, i])
        for i in range(adata.layers["fit_t"].shape[1])
    ]
    time_correlation.append(np.mean(time_corr))
    cnt = cnt + 1

Run 0, dataset 29.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
You are using a CUDA device ('NVIDIA A100 80GB PCIe') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -906.788. Signaling Trainer to stop.
Run 1, dataset 14.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -2117.159. Signaling Trainer to stop.
Run 2, dataset 24.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -2541.006. Signaling Trainer to stop.
Run 3, dataset 28.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1478.948. Signaling Trainer to stop.
Run 4, dataset 6.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -3181.971. Signaling Trainer to stop.
Run 5, dataset 21.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1365.964. Signaling Trainer to stop.
Run 6, dataset 15.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1145.076. Signaling Trainer to stop.
Run 7, dataset 9.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -563.872. Signaling Trainer to stop.
Run 8, dataset 12.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -2137.576. Signaling Trainer to stop.
Run 9, dataset 19.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -473.274. Signaling Trainer to stop.
Run 10, dataset 4.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1411.231. Signaling Trainer to stop.
Run 11, dataset 13.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1826.814. Signaling Trainer to stop.
Run 12, dataset 2.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -2857.797. Signaling Trainer to stop.
Run 13, dataset 16.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1858.875. Signaling Trainer to stop.
Run 14, dataset 1.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1070.268. Signaling Trainer to stop.
Run 15, dataset 18.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1159.840. Signaling Trainer to stop.
Run 16, dataset 5.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -348.245. Signaling Trainer to stop.
Run 17, dataset 10.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -2261.989. Signaling Trainer to stop.
Run 18, dataset 8.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1592.854. Signaling Trainer to stop.
Run 19, dataset 11.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -747.262. Signaling Trainer to stop.
Run 20, dataset 27.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1149.416. Signaling Trainer to stop.
Run 21, dataset 23.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -808.871. Signaling Trainer to stop.
Run 22, dataset 17.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1564.348. Signaling Trainer to stop.
Run 23, dataset 30.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1286.337. Signaling Trainer to stop.
Run 24, dataset 22.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1145.037. Signaling Trainer to stop.
Run 25, dataset 25.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1748.977. Signaling Trainer to stop.
Run 26, dataset 20.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1030.632. Signaling Trainer to stop.
Run 27, dataset 7.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1620.434. Signaling Trainer to stop.
Run 28, dataset 3.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -534.901. Signaling Trainer to stop.
Run 29, dataset 26.

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

Monitored metric elbo_validation did not improve in the last 45 records. Best score: -1075.700. Signaling Trainer to stop.

Data saving

if SAVE_DATA:
    pd.DataFrame({"velocity": velocity_correlation, "time": time_correlation}).to_parquet(
        path=DATA_DIR / DATASET / COMPLEXITY / "results" / "velovi_correlation.parquet"
    )