cystainer 0.1.1

A PyTorch-based tool for predicting missing proteins in cytometry/single-cell data

CyStainer

CyStainer package for cytometry marker imputation

CyStainer is a PyTorch-based deep learning tool for predicting missing proteins and imputing marker expression in cytometry and single-cell data. It utilizes a combination of Variational Autoencoders (VAEs) and Transformer architectures to integrate multiple batches/panels and infer missing markers accurately.

📦 Installation

Since CyStainer is built on PyTorch, ensure you have an environment with Python 3.11+ and the appropriate PyTorch version for your hardware (CUDA recommended).

You can install CyStainer from PyPI:

pip install cystainer

or directly from the source

git clone https://github.com/sysgen-uef/cystainer_package.git
cd cystainer_package
pip install .

🧹 Data Preprocessing (From .fcs to .h5ad)

CyStainer expects your input data to be formatted as AnnData objects (either passed as a list in memory or saved locally as .h5ad files). Before feeding your cytometry data into the model, it must be properly preprocessed.

Mandatory Pre-processing:

• Cleaning: Ensure your data is pre-gated to remove doublets, debris, and dead cells.
• Compensation: Your .fcs files should be already compensated.

Transformation & Scaling (Recommended): You will typically need to transform your fluorescence intensities (e.g., using an arcsinh transformation) and optionally scale them. Removing saturated values (extreme highs or zeros) can also improve model performance depending on your dataset.

Here is a minimal example of how to process a standard .fcs file into a ready-to-use .h5ad file using FlowKit and AnnData:

import flowkit as fk
import pandas as pd
import numpy as np
import anndata as ad

# Load the compensated and cleaned .fcs file
sample = fk.Sample('path/to/cleaned_sample.fcs')
df = sample.as_dataframe('raw')
df.columns = sample.pnn_labels # Set marker names

# Transformation (e.g., arcsinh with a cofactor of 100 or 150)
cofactor = 100
non_scatter = ['FSC' not in c and 'SSC' not in c for c in df.columns]
df.loc[:, non_scatter] = np.arcsinh(df.loc[:, non_scatter] / cofactor)

# Optional: Remove saturation / extreme outliers
# Useful if your instrument records artificial bounds (e.g., exactly 0 or max value)
df = df[~np.any(((df <= 0) | (df >= df.max().max())), axis=1)]

# Optional: Scaling (Min-Max scaling to [0, 1] or Z-score normalization)
# Min-Max Scaling example:
df = (df - df.min()) / (df.max() - df.min())
# Z-score Scaling example (alternatively): 
# df = (df - df.mean()) / df.std()

# Convert to AnnData and save for CyStainer
adata = ad.AnnData(df)
adata.write('preprocessed_sample.h5ad', compression='gzip')

Once your .fcs files are converted into .h5ad objects, you can load them directly into your workflow.

🚀 Quick Start Guide

The primary way to interact with the package is through the CyStainer wrapper class. The standard workflow consists of initializing the model, loading training data, building the network, and running inference.

1. Training a Base Model

You can load your data either from a folder of .h5ad files or by passing a list of AnnData objects directly in memory.

from cystainer import CyStainer

# Initialize the stainer (automatically detects CUDA/CPU)
stainer = CyStainer()

# Load training data
# Alternatively, use: adata_list=[adata1, adata2]
# CyStainer includes a built-in utility to visualize how markers overlap across different panels or batches.
stainer.load_train_data(folder_path='./data_example/train', get_panel_vis=True)

# Build the model 
# You can pass custom hyperparameters here if needed
stainer.build_model()

# Train the model
stainer.train()

# By default the model is automatically saved in the same directory
# as cystainer.pt

2. Predicting Missing Markers

Once a model is trained, you can load inference data. The .load_predict_data() method ensures your cells are not shuffled so that the output matches your input order.

# Load prediction data using the base model's reference markers
stainer.load_predict_data(folder_path='./data_example/test')

# Run predictions and save directly to disk
stainer.predict(output_path='imputed_cells.h5ad')

# Alternatively, return the predictions as a pandas DataFrame:
# df_imputed = stainer.predict(return_pred=True)

3. Fine-Tuning on New Batches

If you receive new data from a different batch or panel, you don't need to retrain from scratch. CyStainer can freeze the main network and fine-tune only the batch embeddings to align the new data.

# Load the pre-trained model
stainer = CyStainer.load('cystainer.pt')

# Load the new data for fine-tuning
# Note, anndata objects must have batch info column
stainer.load_finetune_data(folder_path='./data_example/test', batch_info='batch')

# Fine-tune the batch embeddings
stainer.finetune()

# Predict on the newly fine-tuned data
stainer.predict(output_path='imputed_new_batch.h5ad')

4. Batch Correction

CyStainer allows you to translate cells from their original batch to a specific target batch distribution.

stainer.predict(
    output_path='batch_corrected_cells.h5ad',
    correct_batch=True,
    target_batch_name='single_batch' # Must match a batch name in stainer.batch_dict
)