VeloVAE (fullvb) benchmark on dyngen data

VeloVAE (fullvb) benchmark on dyngen data#

Notebook benchmarks velocity and latent time inference using VeloVAE (fullvb) on dyngen-generated data.

Library imports#

import velovae as vv

import numpy as np
import pandas as pd
import torch

import anndata as ad
import scvelo as scv

from rgv_tools import DATA_DIR
from rgv_tools.benchmarking import get_velocity_correlation
2024-12-14 01:30:12.923192: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-12-14 01:30:12.938701: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1734136212.957171 3032484 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1734136212.962640 3032484 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2024-12-14 01:30:12.983527: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
/home/icb/weixu.wang/miniconda3/envs/regvelo_test/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/weixu.wang/miniconda3/envs/regvelo_test/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing read_excel from `anndata` is deprecated. Import anndata.io.read_excel instead.
  warnings.warn(msg, FutureWarning)
/home/icb/weixu.wang/miniconda3/envs/regvelo_test/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing read_hdf from `anndata` is deprecated. Import anndata.io.read_hdf instead.
  warnings.warn(msg, FutureWarning)
/home/icb/weixu.wang/miniconda3/envs/regvelo_test/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/weixu.wang/miniconda3/envs/regvelo_test/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing read_mtx from `anndata` is deprecated. Import anndata.io.read_mtx instead.
  warnings.warn(msg, FutureWarning)
/home/icb/weixu.wang/miniconda3/envs/regvelo_test/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/weixu.wang/miniconda3/envs/regvelo_test/lib/python3.10/site-packages/anndata/utils.py:429: FutureWarning: Importing read_umi_tools from `anndata` is deprecated. Import anndata.io.read_umi_tools instead.
  warnings.warn(msg, FutureWarning)

General settings#

scv.settings.verbosity = 3

Constants#

torch.manual_seed(0)
np.random.seed(0)
DATASET = "dyngen"
SAVE_DATA = True
if SAVE_DATA:
    (DATA_DIR / DATASET / "results").mkdir(parents=True, exist_ok=True)
    (DATA_DIR / DATASET / "processed" / "velovae_vae").mkdir(parents=True, exist_ok=True)

Velocity pipeline#

velocity_correlation = []

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

    adata = ad.io.read_zarr(filename)

    try:
        rate_prior = {"alpha": (0.0, 1.0), "beta": (0.0, 0.5), "gamma": (0.0, 0.5)}
        vae = vv.VAE(adata, tmax=20, dim_z=5, device="cuda:0", full_vb=True, rate_prior=rate_prior)
        config = {}
        vae.train(adata, config=config, plot=False, embed="pca")

        # Output velocity to adata object
        vae.save_anndata(adata, "fullvb", DATA_DIR / DATASET / "processed" / "velovae_vae", file_name="velovae.h5ad")

        adata.layers["velocity"] = adata.layers["fullvb_velocity"].copy()

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

    except Exception as e:  # noqa: BLE001
        # Append np.nan in case of an error and optionally log the error
        print(f"An error occurred: {e}")
        velocity_correlation.append(np.nan)
Estimating ODE parameters...
Detected 59 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
(0.42, 0.2890740448785467), (0.58, 0.7827494240291312)
KS-test result: [1. 0. 1.]
Initial induction: 84, repression: 21/105
Learning Rate based on Data Sparsity: 0.0001
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 6, test iteration: 10
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.040
Average Set Size: 20
Change in noise variance: 0.0455
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 1931.    *********
Change in noise variance: 0.0177
Change in x0: 0.2938
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 2287.    *********
Change in noise variance: 0.0026
Change in x0: 0.2331
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 2363.    *********
Change in noise variance: 0.0029
Change in x0: 0.1918
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 2389.    *********
Change in noise variance: 0.0009
Change in x0: 0.1528
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 2425.    *********
Change in noise variance: 0.0000
Change in x0: 0.0943
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 2437.    *********
Change in noise variance: 0.0000
Change in x0: 0.0717
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 2614.    *********
Change in noise variance: 0.0000
Change in x0: 0.0580
*********             Velocity Refinement Round 9             *********
*********     Round 9: Early Stop Triggered at epoch 2718.    *********
Change in noise variance: 0.0000
Change in x0: 0.0564
*********             Velocity Refinement Round 10             *********
Stage 2: Early Stop Triggered at round 9.
*********              Finished. Total Time =   0 h :  5 m : 26 s             *********
Final: Train ELBO = 91.494,	Test ELBO = 84.699
Estimating ODE parameters...
Detected 25 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
KS-test result: [1. 2. 1.]
Initial induction: 35, repression: 24/59
Learning Rate based on Data Sparsity: 0.0000
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 3, test iteration: 4
*********       Stage 1: Early Stop Triggered at epoch 6.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.058
Average Set Size: 18
*********     Round 1: Early Stop Triggered at epoch 13.    *********
Change in noise variance: 0.2737
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 87.    *********
Change in noise variance: 0.0610
Change in x0: 0.8688
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 94.    *********
Change in noise variance: 0.0019
Change in x0: 0.3654
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 101.    *********
Change in noise variance: 0.0010
Change in x0: 0.2517
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 108.    *********
Change in noise variance: 0.0000
Change in x0: 0.2153
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 134.    *********
Change in noise variance: 0.0000
Change in x0: 0.2039
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 157.    *********
Change in noise variance: 0.0000
Change in x0: 0.1966
*********             Velocity Refinement Round 8             *********
Stage 2: Early Stop Triggered at round 7.
*********              Finished. Total Time =   0 h :  0 m : 20 s             *********
Final: Train ELBO = -13651.812,	Test ELBO = -13606.135
Estimating ODE parameters...
Detected 37 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
(0.39, 0.1200757575748825), (0.61, 0.7766228691348986)
KS-test result: [0. 1. 1.]
Initial induction: 73, repression: 13/86
Learning Rate based on Data Sparsity: 0.0000
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 2, test iteration: 2
*********       Stage 1: Early Stop Triggered at epoch 8.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.057
Average Set Size: 16
*********     Round 1: Early Stop Triggered at epoch 14.    *********
Change in noise variance: 0.3585
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 19.    *********
Change in noise variance: 0.0663
Change in x0: 0.2847
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 24.    *********
Change in noise variance: 0.0642
Change in x0: 0.2418
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 30.    *********
Change in noise variance: 0.0075
Change in x0: 0.1653
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 36.    *********
Change in noise variance: 0.0122
Change in x0: 0.1416
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 42.    *********
Change in noise variance: 0.0046
Change in x0: 0.1361
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 48.    *********
Change in noise variance: 0.0025
Change in x0: 0.1304
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 54.    *********
Change in noise variance: 0.0020
Change in x0: 0.1185
*********             Velocity Refinement Round 9             *********
*********     Round 9: Early Stop Triggered at epoch 59.    *********
Change in noise variance: 0.0012
Change in x0: 0.1085
*********             Velocity Refinement Round 10             *********
*********     Round 10: Early Stop Triggered at epoch 64.    *********
Change in noise variance: 0.0006
Change in x0: 0.1040
*********             Velocity Refinement Round 11             *********
Stage 2: Early Stop Triggered at round 10.
*********              Finished. Total Time =   0 h :  0 m : 19 s             *********
Final: Train ELBO = -18201.777,	Test ELBO = -18397.561
Estimating ODE parameters...
Detected 35 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
KS-test result: [1. 1. 1.]
Assign cluster 0 to repressive
Initial induction: 50, repression: 31/81
Learning Rate based on Data Sparsity: 0.0000
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 2, test iteration: 2
*********       Stage 1: Early Stop Triggered at epoch 458.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.120
Average Set Size: 32
Change in noise variance: 0.1604
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 1025.    *********
Change in noise variance: 0.1941
Change in x0: 1.2177
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 1111.    *********
Change in noise variance: 0.0061
Change in x0: 0.9228
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 1161.    *********
Change in noise variance: 0.0068
Change in x0: 0.4832
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 1240.    *********
Change in noise variance: 0.0064
Change in x0: 0.2474
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 1332.    *********
Change in noise variance: 0.0030
Change in x0: 0.1310
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 1522.    *********
Change in noise variance: 0.0028
Change in x0: 0.0764
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 1551.    *********
Change in noise variance: 0.0008
Change in x0: 0.0919
*********             Velocity Refinement Round 9             *********
Stage 2: Early Stop Triggered at round 8.
*********              Finished. Total Time =   0 h :  1 m : 15 s             *********
Final: Train ELBO = -2424.000,	Test ELBO = -2521.890
Estimating ODE parameters...
Detected 35 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
KS-test result: [1. 1. 1.]
Assign cluster 0 to repressive
Initial induction: 51, repression: 29/80
Learning Rate based on Data Sparsity: 0.0000
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 2, test iteration: 2
*********       Stage 1: Early Stop Triggered at epoch 7.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.033
Average Set Size: 17
*********     Round 1: Early Stop Triggered at epoch 13.    *********
Change in noise variance: 0.3382
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 322.    *********
Change in noise variance: 0.0977
Change in x0: 0.5931
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 521.    *********
Change in noise variance: 0.0098
Change in x0: 0.3892
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 591.    *********
Change in noise variance: 0.0020
Change in x0: 0.3135
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 669.    *********
Change in noise variance: 0.0009
Change in x0: 0.2734
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 712.    *********
Change in noise variance: 0.0000
Change in x0: 0.2453
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 722.    *********
Change in noise variance: 0.0000
Change in x0: 0.2222
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 743.    *********
Change in noise variance: 0.0000
Change in x0: 0.2082
*********             Velocity Refinement Round 9             *********
*********     Round 9: Early Stop Triggered at epoch 780.    *********
Change in noise variance: 0.0000
Change in x0: 0.2067
*********             Velocity Refinement Round 10             *********
Stage 2: Early Stop Triggered at round 9.
*********              Finished. Total Time =   0 h :  0 m : 53 s             *********
Final: Train ELBO = -14614.414,	Test ELBO = -14580.146
Estimating ODE parameters...
Detected 67 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
KS-test result: [1. 1. 1.]
Assign cluster 1 to repressive
Initial induction: 72, repression: 37/109
Learning Rate based on Data Sparsity: 0.0001
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 6, test iteration: 10
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.062
Average Set Size: 20
Change in noise variance: 0.0388
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 1599.    *********
Change in noise variance: 0.0101
Change in x0: 0.2792
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 1628.    *********
Change in noise variance: 0.0018
Change in x0: 0.2092
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 1757.    *********
Change in noise variance: 0.0020
Change in x0: 0.1938
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 1768.    *********
Change in noise variance: 0.0011
Change in x0: 0.1624
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 1874.    *********
Change in noise variance: 0.0003
Change in x0: 0.1271
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 1906.    *********
Change in noise variance: 0.0000
Change in x0: 0.1167
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 1930.    *********
Change in noise variance: 0.0000
Change in x0: 0.0996
*********             Velocity Refinement Round 9             *********
*********     Round 9: Early Stop Triggered at epoch 1986.    *********
Change in noise variance: 0.0000
Change in x0: 0.0815
*********             Velocity Refinement Round 10             *********
*********     Round 10: Early Stop Triggered at epoch 2030.    *********
Change in noise variance: 0.0000
Change in x0: 0.0677
*********             Velocity Refinement Round 11             *********
*********     Round 11: Early Stop Triggered at epoch 2047.    *********
Change in noise variance: 0.0000
Change in x0: 0.0574
*********             Velocity Refinement Round 12             *********
*********     Round 12: Early Stop Triggered at epoch 2064.    *********
Change in noise variance: 0.0000
Change in x0: 0.0484
*********             Velocity Refinement Round 13             *********
Stage 2: Early Stop Triggered at round 12.
*********              Finished. Total Time =   0 h :  3 m : 51 s             *********
Final: Train ELBO = 113.724,	Test ELBO = 119.769
Estimating ODE parameters...
Detected 18 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
(0.47, 0.8701520422929223), (0.53, 0.2900017161456463)
KS-test result: [1. 2. 0.]
Initial induction: 20, repression: 41/61
Learning Rate based on Data Sparsity: 0.0001
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 2, test iteration: 2
*********       Stage 1: Early Stop Triggered at epoch 5.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.030
Average Set Size: 16
*********     Round 1: Early Stop Triggered at epoch 116.    *********
Change in noise variance: 0.5613
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 151.    *********
Change in noise variance: 0.0224
Change in x0: 0.5994
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 157.    *********
Change in noise variance: 0.0021
Change in x0: 0.4098
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 170.    *********
Change in noise variance: 0.0018
Change in x0: 0.3452
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 176.    *********
Change in noise variance: 0.0009
Change in x0: 0.2928
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 182.    *********
Change in noise variance: 0.0000
Change in x0: 0.2531
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 188.    *********
Change in noise variance: 0.0000
Change in x0: 0.2259
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 194.    *********
Change in noise variance: 0.0000
Change in x0: 0.2183
*********             Velocity Refinement Round 9             *********
Stage 2: Early Stop Triggered at round 8.
*********              Finished. Total Time =   0 h :  0 m : 25 s             *********
Final: Train ELBO = -13386.965,	Test ELBO = -13522.012
Estimating ODE parameters...
Detected 34 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
KS-test result: [1. 1. 1.]
Assign cluster 2 to repressive
Initial induction: 45, repression: 18/63
Learning Rate based on Data Sparsity: 0.0000
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 2, test iteration: 2
*********       Stage 1: Early Stop Triggered at epoch 8.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.027
Average Set Size: 17
*********     Round 1: Early Stop Triggered at epoch 14.    *********
Change in noise variance: 0.1701
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 19.    *********
Change in noise variance: 0.0061
Change in x0: 0.6860
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 25.    *********
Change in noise variance: 0.0006
Change in x0: 0.4838
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 31.    *********
Change in noise variance: 0.0000
Change in x0: 0.3766
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 36.    *********
Change in noise variance: 0.0000
Change in x0: 0.2992
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 41.    *********
Change in noise variance: 0.0000
Change in x0: 0.2497
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 46.    *********
Change in noise variance: 0.0000
Change in x0: 0.2236
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 51.    *********
Change in noise variance: 0.0000
Change in x0: 0.2116
*********             Velocity Refinement Round 9             *********
*********     Round 9: Early Stop Triggered at epoch 61.    *********
Change in noise variance: 0.0000
Change in x0: 0.2000
*********             Velocity Refinement Round 10             *********
*********     Round 10: Early Stop Triggered at epoch 73.    *********
Change in noise variance: 0.0000
Change in x0: 0.1865
*********             Velocity Refinement Round 11             *********
*********     Round 11: Early Stop Triggered at epoch 78.    *********
Change in noise variance: 0.0000
Change in x0: 0.1758
*********             Velocity Refinement Round 12             *********
*********     Round 12: Early Stop Triggered at epoch 90.    *********
Change in noise variance: 0.0000
Change in x0: 0.1698
*********             Velocity Refinement Round 13             *********
Stage 2: Early Stop Triggered at round 12.
*********              Finished. Total Time =   0 h :  0 m : 21 s             *********
Final: Train ELBO = -15680.239,	Test ELBO = -15800.924
Estimating ODE parameters...
Detected 32 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
(0.56, 0.3239376112994556), (0.44, 0.9106431281299979)
(0.52, 0.2605808478636782), (0.48, 0.805314977942315)
KS-test result: [1. 0. 0.]
Initial induction: 43, repression: 37/80
Learning Rate based on Data Sparsity: 0.0000
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 2, test iteration: 2
*********       Stage 1: Early Stop Triggered at epoch 396.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.010
Average Set Size: 15
*********     Round 1: Early Stop Triggered at epoch 402.    *********
Change in noise variance: 0.4139
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 408.    *********
Change in noise variance: 0.3447
Change in x0: 0.7088
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 444.    *********
Change in noise variance: 0.0068
Change in x0: 0.2237
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 450.    *********
Change in noise variance: 0.0017
Change in x0: 0.1605
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 456.    *********
Change in noise variance: 0.0015
Change in x0: 0.0966
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 462.    *********
Change in noise variance: 0.0017
Change in x0: 0.0794
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 468.    *********
Change in noise variance: 0.0010
Change in x0: 0.0909
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 474.    *********
Change in noise variance: 0.0009
Change in x0: 0.0889
*********             Velocity Refinement Round 9             *********
Stage 2: Early Stop Triggered at round 8.
*********              Finished. Total Time =   0 h :  0 m : 38 s             *********
Final: Train ELBO = -3722.105,	Test ELBO = -3882.429
Estimating ODE parameters...
Detected 20 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
(0.53, 0.06264460485018672), (0.47, 0.7010424470363079)
(0.58, 0.707388107196621), (0.42, 0.09198273775867975)
KS-test result: [0. 0. 2.]
Initial induction: 27, repression: 49/76
Learning Rate based on Data Sparsity: 0.0000
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 2, test iteration: 2
*********       Stage 1: Early Stop Triggered at epoch 7.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********
Percentage of Invalid Sets: 0.030
Average Set Size: 15
*********     Round 1: Early Stop Triggered at epoch 13.    *********
Change in noise variance: 0.3139
*********             Velocity Refinement Round 2             *********
*********     Round 2: Early Stop Triggered at epoch 144.    *********
Change in noise variance: 0.0408
Change in x0: 0.5010
*********             Velocity Refinement Round 3             *********
*********     Round 3: Early Stop Triggered at epoch 180.    *********
Change in noise variance: 0.0073
Change in x0: 0.5771
*********             Velocity Refinement Round 4             *********
*********     Round 4: Early Stop Triggered at epoch 186.    *********
Change in noise variance: 0.0269
Change in x0: 0.3684
*********             Velocity Refinement Round 5             *********
*********     Round 5: Early Stop Triggered at epoch 192.    *********
Change in noise variance: 0.0052
Change in x0: 0.3283
*********             Velocity Refinement Round 6             *********
*********     Round 6: Early Stop Triggered at epoch 198.    *********
Change in noise variance: 0.0027
Change in x0: 0.2698
*********             Velocity Refinement Round 7             *********
*********     Round 7: Early Stop Triggered at epoch 204.    *********
Change in noise variance: 0.0018
Change in x0: 0.2269
*********             Velocity Refinement Round 8             *********
*********     Round 8: Early Stop Triggered at epoch 210.    *********
Change in noise variance: 0.0011
Change in x0: 0.2009
*********             Velocity Refinement Round 9             *********
*********     Round 9: Early Stop Triggered at epoch 216.    *********
Change in noise variance: 0.0005
Change in x0: 0.1882
*********             Velocity Refinement Round 10             *********
*********     Round 10: Early Stop Triggered at epoch 221.    *********
Change in noise variance: 0.0000
Change in x0: 0.1845
*********             Velocity Refinement Round 11             *********
Stage 2: Early Stop Triggered at round 10.
*********              Finished. Total Time =   0 h :  0 m : 26 s             *********
Final: Train ELBO = -13665.999,	Test ELBO = -13649.533
Estimating ODE parameters...
Detected 35 velocity genes.
Estimating the variance...
Initialization using the steady-state and dynamical models.
Reinitialize the regular ODE parameters based on estimated global latent time.
3 clusters detected based on gene co-expression.
(0.56, 0.24297477479231852), (0.44, 0.7779664767114342)
KS-test result: [0. 1. 1.]
Initial induction: 60, repression: 23/83
Learning Rate based on Data Sparsity: 0.0000
--------------------------- Train a VeloVAE ---------------------------
*********        Creating Training/Validation Datasets        *********
*********                      Finished.                      *********
*********                 Creating optimizers                 *********
*********                      Finished.                      *********
*********                    Start training                   *********
*********                      Stage  1                       *********
Total Number of Iterations Per Epoch: 2, test iteration: 2
*********       Stage 1: Early Stop Triggered at epoch 9.       *********
*********                      Stage  2                       *********
*********             Velocity Refinement Round 1             *********

Data saving#

if SAVE_DATA:
    pd.DataFrame({"velocity": velocity_correlation}).to_parquet(
        path=DATA_DIR / DATASET / "results" / "velovae_fullvb_correlation.parquet"
    )