Note: This environment variable is required for fully deterministic CuBLAS ops on CUDA >= 10.2 when reproducible=True is set below. Without it, PyTorch raises a RuntimeError instead of training deterministically. It must be set before torch is imported. See the README FAQ for details.

In [ ]:

%env CUBLAS_WORKSPACE_CONFIG=:16:8

%env CUBLAS_WORKSPACE_CONFIG=:16:8

Tutorial showing usage of pretrained autoencoder models¶

This tutorial focus on the usage of pretrained models for scRNASeq data and provided under: https://huggingface.co/collections/autoencodix/acx-pretrained-models

Download the model files¶

In this tutorial we will use the a Ontix model with explainable latent space based on ontologies generated with Gemini3ProPreview:

https://huggingface.co/autoencodix/Ontix-Dim24-Gemini3ProPreview

In [1]:

import os
from huggingface_hub import snapshot_download

model_name = "Ontix-Dim24-Gemini3ProPreview"
repo_id = f"autoencodix/{model_name}"
target_folder = f"./acx_pretrained_models/{model_name}"
private = True  # Set to True if the repository is private
if private:
	token = os.environ.get("HF_TOKEN")
	if token is None:
		raise EnvironmentError(
			"HF_TOKEN is not set."
		)
else:
	token = None

local_dir = snapshot_download(
    repo_id=repo_id,
    local_dir=target_folder,
    local_dir_use_symlinks=False,
    token=token,  # uncomment if repo is private
)

print(f"All repo files downloaded to: {local_dir}")

import os from huggingface_hub import snapshot_download model_name = "Ontix-Dim24-Gemini3ProPreview" repo_id = f"autoencodix/{model_name}" target_folder = f"./acx_pretrained_models/{model_name}" private = True # Set to True if the repository is private if private: token = os.environ.get("HF_TOKEN") if token is None: raise EnvironmentError( "HF_TOKEN is not set." ) else: token = None local_dir = snapshot_download( repo_id=repo_id, local_dir=target_folder, local_dir_use_symlinks=False, token=token, # uncomment if repo is private ) print(f"All repo files downloaded to: {local_dir}")

Fetching 6 files:   0%|          | 0/6 [00:00<?, ?it/s]

All repo files downloaded to: /home/alicia/dev/biomarker_autoencoder/autoencodix_package/Tutorials/DeepDives/acx_pretrained_models/Ontix-Dim24-Gemini3ProPreview

/home/alicia/dev/biomarker_autoencoder/autoencodix_package/.venv/lib/python3.10/site-packages/huggingface_hub/file_download.py:986: UserWarning: `local_dir_use_symlinks` parameter is deprecated and will be ignored. The process to download files to a local folder has been updated and do not rely on symlinks anymore. You only need to pass a destination folder as`local_dir`.
For more details, check out https://huggingface.co/docs/huggingface_hub/main/en/guides/download#download-files-to-local-folder.
  warnings.warn(

Download and prepare example scRNASeq data¶

To show how to use scRNASeq with pretrained models, we use an example lung tissue dataset originally extracted from the CZ CELLxGENE Census database. It is hosted on Hugging Face Hub (autoencodix/census-lung) and is downloaded automatically in the cell below on first run — no manual download needed.

As an example we will compare malignant cells from lung adenocarcinoma to pulmonary alveolar type 2 (AT2) cells from normal lung tissue — AT2 cells are the well-established cell-of-origin for lung adenocarcinoma, so this contrasts tumor cells with their normal counterpart.

In [2]:

import pandas as pd

# Get the gene space of the autoencoder to subset to only the genes it was trained on
ont_file = f"./acx_pretrained_models/{model_name}/ontology/Dim24_Gemini3ProPreview_ontology_task__ensembl_level2.tsv"
ont_features = pd.read_csv(ont_file, sep='\t', usecols=[0], header=None)
ont_features.columns = ['feature_id']

# Define the observation filter to get only the relevant cell types and disease states
obs_value_filter = "tissue_general == 'lung' and cell_type in ['malignant cell', 'pulmonary alveolar type 2 cell'] and is_primary_data == True"

import pandas as pd # Get the gene space of the autoencoder to subset to only the genes it was trained on ont_file = f"./acx_pretrained_models/{model_name}/ontology/Dim24_Gemini3ProPreview_ontology_task__ensembl_level2.tsv" ont_features = pd.read_csv(ont_file, sep='\t', usecols=[0], header=None) ont_features.columns = ['feature_id'] # Define the observation filter to get only the relevant cell types and disease states obs_value_filter = "tissue_general == 'lung' and cell_type in ['malignant cell', 'pulmonary alveolar type 2 cell'] and is_primary_data == True"

In [3]:

from huggingface_hub import hf_hub_download
import anndata as ad

census_path = hf_hub_download(
    repo_id="autoencodix/census-lung", repo_type="dataset", filename="census_lung.h5ad"
)
adata = ad.read_h5ad(census_path)
adata = adata[adata.obs.query(obs_value_filter).index, adata.var.feature_id.isin(ont_features.feature_id)].copy()

adata

from huggingface_hub import hf_hub_download import anndata as ad census_path = hf_hub_download( repo_id="autoencodix/census-lung", repo_type="dataset", filename="census_lung.h5ad" ) adata = ad.read_h5ad(census_path) adata = adata[adata.obs.query(obs_value_filter).index, adata.var.feature_id.isin(ont_features.feature_id)].copy() adata

Out[3]:

AnnData object with n_obs × n_vars = 19908 × 8189
    obs: 'soma_joinid', 'dataset_id', 'assay', 'assay_ontology_term_id', 'cell_type', 'cell_type_ontology_term_id', 'development_stage', 'development_stage_ontology_term_id', 'disease', 'disease_ontology_term_id', 'donor_id', 'is_primary_data', 'observation_joinid', 'self_reported_ethnicity', 'self_reported_ethnicity_ontology_term_id', 'sex', 'sex_ontology_term_id', 'suspension_type', 'tissue', 'tissue_ontology_term_id', 'tissue_type', 'tissue_general', 'tissue_general_ontology_term_id', 'raw_sum', 'nnz', 'raw_mean_nnz', 'raw_variance_nnz', 'n_measured_vars'
    var: 'soma_joinid', 'feature_id', 'feature_name', 'feature_type', 'feature_length', 'nnz', 'n_measured_obs'

Importantly, models have been trained on log1p normalized counts. We do this manually and then create a data package for autoencodix:

In [4]:

import scanpy 
from autoencodix.data._numeric_dataset import NumericDataset
from autoencodix.data._datasetcontainer import DatasetContainer
from autoencodix.configs.ontix_config import OntixConfig

scanpy.pp.log1p(adata, copy=False)

test_dataset = NumericDataset(
			data=adata.X, # the normalized expression data per cell
			config=OntixConfig(), # Empty placeholder
			sample_ids=adata.obs.index, # cell ids
			metadata=adata.obs.loc[adata.obs.index,:], # cell annotation data for plotting etc.
			split_indices=None, # Empty placeholder
			feature_ids=adata.var.feature_id, # Gene ids
		)

acx_container = DatasetContainer(train=None, valid=None, test=test_dataset)

# Save the container 
import pickle
# Create directory
os.makedirs("./TutData", exist_ok=True)

with open(f"./TutData/scRNASeq_malignant_vs_AT2_{model_name}.pkl", "wb") as f:
	pickle.dump(acx_container, f)

import scanpy from autoencodix.data._numeric_dataset import NumericDataset from autoencodix.data._datasetcontainer import DatasetContainer from autoencodix.configs.ontix_config import OntixConfig scanpy.pp.log1p(adata, copy=False) test_dataset = NumericDataset( data=adata.X, # the normalized expression data per cell config=OntixConfig(), # Empty placeholder sample_ids=adata.obs.index, # cell ids metadata=adata.obs.loc[adata.obs.index,:], # cell annotation data for plotting etc. split_indices=None, # Empty placeholder feature_ids=adata.var.feature_id, # Gene ids ) acx_container = DatasetContainer(train=None, valid=None, test=test_dataset) # Save the container import pickle # Create directory os.makedirs("./TutData", exist_ok=True) with open(f"./TutData/scRNASeq_malignant_vs_AT2_{model_name}.pkl", "wb") as f: pickle.dump(acx_container, f)

Calculate embeddings using a pretrained model¶

Now we can load the data and a model to calculate embeddings and visualize them

In [5]:

model_name = "Ontix-Dim24-Gemini3ProPreview"
import autoencodix as acx
import pickle

ontix_file_path = f"./acx_pretrained_models/{model_name}/large_ontix_final_model_Dim24_Gemini3ProPreview_ontology_task__.pkl"

print("Loading trained model ...")
loaded_ontix = acx.Ontix.load(file_path=ontix_file_path)
loaded_ontix._trainer._config.save_vram = True  # Enable memory saving for prediction

print("Load test data container ...")
with open(f"./TutData/scRNASeq_malignant_vs_AT2_{model_name}.pkl", "rb") as f:
	acx_container = pickle.load(f)

print("Calculate embeddings ...")
result = loaded_ontix.predict(data=acx_container)

model_name = "Ontix-Dim24-Gemini3ProPreview" import autoencodix as acx import pickle ontix_file_path = f"./acx_pretrained_models/{model_name}/large_ontix_final_model_Dim24_Gemini3ProPreview_ontology_task__.pkl" print("Loading trained model ...") loaded_ontix = acx.Ontix.load(file_path=ontix_file_path) loaded_ontix._trainer._config.save_vram = True # Enable memory saving for prediction print("Load test data container ...") with open(f"./TutData/scRNASeq_malignant_vs_AT2_{model_name}.pkl", "rb") as f: acx_container = pickle.load(f) print("Calculate embeddings ...") result = loaded_ontix.predict(data=acx_container)

Loading trained model ...
Attempting to load a pipeline from acx_pretrained_models/Ontix-Dim24-Gemini3ProPreview/large_ontix_final_model_Dim24_Gemini3ProPreview_ontology_task__.pkl...
Pipeline object loaded successfully. Actual type: Ontix
Preprocessor loaded successfully.
Load test data container ...
Calculate embeddings ...
Processed 19908 / 19908 samples

In [ ]:

# Create a latent space visualization
loaded_ontix.visualizer.show_latent_space(
	result=result,
	plot_type="Ridgeline",
	param=['cell_type'], 
	split="test",
)

# Create a latent space visualization loaded_ontix.visualizer.show_latent_space( result=result, plot_type="Ridgeline", param=['cell_type'], split="test", )

Using explain() functionality for xAI¶

We can leverage posthoc xAI and LLMs to get a better understanding which genes and processes drive differences between malignant cells and their normal cell-of-origin, alveolar type 2 (AT2) cells

For this we will firstly determine which latent dimensions show highest class separation:

In [ ]:

# Perform Linear Discriminant Analysis (LDA) to find the latent dimension that best separates the groups
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.model_selection import cross_validate

import numpy as np
import pandas as pd

anno_col = "cell_type"
groups_list = [
	"malignant cell",
	"pulmonary alveolar type 2 cell"
]
# Get the latent space as df
df_latent = loaded_ontix.result.get_latent_df(split="test", epoch=-1)

scores = {}
# for each latent dimension
for latent_dim in df_latent.columns:
	lda = LinearDiscriminantAnalysis()
	X = df_latent[latent_dim].values.reshape(-1, 1)
	y = loaded_ontix.result.new_datasets.test.metadata[anno_col].values
	# Reduce x and y to only include samples from the groups of interest
	mask = np.isin(y, groups_list)
	X = X[mask]
	y = y[mask]
	metric = "roc_auc_ovo"
 	# Do CV-5 fold cross-validation and calculate the AUC-ROC score for this latent dimension
	scores[latent_dim] = cross_validate(lda, X, y, cv=5, scoring=metric, return_train_score=True, n_jobs=-1)

scores_df = pd.DataFrame({
	'latent_dim': list(scores.keys()),
	'test_score_mean': [scores[ld]['test_score'].mean() for ld in scores.keys()],
	'test_score_std': [scores[ld]['test_score'].std() for ld in scores.keys()],
	'train_score_mean': [scores[ld]['train_score'].mean() for ld in scores.keys()],
	'train_score_std': [scores[ld]['train_score'].std() for ld in scores.keys()],
})
scores_df = scores_df.sort_values(by='test_score_mean', ascending=False)
scores_df

# Perform Linear Discriminant Analysis (LDA) to find the latent dimension that best separates the groups from sklearn.discriminant_analysis import LinearDiscriminantAnalysis from sklearn.model_selection import cross_validate import numpy as np import pandas as pd anno_col = "cell_type" groups_list = [ "malignant cell", "pulmonary alveolar type 2 cell" ] # Get the latent space as df df_latent = loaded_ontix.result.get_latent_df(split="test", epoch=-1) scores = {} # for each latent dimension for latent_dim in df_latent.columns: lda = LinearDiscriminantAnalysis() X = df_latent[latent_dim].values.reshape(-1, 1) y = loaded_ontix.result.new_datasets.test.metadata[anno_col].values # Reduce x and y to only include samples from the groups of interest mask = np.isin(y, groups_list) X = X[mask] y = y[mask] metric = "roc_auc_ovo" # Do CV-5 fold cross-validation and calculate the AUC-ROC score for this latent dimension scores[latent_dim] = cross_validate(lda, X, y, cv=5, scoring=metric, return_train_score=True, n_jobs=-1) scores_df = pd.DataFrame({ 'latent_dim': list(scores.keys()), 'test_score_mean': [scores[ld]['test_score'].mean() for ld in scores.keys()], 'test_score_std': [scores[ld]['test_score'].std() for ld in scores.keys()], 'train_score_mean': [scores[ld]['train_score'].mean() for ld in scores.keys()], 'train_score_std': [scores[ld]['train_score'].std() for ld in scores.keys()], }) scores_df = scores_df.sort_values(by='test_score_mean', ascending=False) scores_df

The highest separation of malignant cells vs. AT2 cells is observed for the top-ranked latent dimensions in the table above. Now let's calculatue the feature (gene) attribution to those dimensions and identify potential marker genes and get a functional analyses via an LLM.

To use the openrouter API store make yor API key available as environment variable OPENROUTER_PREMIUM_API_KEY. Accordingly, using SCADS-LLM server with API key under SCADS_LLM_API_KEY. See also LLM_Setup.md.

In [ ]:

selected_dim = list(scores_df.iloc[0:2]['latent_dim'].values) # Top latent dimensions with highest AUC-ROC score for separating the groups
top_n_genes = 15 # Number of top contributing genes per latent dimension

latent_contributions = loaded_ontix.explain(
    split="test",
    method="IntegratedGradients", # "DeepLiftShap" or "IntegratedGradients"
    n_subset=400, # Randomly sample in input and baseline space for computational efficiency
    sel_latent_dim=selected_dim, # Or specify a single latent dimension to explain    
	input_type="grouped", # We will test group 'malignant cell' (input) vs. 'pulmonary alveolar type 2 cell' (baseline)
	input_group=groups_list[0], # "malignant cell"
	baseline_type="random",
	baseline_group=groups_list[1], # "pulmonary alveolar type 2 cell"
	anno_col=anno_col,
	llm_explain=True, 	# Explain biological functions of the contributing genes using a LLM
	llm_client="openrouter",
	llm_model="google/gemma-4-31b-it",
	# llm_client="scads-llm",
	# llm_model="google/gemma-4-31b-it",
	# llm_model="alias-ha",
	top_n_genes=top_n_genes,
)

selected_dim = list(scores_df.iloc[0:2]['latent_dim'].values) # Top latent dimensions with highest AUC-ROC score for separating the groups top_n_genes = 15 # Number of top contributing genes per latent dimension latent_contributions = loaded_ontix.explain( split="test", method="IntegratedGradients", # "DeepLiftShap" or "IntegratedGradients" n_subset=400, # Randomly sample in input and baseline space for computational efficiency sel_latent_dim=selected_dim, # Or specify a single latent dimension to explain input_type="grouped", # We will test group 'malignant cell' (input) vs. 'pulmonary alveolar type 2 cell' (baseline) input_group=groups_list[0], # "malignant cell" baseline_type="random", baseline_group=groups_list[1], # "pulmonary alveolar type 2 cell" anno_col=anno_col, llm_explain=True, # Explain biological functions of the contributing genes using a LLM llm_client="openrouter", llm_model="google/gemma-4-31b-it", # llm_client="scads-llm", # llm_model="google/gemma-4-31b-it", # llm_model="alias-ha", top_n_genes=top_n_genes, )

The LLM-based verbal explaination and summary is stored as a markdown file:

In [ ]:

from pathlib import Path
from IPython.display import Markdown, display

md_path = Path("latent_explanations.md")
display(Markdown(md_path.read_text(encoding="utf-8")))

from pathlib import Path from IPython.display import Markdown, display md_path = Path("latent_explanations.md") display(Markdown(md_path.read_text(encoding="utf-8")))

Using the .generate() functionality for synthetic data generation¶

Leveraging our pretrained models we can sample from the latent space artificial cells.

We can do this either randomly across the whole latent space or from pre-defined points of the latent space.

We will do here the latter by calculating the mean of our malignant cells for each latent dimension and generate 1000 synthetic cells from this point of the latent space.

In [ ]:

df_latent_malignant = df_latent.loc[
	result.new_datasets.test.metadata.cell_type == 'malignant cell',
 :
]

latent_prior_means = df_latent_malignant.mean()

# Generate 1000 latent prior samples based on means and some noise
num_samples = 1000
noise = np.random.normal(0, 0.1, (num_samples, len(latent_prior_means)))
synthetic_latent = np.tile(latent_prior_means.values, (num_samples, 1)) + noise

df_latent_malignant = df_latent.loc[ result.new_datasets.test.metadata.cell_type == 'malignant cell', : ] latent_prior_means = df_latent_malignant.mean() # Generate 1000 latent prior samples based on means and some noise num_samples = 1000 noise = np.random.normal(0, 0.1, (num_samples, len(latent_prior_means))) synthetic_latent = np.tile(latent_prior_means.values, (num_samples, 1)) + noise

In [ ]:

# Generate reconstructions (gene expression) from the synthetic latent samples
generated_reconstructions = loaded_ontix.generate(latent_prior=synthetic_latent)
print("Generated reconstructions shape:", generated_reconstructions.shape)

# Generate reconstructions (gene expression) from the synthetic latent samples generated_reconstructions = loaded_ontix.generate(latent_prior=synthetic_latent) print("Generated reconstructions shape:", generated_reconstructions.shape)