{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Stability Analysis\n",
"\n",
"Stability Analysis\n",
"\n",
"Analyze local stability of cellular states using Jacobian matrices and eigenvalue decomposition.\n",
"\n",
"## Overview\n",
"\n",
"The Jacobian matrix captures local dynamics at each cell state:\n",
"\n",
"\n",
"\n",
"$$\n",
"J = W \\cdot \\text{diag}\\left(\\frac{d\\sigma}{dx}\\right) - \\text{diag}(\\gamma)\n",
"$$\n",
"\n",
"Eigenvalues of J indicate:\n",
"\n",
"- **Positive real part**: Unstable directions (repulsion)\n",
"- **Negative real part**: Stable directions (attraction)\n",
"- **Imaginary part**: Oscillatory dynamics\n",
"\n",
"## Computing Jacobians\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The Jacobian matrix **J** at a cell state **x** characterises local dynamics:\n",
"\n",
"\n",
"\n",
"$$J_{ij}(x) = \\frac{\\partial \\dot{x}_i}{\\partial x_j} = W_{ij} \\cdot \\sigma'(x_j) - \\gamma_i \\delta_{ij}$$\n",
"\n",
"\n",
"\n",
"Key stability indicators:\n",
"\n",
"- **Trace(J)** < 0 → contracting (stabilising) dynamics\n",
"\n",
"- **Leading Re(λ)** < 0 → local attractor; > 0 → local repeller / saddle\n",
"\n",
"- **Rotational part** → oscillatory tendency\n",
"\n",
"\n",
"\n",
"Topics covered:\n",
"\n",
"1. Compute Jacobian matrices\n",
"\n",
"2. Jacobian summary statistics on UMAP\n",
"\n",
"3. Eigenvalue spectra per cluster\n",
"\n",
"4. Rotational dynamics\n",
"\n",
"5. Element-wise Jacobian analysis"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/tmp/ipykernel_1221845/3765484981.py:17: UserWarning: adata has 1956 genes but the model was trained on 1728. A subsetted copy is being returned; the original adata is NOT modified. Reassign the return value:\n",
" adata = sch.tl.load_model(adata, filename)\n",
" adata = sch.tl.load_model(adata, MODEL_FILE)\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Model loaded from 'model.h5sch' | clusters=['Bas', 'Ery', 'GMP-like', 'HSC', 'MEP-like', 'Meg', 'Mon', 'Neu', 'all'] | genes=1728\n",
"AnnData object with n_obs × n_vars = 1947 × 1728\n",
" obs: 'batch', 'time', 'cell_type', 'nGenes', 'nCounts', 'pMito', 'pass_basic_filter', 'new_Size_Factor', 'initial_new_cell_size', 'total_Size_Factor', 'initial_total_cell_size', 'spliced_Size_Factor', 'initial_spliced_cell_size', 'unspliced_Size_Factor', 'initial_unspliced_cell_size', 'Size_Factor', 'initial_cell_size', 'ntr', 'cell_cycle_phase', 'leiden', 'control_point_pca', 'inlier_prob_pca', 'obs_vf_angle_pca', 'pca_ddhodge_div', 'pca_ddhodge_potential', 'acceleration_pca', 'curvature_pca', 'n_counts', 'mt_frac', 'jacobian_det_pca', 'manual_selection', 'divergence_pca', 'curv_leiden', 'curv_louvain', 'SPI1->GATA1_jacobian', 'jacobian', 'umap_ori_leiden', 'umap_ori_louvain', 'umap_ddhodge_div', 'umap_ddhodge_potential', 'curl_umap', 'divergence_umap', 'acceleration_umap', 'control_point_umap_ori', 'inlier_prob_umap_ori', 'obs_vf_angle_umap_ori', 'curvature_umap_ori'\n",
" var: 'gene_name', 'gene_id', 'nCells', 'nCounts', 'pass_basic_filter', 'use_for_pca', 'frac', 'ntr', 'time_3_alpha', 'time_3_beta', 'time_3_gamma', 'time_3_half_life', 'time_3_alpha_b', 'time_3_alpha_r2', 'time_3_gamma_b', 'time_3_gamma_r2', 'time_3_gamma_logLL', 'time_3_delta_b', 'time_3_delta_r2', 'time_3_bs', 'time_3_bf', 'time_3_uu0', 'time_3_ul0', 'time_3_su0', 'time_3_sl0', 'time_3_U0', 'time_3_S0', 'time_3_total0', 'time_3_beta_k', 'time_3_gamma_k', 'time_5_alpha', 'time_5_beta', 'time_5_gamma', 'time_5_half_life', 'time_5_alpha_b', 'time_5_alpha_r2', 'time_5_gamma_b', 'time_5_gamma_r2', 'time_5_gamma_logLL', 'time_5_bs', 'time_5_bf', 'time_5_uu0', 'time_5_ul0', 'time_5_su0', 'time_5_sl0', 'time_5_U0', 'time_5_S0', 'time_5_total0', 'time_5_beta_k', 'time_5_gamma_k', 'use_for_dynamics', 'gamma', 'gamma_r2', 'use_for_transition', 'gamma_k', 'gamma_b', 'I_Bas', 'I_Ery', 'I_GMP-like', 'I_HSC', 'I_MEP-like', 'I_Meg', 'I_Mon', 'I_Neu', 'I_all', 'scHopfield_used', 'sigmoid_exponent', 'sigmoid_mse', 'sigmoid_offset', 'sigmoid_threshold'\n",
" uns: 'PCs', 'VecFld_pca', 'VecFld_umap', 'X_umap_neighbors', 'cell_phase_genes', 'cell_type_colors', 'dynamics', 'explained_variance_ratio_', 'feature_selection', 'grid_velocity_pca', 'grid_velocity_umap', 'grid_velocity_umap_ori_perturbation', 'grid_velocity_umap_test', 'jacobian_pca', 'leiden', 'neighbors', 'pca_mean', 'pp', 'response', 'scHopfield'\n",
" obsm: 'X', 'X_pca', 'X_pca_SparseVFC', 'X_umap', 'X_umap_SparseVFC', 'X_umap_ori_perturbation', 'X_umap_test', 'acceleration_pca', 'acceleration_umap', 'cell_cycle_scores', 'curvature_pca', 'curvature_umap', 'j_delta_x_perturbation', 'velocity_pca', 'velocity_pca_SparseVFC', 'velocity_umap', 'velocity_umap_SparseVFC', 'velocity_umap_ori_perturbation', 'velocity_umap_test'\n",
" layers: 'M_n', 'M_nn', 'M_t', 'M_tn', 'M_tt', 'X_new', 'X_total', 'velocity_alpha_minus_gamma_s'\n",
" obsp: 'X_umap_connectivities', 'X_umap_distances', 'connectivities', 'cosine_transition_matrix', 'distances', 'fp_transition_rate', 'moments_con', 'pca_ddhodge', 'perturbation_transition_matrix', 'umap_ddhodge'\n",
" varp: 'W_Bas', 'W_Ery', 'W_GMP-like', 'W_HSC', 'W_MEP-like', 'W_Meg', 'W_Mon', 'W_Neu', 'W_all'\n",
"Using device: cuda\n"
]
}
],
"source": [
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"import scanpy as sc\n",
"import torch\n",
"import scHopfield as sch\n",
" \n",
"DATA_PATH = './scratch/Data/'\n",
"DATASET_FILE = 'hematopoiesis.h5ad'\n",
"MODEL_FILE = 'model.h5sch'\n",
"\n",
"CLUSTER_KEY = 'cell_type'\n",
"SPLICED_KEY = 'M_t'\n",
"\n",
"CELL_TYPE_ORDER = ['Meg', 'Ery', 'MEP-like', 'HSC', 'GMP-like', 'Mon', 'Bas', 'Neu']\n",
"\n",
"adata = sc.read_h5ad(DATA_PATH + DATASET_FILE)\n",
"adata = sch.tl.load_model(adata, MODEL_FILE)\n",
"print(adata)\n",
"\n",
"colors = dict(zip(CELL_TYPE_ORDER, adata.uns['cell_type_colors']))\n",
"device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
"print(f'Using device: {device}')\n",
"\n",
"\n",
"# Set seed for reproducibility\n",
"np.random.seed(42)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4.1 Compute Jacobian Matrices\n",
"\n",
"\n",
"\n",
"**Warning**: This is memory-intensive for large datasets.\n",
"\n",
"Save immediately after computation."
]
},
{
"cell_type": "code",
"execution_count": 34,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Computing Jacobians for cluster Mon\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 423/423 [00:10<00:00, 40.81it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Computing Jacobians for cluster Meg\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 154/154 [00:03<00:00, 43.17it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Computing Jacobians for cluster MEP-like\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 457/457 [00:11<00:00, 39.39it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Computing Jacobians for cluster Ery\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 234/234 [00:05<00:00, 42.71it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Computing Jacobians for cluster Bas\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 177/177 [00:04<00:00, 41.94it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Computing Jacobians for cluster GMP-like\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 161/161 [00:03<00:00, 40.32it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Computing Jacobians for cluster HSC\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 309/309 [00:07<00:00, 40.43it/s]\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Computing Jacobians for cluster Neu\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 32/32 [00:00<00:00, 41.02it/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Jacobian eigenvalues stored in adata.obsm['jacobian_eigenvalues']\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n"
]
}
],
"source": [
"sch.tl.compute_jacobians(\n",
" adata,\n",
" spliced_key=SPLICED_KEY,\n",
" cluster_key=CLUSTER_KEY,\n",
" compute_eigenvectors=False, # Set True only if eigenvectors are needed\n",
" device=device\n",
")\n",
"\n",
"print(\"Jacobian eigenvalues stored in adata.obsm['jacobian_eigenvalues']\")\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Save to disk to avoid re-computation\n",
"# sch.tl.save_jacobians(\n",
"# adata,\n",
"# filename='jacobian_eigenvalues.h5',\n",
"# cluster_key=CLUSTER_KEY,\n",
"# compression='gzip'\n",
"# )\n",
"# print(\"Jacobians saved to disk.\")\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Load from disk later (if previously saved):\n",
"# sch.tl.load_jacobians(\n",
"# adata,\n",
"# filename='jacobian_eigenvalues.h5',\n",
"# load_eigenvectors=False\n",
"# )\n",
"# sch.tl.compute_jacobian_stats(adata, store_in_obs=True)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4.2 Jacobian Statistics\n",
"\n",
"\n",
"\n",
"Extracts per-cell scalar summaries from eigenvalue arrays and stores them in\n",
"\n",
"`adata.obs`."
]
},
{
"cell_type": "code",
"execution_count": 35,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Jacobian statistics per cell:\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" jacobian_eig1_real | \n",
" jacobian_eig1_imag | \n",
" jacobian_positive_evals | \n",
" jacobian_negative_evals | \n",
" jacobian_trace | \n",
"
\n",
" \n",
" \n",
" \n",
" | count | \n",
" 1947.000000 | \n",
" 1947.0 | \n",
" 1947.000000 | \n",
" 1947.000000 | \n",
" 1947.000000 | \n",
"
\n",
" \n",
" | mean | \n",
" -0.068015 | \n",
" 0.0 | \n",
" 17.364664 | \n",
" 1710.635336 | \n",
" -296.985982 | \n",
"
\n",
" \n",
" | std | \n",
" 0.031078 | \n",
" 0.0 | \n",
" 3.917776 | \n",
" 3.917776 | \n",
" 2.714663 | \n",
"
\n",
" \n",
" | min | \n",
" -0.646356 | \n",
" 0.0 | \n",
" 4.000000 | \n",
" 1698.000000 | \n",
" -306.789470 | \n",
"
\n",
" \n",
" | 25% | \n",
" -0.064679 | \n",
" 0.0 | \n",
" 15.000000 | \n",
" 1708.000000 | \n",
" -298.732645 | \n",
"
\n",
" \n",
" | 50% | \n",
" -0.064679 | \n",
" 0.0 | \n",
" 17.000000 | \n",
" 1711.000000 | \n",
" -297.255203 | \n",
"
\n",
" \n",
" | 75% | \n",
" -0.064679 | \n",
" 0.0 | \n",
" 20.000000 | \n",
" 1713.000000 | \n",
" -295.487251 | \n",
"
\n",
" \n",
" | max | \n",
" -0.005080 | \n",
" 0.0 | \n",
" 30.000000 | \n",
" 1724.000000 | \n",
" -287.853683 | \n",
"
\n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"# Visualise on UMAP\n",
"sc.pl.umap(\n",
" adata,\n",
" color=['jacobian_trace', 'jacobian_positive_evals', 'jacobian_eig1_real'],\n",
" ncols=3,\n",
" cmap='RdBu_r',\n",
" vcenter=0\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 37,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Jacobian Statistics Summary:\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" jacobian_positive_evals | \n",
" jacobian_negative_evals | \n",
" jacobian_trace | \n",
"
\n",
" \n",
" | \n",
" mean | \n",
" std | \n",
" mean | \n",
" std | \n",
" mean | \n",
" std | \n",
"
\n",
" \n",
" | cell_type | \n",
" | \n",
" | \n",
" | \n",
" | \n",
" | \n",
" | \n",
"
\n",
" \n",
" \n",
" \n",
" | Bas | \n",
" 15.492 | \n",
" 3.010 | \n",
" 1712.508 | \n",
" 3.010 | \n",
" -297.896 | \n",
" 1.907 | \n",
"
\n",
" \n",
" | Ery | \n",
" 17.350 | \n",
" 3.096 | \n",
" 1710.650 | \n",
" 3.096 | \n",
" -296.432 | \n",
" 2.021 | \n",
"
\n",
" \n",
" | GMP-like | \n",
" 16.832 | \n",
" 3.218 | \n",
" 1711.168 | \n",
" 3.218 | \n",
" -297.408 | \n",
" 1.934 | \n",
"
\n",
" \n",
" | HSC | \n",
" 17.317 | \n",
" 3.969 | \n",
" 1710.683 | \n",
" 3.969 | \n",
" -297.905 | \n",
" 1.739 | \n",
"
\n",
" \n",
" | MEP-like | \n",
" 17.998 | \n",
" 3.471 | \n",
" 1710.002 | \n",
" 3.471 | \n",
" -297.345 | \n",
" 1.852 | \n",
"
\n",
" \n",
" | Meg | \n",
" 14.649 | \n",
" 2.138 | \n",
" 1713.351 | \n",
" 2.138 | \n",
" -298.837 | \n",
" 1.759 | \n",
"
\n",
" \n",
" | Mon | \n",
" 19.390 | \n",
" 3.950 | \n",
" 1708.610 | \n",
" 3.950 | \n",
" -294.455 | \n",
" 2.827 | \n",
"
\n",
" \n",
" | Neu | \n",
" 8.219 | \n",
" 3.405 | \n",
" 1719.781 | \n",
" 3.405 | \n",
" -304.415 | \n",
" 2.780 | \n",
"
\n",
" \n",
"
\n",
"