Metadata-Version: 2.4
Name: cystainer
Version: 0.1.1
Summary: A PyTorch-based tool for predicting missing proteins in cytometry/single-cell data
Author: Konstantin Ivanov
Author-email: kivanov@uef.fi
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.11.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.7.1
Requires-Dist: numpy>=1.26.4
Requires-Dist: pandas>=2.2.3
Requires-Dist: anndata>=0.12.4
Requires-Dist: scipy>=1.14.1
Requires-Dist: scikit-learn>=1.5.1
Requires-Dist: tqdm>=4.66.5
Requires-Dist: matplotlib>=3.9.2
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: license-file
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# 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

```bash
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`:

```python
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.

```python
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)
```
![image](image/panel_alignment.png)

```python
# 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.

```python
# 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.

```python
# 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.

```python
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
)
```
