cystainer 0.1.0__py3-none-any.whl

cystainer/__init__.py

@@ -0,0 +1,16 @@
+# cystainer/__init__.py
+
+# Package version
+__version__ = "0.1.0"
+
+# Imports
+from .runner import CyStainer
+from .data import load_data_from_folder
+from .modules import CyStainerModel
+
+# Imports `from cystainer import *`
+__all__ = [
+    "CyStainerModel",
+    "CyStainer",
+    "load_data_from_folder",
+]

cystainer/data.py

@@ -0,0 +1,204 @@
+import os
+import itertools
+import torch
+import numpy as np
+import pandas as pd
+import anndata as ad
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from torch.utils.data import DataLoader
+from .modules import get_marker_list, pad_and_get_pad_mask
+
+def load_data_from_folder(folder_path=None, adata_list=None, exclude_files=[], max_n_cells=int(1e+6), batch_size=1024, is_train=True, drop_markers=None,
+                          reference_markers=None, reference_batch_dict=None, normalize=False, batch_info=None, get_panel_vis=False):
+    """Reads .h5ad files from a folder OR accepts a list of AnnData objects directly."""
+    
+    # Establish the adata_list
+    if adata_list is None:
+        if folder_path is None:
+            raise ValueError("You must provide either 'folder_path' or 'adata_list'.")
+            
+        sample_names = [f for f in os.listdir(folder_path) if f.endswith('.h5ad') and (f not in exclude_files)]
+        if not sample_names:
+            raise ValueError(f"No files found in {folder_path}")
+        
+        # Load adatas from disk
+        loaded_adatas = [ad.read_h5ad(os.path.join(folder_path, f)) for f in sample_names]
+    else:
+        # Use the provided list of AnnData objects
+        loaded_adatas = adata_list
+
+    # Extract DataFrames and Batch Info
+    df_list = [adata.to_df().iloc[:max_n_cells,:] for adata in loaded_adatas]
+    
+    if batch_info is not None:
+        batch_list = [adata.obs[batch_info].tolist() for adata in loaded_adatas]
+        batch_names = list(itertools.chain.from_iterable(batch_list))
+    else:
+        print("No batch info passed, setting 'single_batch' for the data.")
+        batch_names = list(itertools.chain.from_iterable([['single_batch'] * df.shape[0] for df in df_list]))
+
+    # Visualization, Processing, and Normalization
+    if get_panel_vis:
+        visualize_panel_overlap(df_list)
+        
+    if drop_markers is not None:
+        df_list = [df.loc[:,~df.columns.isin(drop_markers)] for df in df_list]
+    
+    if normalize:
+        df_list = [(df - df.mean()) / df.std() for df in df_list]
+
+    # Handle Batches
+    if reference_batch_dict is None:
+        # Sort for deterministic behavior
+        batch_dict = {batch: i for i, batch in enumerate(sorted(list(set(batch_names))))}
+    else:
+        batch_dict = reference_batch_dict.copy()
+        current_max_idx = max(batch_dict.values()) if batch_dict else -1
+        
+        # Find only the new batches and sort them
+        new_batches = sorted(list(set(batch_names) - set(batch_dict.keys())))
+        
+        for batch in new_batches:
+            current_max_idx += 1
+            batch_dict[batch] = current_max_idx
+
+    batch_ids = pd.Series(batch_names).apply(lambda x: batch_dict[x]).values
+
+    # Handle Markers and Padding
+    if reference_markers is None:
+        marker_list, shared_markers, unique_markers = get_marker_list(df_list)
+    else:
+        marker_list = reference_markers['marker_list']
+        shared_markers = reference_markers['shared_markers']
+        unique_markers = reference_markers['unique_markers']
+
+    df_list, pad_mask = pad_and_get_pad_mask(df_list, marker_list)
+    
+    # PyTorch Dataloader
+    data = [{'x': x.astype('float32'), 'batch': b, 'pad': p} 
+            for x, b, p in zip(pd.concat(df_list).values, batch_ids, np.concatenate(pad_mask, axis=0))]
+    
+    dataloader = DataLoader(data, batch_size=batch_size, shuffle=is_train)
+
+    markers_info = {
+        'marker_list': marker_list,
+        'shared_markers': shared_markers,
+        'unique_markers': unique_markers
+    }
+
+    return dataloader, markers_info, batch_dict
+
+def visualize_panel_overlap(df_list, alignment='left'):
+    """
+    Creates a block alignment plot showing shared and unique markers 
+    across different panels.
+    
+    Args:
+        df_list (list): List of pandas DataFrames.
+        panel_names (list): List of panel names as strings.
+        alignment (str): 'center' for a center-outwards pyramid style, 
+                         'left' for a cascading left-aligned style.
+    """
+    # Extract markers and build presence matrix
+    panel_markers = set([frozenset(df.columns) for df in df_list])
+    all_markers = set().union(*panel_markers)
+    num_panels = len(panel_markers)
+    panel_names = [f'Panel {i+1}' for i in range(num_panels)]
+    
+    matrix = pd.DataFrame(index=list(all_markers), columns=panel_names)
+    for name, markers in zip(panel_names, panel_markers):
+        matrix[name] = matrix.index.isin(markers).astype(int)
+        
+    # Sort markers based on requested alignment
+    if alignment == 'left':
+        # Simple cascade sort
+        matrix = matrix.sort_values(by=panel_names, ascending=False)
+        sorted_markers = matrix.index.tolist()
+        
+    elif alignment == 'center':
+        # Center-outwards sort (Frequency and Center of Mass)
+        matrix['freq'] = matrix.sum(axis=1)
+        col_indices = np.arange(num_panels)
+        
+        matrix['com'] = (matrix[panel_names] * col_indices).sum(axis=1) / matrix['freq']
+        matrix['pattern'] = matrix[panel_names].astype(str).agg(''.join, axis=1)
+        
+        grouped = matrix.groupby('pattern')
+        group_stats = grouped.first()[['freq', 'com']].sort_values(by=['freq', 'com'], ascending=[False, True])
+        
+        left_part = []
+        right_part = []
+        center_part = []
+        mid_point = (num_panels - 1) / 2.0
+        
+        for pattern, row in group_stats.iterrows():
+            group_markers = matrix[matrix['pattern'] == pattern].index.tolist()
+            group_markers.sort() # Alphabetical fallback
+            
+            if row['freq'] == num_panels:
+                center_part.extend(group_markers)
+            else:
+                if row['com'] < mid_point:
+                    left_part = group_markers + left_part
+                elif row['com'] > mid_point:
+                    right_part = right_part + group_markers
+                else:
+                    if len(left_part) <= len(right_part):
+                        left_part = group_markers + left_part
+                    else:
+                        right_part = right_part + group_markers
+                        
+        sorted_markers = left_part + center_part + right_part
+        
+    else:
+        raise ValueError("The 'alignment' parameter must be either 'left' or 'center'.")
+
+    # Setup the plot
+    fig, ax = plt.subplots(figsize=(14, len(panel_names) * 1.5))
+    rect_height = 0.6
+    colors = ['#72A0C1', '#D65A61', '#E8B358', '#7CE0C9', '#8291A8', '#C48CB3'] 
+    
+    # Draw the blocks
+    for i, (name, markers) in enumerate(zip(panel_names, panel_markers)):
+        presence = [1 if m in markers else 0 for m in sorted_markers]
+        start_idx = None
+        for j, val in enumerate(presence):
+            if val == 1 and start_idx is None:
+                start_idx = j
+            elif val == 0 and start_idx is not None:
+                rect = patches.Rectangle(
+                    (start_idx, i - rect_height/2), j - start_idx, rect_height,
+                    edgecolor='black', facecolor=colors[i % len(colors)], alpha=0.9, linewidth=0.8
+                )
+                ax.add_patch(rect)
+                start_idx = None
+                
+        # Catch end blocks
+        if start_idx is not None:
+            rect = patches.Rectangle(
+                (start_idx, i - rect_height/2), len(presence) - start_idx, rect_height,
+                edgecolor='black', facecolor=colors[i % len(colors)], alpha=0.9, linewidth=0.8
+            )
+            ax.add_patch(rect)
+            
+        ax.text(-0.8, i, name, va='center', ha='right', fontsize=12, fontweight='bold')
+
+    # Formatting
+    ax.set_xlim(0, len(sorted_markers))
+    ax.set_ylim(-1, len(panel_names))
+    ax.set_yticks([])
+    ax.set_xticks(np.arange(len(sorted_markers)) + 0.5)
+    ax.set_xticklabels(sorted_markers, rotation=90, ha='center', fontsize=10)
+    
+    # Remove borders
+    for spine in ['top', 'right', 'left']:
+        ax.spines[spine].set_visible(False)
+    
+    ax.xaxis.grid(True, linestyle='--', alpha=0.4)
+    ax.set_axisbelow(True)
+    
+    title_align = "Center" if alignment == 'center' else "Left"
+    plt.title(f"Panel Marker Alignment", fontsize=16, pad=25, fontweight='bold')
+    plt.tight_layout()
+    plt.show()

cystainer/modules.py

@@ -0,0 +1,277 @@
+# modules
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Normal
+from torch.distributions import kl_divergence as kl
+
+
+class CyStainerModel(nn.Module):
+    def __init__(self, marker_info, batch_dict, hidden_size=[512, 256, 128], feature_hidden_size=[128, 64, 32],
+                 num_layers=3, num_heads=8, value_emb_dim=8, marker_emb_dim=8,
+                 batch_emb_dim=8, latent_size=16, dropout=0.1):
+        super(CyStainerModel, self).__init__()
+        self.latent_size = latent_size
+        self.batch_dict = batch_dict
+        self.num_batches = len(batch_dict)
+        self.marker_info = marker_info
+        self.n_marker = len(marker_info['marker_list'])
+        self.value_emb_dim = value_emb_dim
+        self.marker_emb_dim = marker_emb_dim
+        self.batch_emb_dim = batch_emb_dim
+
+        self.marker_embedding = nn.Embedding(self.n_marker, marker_emb_dim)
+        self.batch_embedding = nn.Embedding(self.num_batches, batch_emb_dim)
+        self.value_projector = nn.Linear(1, value_emb_dim)
+        self.value_layer_norm = nn.LayerNorm(value_emb_dim)
+
+        transformer_input_dim = value_emb_dim + marker_emb_dim + batch_emb_dim
+
+        transformer_encoder_layer = nn.TransformerEncoderLayer(
+            d_model=transformer_input_dim,
+            nhead=num_heads,
+            dim_feedforward=hidden_size[0],
+            dropout=dropout,
+            activation='gelu',
+            batch_first=True
+        )
+        self.transformer_encoder = nn.TransformerEncoder(transformer_encoder_layer,
+                                                         num_layers=num_layers,
+                                                         enable_nested_tensor=False)
+        self.transformer_layer_norm = nn.LayerNorm(transformer_input_dim)
+
+        self.feature_encoder = nn.Sequential(
+            nn.Linear(transformer_input_dim, feature_hidden_size[0]),
+            nn.Dropout(dropout),
+            nn.SiLU(),
+            nn.Linear(feature_hidden_size[0], feature_hidden_size[1]),
+            nn.Dropout(dropout),
+            nn.SiLU(),
+            nn.Linear(feature_hidden_size[1], feature_hidden_size[2]))
+
+        self.latent_projector = nn.Sequential(
+            nn.Linear(feature_hidden_size[2] * self.n_marker, hidden_size[0]),
+            nn.Dropout(dropout),
+            nn.SiLU(),
+            nn.Linear(hidden_size[0], hidden_size[1]),
+            nn.Dropout(dropout),
+            nn.SiLU())
+
+        self.mu_encoder = nn.Linear(hidden_size[1], latent_size)
+        self.logvar_encoder = nn.Linear(hidden_size[1], latent_size)
+
+        decoder_input_dim = latent_size + batch_emb_dim
+
+        self.fc_decoder = nn.Sequential(
+            nn.Linear(decoder_input_dim, self.n_marker),
+            nn.Dropout(dropout),
+            nn.SiLU(),
+            nn.Linear(self.n_marker, self.n_marker),
+            nn.Dropout(dropout),
+            nn.SiLU(),
+            nn.Linear(self.n_marker, self.n_marker))
+
+    def encode(self, x, batch_idx, attention_mask=None):
+        batch_size, seq_len = x.shape
+        device = x.device
+
+        marker_indices = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)
+        marker_emb = self.marker_embedding(marker_indices)
+
+        batch_emb_flat = self.batch_embedding(batch_idx)
+        batch_emb = batch_emb_flat.unsqueeze(1).expand(-1, seq_len, -1)
+
+        value_emb = self.value_projector(x.unsqueeze(-1))
+        value_emb = self.value_layer_norm(value_emb)
+
+        combined_emb = torch.cat([value_emb, marker_emb, batch_emb], dim=-1)
+
+        transformer_output = self.transformer_encoder(
+            combined_emb, src_key_padding_mask=attention_mask
+        )
+        normalized_output = self.transformer_layer_norm(transformer_output)
+
+        encoded_features = self.feature_encoder(normalized_output).flatten(1,-1)
+        latent_projection = self.latent_projector(encoded_features)
+        mu = self.mu_encoder(latent_projection)
+        logvar = self.logvar_encoder(latent_projection)
+        var = torch.exp(logvar) + 1e-4
+        dist = Normal(mu, var.sqrt())
+        latent_sample = dist.rsample()
+        return mu, var, latent_sample
+
+    def decode(self, z, batch_idx):
+        batch_emb = self.batch_embedding(batch_idx)
+        decode_input = torch.cat([z, batch_emb], dim=-1)
+        x = self.fc_decoder(decode_input)
+        return x
+
+    def forward(self, x, batch_idx, attention_mask=None):
+        mu, var, latent_sample = self.encode(x, batch_idx, attention_mask)
+        z = mu if not self.training else latent_sample
+        x = self.decode(z, batch_idx)
+        return x, mu, var, latent_sample
+
+    def add_new_batches(self, n_new_batches):
+        """
+        Expands the embedding layer to accommodate new batches.
+        Useful for fine-tuning on new datasets without retraining entirely.
+        """
+        old_embedding = self.batch_embedding
+        old_num_embeddings = old_embedding.num_embeddings
+        new_num_embeddings = old_num_embeddings + n_new_batches
+
+        # Create new embedding layer
+        device = old_embedding.weight.device
+        new_embedding = nn.Embedding(new_num_embeddings, old_embedding.embedding_dim, device=device)
+
+        # Initialize weights (copy old ones, randomize new ones)
+        with torch.no_grad():
+            new_embedding.weight[:old_num_embeddings].copy_(old_embedding.weight)
+
+        self.batch_embedding = new_embedding
+        return new_embedding
+
+    @torch.no_grad()
+    def correct_batch(self, x, source_batch_idx, target_batch_idx, attention_mask=None):
+        """
+        Translates cells from their original batch to a reference batch distribution.
+        """
+        self.eval()
+        mu, _, _ = self.encode(x, source_batch_idx, attention_mask)
+        corrected_x = self.decode(mu, target_batch_idx)
+        return corrected_x
+
+
+def get_marker_list(df_list):
+    all_features = []
+    for df in df_list:
+        all_features.extend(df.columns.tolist())
+    marker_list = sorted(list(pd.Series(all_features).unique()))
+    is_overlap = [
+        all(marker in df.columns for df in df_list) for marker in marker_list
+    ]
+    shared_markers = [marker for marker, is_o in zip(marker_list, is_overlap) if is_o]
+    unique_markers = [marker for marker, is_o in zip(marker_list, is_overlap) if not is_o]
+
+    return marker_list, shared_markers, unique_markers
+
+
+def pad_and_get_pad_mask(df_list, marker_list):
+    df_list_copy = df_list.copy()
+    pad_mask = [
+        np.tile(
+            np.array([0.0 if marker not in df.columns else 1.0 for marker in marker_list], dtype=np.float32), 
+            (len(df), 1)
+        ) 
+        for df in df_list_copy
+    ]
+    
+    for i, df in enumerate(df_list_copy):
+        df = df.reindex(columns=marker_list, fill_value=0.0)
+        df_list_copy[i] = df
+
+    return df_list_copy, pad_mask
+
+
+def mask_batch(batch,
+               marker_list=None,
+               max_mask_rate_overlap=0.3,
+               separated_mask_rate=True,
+               max_mask_rate_unique=1.0,
+               shared_markers=None,
+               unique_markers=None):
+
+    batch_x_clone = batch['x'].clone()
+    pad_mask = batch['pad'].bool()
+
+    batch_size, seq_len = batch_x_clone.shape
+    device = batch_x_clone.device
+
+    def get_mask_for_top_k(scores, is_valid_token_mask, n_to_mask):
+        """Helper to select top-k random scores for masking."""
+        scores[~is_valid_token_mask] = float('inf')
+        sorted_indices = torch.argsort(scores, dim=1)
+        ranks = torch.empty_like(sorted_indices)
+        row_ranks = torch.arange(seq_len, device=device).expand(batch_size, -1)
+        ranks.scatter_(1, sorted_indices, row_ranks)
+        return ranks < n_to_mask.unsqueeze(1)
+
+    rand_scores = torch.rand(batch_size, seq_len, device=device)
+    final_mask = torch.zeros(batch_size, seq_len, dtype=torch.bool, device=device)
+
+    if separated_mask_rate:
+        if marker_list is None or shared_markers is None or unique_markers is None:
+            raise ValueError(
+                "marker_list, shared_markers, and unique_markers must be provided "
+                "when separated_mask_rate is True."
+            )
+
+        # Convert to sets for faster lookup
+        shared_set = set(shared_markers)
+        unique_set = set(unique_markers)
+
+        # Create 1D boolean masks for the columns
+        is_shared_col = torch.tensor([m in shared_set for m in marker_list], device=device)
+        is_unique_col = torch.tensor([m in unique_set for m in marker_list], device=device)
+
+        # A marker is valid to mask if it matches the group AND is present in the sample
+        is_overlap = is_shared_col.unsqueeze(0) & pad_mask
+        is_unique = is_unique_col.unsqueeze(0) & pad_mask
+
+        # Calculate number of markers to mask for overlap and unique sets
+        max_n_overlap = is_overlap.sum(dim=1) * max_mask_rate_overlap
+        max_n_unique = is_unique.sum(dim=1) * max_mask_rate_unique
+
+        n_to_mask_overlap = torch.floor(torch.rand(batch_size, device=device) * (max_n_overlap.float() + 1))
+        n_to_mask_unique = torch.floor(torch.rand(batch_size, device=device) * (max_n_unique.float() + 1))
+
+        # Get masks and combine them
+        mask_overlap = get_mask_for_top_k(rand_scores.clone(), is_overlap, n_to_mask_overlap)
+        mask_unique = get_mask_for_top_k(rand_scores.clone(), is_unique, n_to_mask_unique)
+        final_mask = mask_overlap | mask_unique
+
+    else: # Standard masking mode
+        max_n_to_mask = pad_mask.sum(dim=1) * max_mask_rate_overlap
+        n_to_mask = torch.floor(torch.rand(batch_size, device=device) * (max_n_to_mask.float() + 1))
+        final_mask = get_mask_for_top_k(rand_scores, pad_mask, n_to_mask)
+
+    # Apply the final mask: Set masked expression values to 0.0
+    batch_x_clone[final_mask] = 0.0
+
+    return {
+        'x': batch_x_clone,
+        'batch': batch['batch'].clone(),
+        'pad': batch['pad'].clone(),
+        'mask': final_mask | ~batch['pad'].bool()
+    }
+
+
+def loss_function(recon_x, x, pad, mu, var, kl_weight):
+    """
+    Calculates the VAE loss, combining reconstruction and KL divergence.
+
+    Args:
+        recon_x (torch.Tensor): The reconstructed data from the model.
+        x (torch.Tensor): The original input data.
+        pad_mask (torch.Tensor): A boolean mask where True indicates non-padded
+                                 values that should be included in the loss.
+        mu (torch.Tensor): The mean of the latent distribution.
+        var (torch.Tensor): The variance of the latent distribution.
+        kl_weight (float): A weight factor for the KL divergence term, often
+                           used for annealing.
+
+    Returns:
+        torch.Tensor: The total calculated loss.
+    """
+    # Reconstruction loss (Huber loss) on non-padded values
+    recon_loss = F.huber_loss(recon_x[pad], x[pad])
+
+    # KL divergence between the learned distribution and a standard normal
+    prior = Normal(torch.zeros_like(mu), torch.ones_like(var))
+    kl_div = kl(Normal(mu, var.sqrt()), prior).sum(dim=1).mean()
+    kl_loss = kl_div * kl_weight
+
+    return recon_loss + kl_loss

cystainer/runner.py

@@ -0,0 +1,318 @@
+import os
+import torch
+from torch.utils.data import RandomSampler
+import pandas as pd
+import anndata as ad
+import numpy as np
+from tqdm import tqdm
+
+from .modules import CyStainerModel, mask_batch, loss_function
+from .data import load_data_from_folder 
+
+def get_default_device():
+    """Helper function to automatically detect CUDA, MPS, or CPU."""
+    if torch.cuda.is_available():
+        return torch.device('cuda')
+    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        return torch.device('mps')
+    return torch.device('cpu')
+
+class CyStainer:
+    def __init__(self, marker_info=None, batch_dict=None, model=None, device=None, **model_kwargs):
+        """
+        Initializes the CyStainer wrapper. Can be initialized empty and populated later
+        via .load_train_data() and .build_model().
+        """
+        self.device = device or get_default_device()
+        
+        self.marker_info = marker_info
+        self.batch_dict = batch_dict
+        self.train_dataloader = None
+        self.finetune_dataloader = None
+        self.predict_dataloader = None
+        self.model = None
+
+        if model is not None:
+            self.model = model.to(self.device)
+            self.marker_info = model.marker_info
+            self.batch_dict = model.batch_dict
+        elif marker_info is not None and batch_dict is not None:
+            self.build_model(**model_kwargs)
+
+    def __getstate__(self):
+        """
+        Dictates what gets saved when torch.save(stainer) is called.
+        We create a copy of the object's dictionary and remove the dataloader.
+        """
+        state = self.__dict__.copy()
+        state['train_dataloader'] = None
+        state['finetune_dataloader'] = None 
+        state['predict_dataloader'] = None
+        return state
+
+    def load_train_data(self, folder_path=None, adata_list=None, **kwargs):
+        """
+        Convenience method to load data from a folder or an AnnData list 
+        and populate the class state.
+        """
+        source_msg = folder_path if folder_path else "provided AnnData list"
+        print(f"Loading data from {source_msg}...")
+        
+        self.train_dataloader, self.marker_info, self.batch_dict = load_data_from_folder(
+            folder_path=folder_path,
+            adata_list=adata_list,
+            **kwargs
+        )
+        print(f"Data loaded. Found {len(self.marker_info['marker_list'])} markers and {len(self.batch_dict)} batches.")
+        return self
+    
+    def load_finetune_data(self, folder_path=None, adata_list=None, **kwargs):
+        """
+        Loads a new dataset for fine-tuning, using the existing model's 
+        markers and batches as a reference to align the new data.
+        """
+        if self.marker_info is None or self.batch_dict is None:
+            raise RuntimeError("Base model state not found. Train or load a base model first before loading fine-tuning data.")
+
+        source_msg = folder_path if folder_path else "provided AnnData list"
+        print(f"Loading fine-tuning data from {source_msg}...")
+        
+        self.finetune_dataloader, _, self.ft_batch_dict = load_data_from_folder(
+            folder_path=folder_path,
+            adata_list=adata_list,
+            reference_markers=self.marker_info,
+            reference_batch_dict=self.batch_dict,
+            **kwargs
+        )
+        
+        ft_batches_count = len(self.ft_batch_dict) - len(self.batch_dict)
+        print(f"Fine-tuning data loaded. Found {ft_batches_count} new batches to integrate.")
+        return self
+
+    def load_predict_data(self, folder_path=None, adata_list=None, **kwargs):
+        """
+        Loads data specifically for inference. Forces is_train=False to 
+        ensure the predicted cells maintain their original order.
+        """
+        if self.marker_info is None or self.batch_dict is None:
+            raise RuntimeError("Model state missing. Train or load a model first.")
+
+        source_msg = folder_path if folder_path else "provided AnnData list"
+        print(f"Loading prediction data from {source_msg}...")
+        
+        # Force is_train to False to prevent shuffling!
+        kwargs['is_train'] = False 
+        
+        self.predict_dataloader, _, _ = load_data_from_folder(
+            folder_path=folder_path,
+            adata_list=adata_list,
+            reference_markers=self.marker_info,
+            reference_batch_dict=self.batch_dict,
+            **kwargs
+        )
+        print("Prediction data loaded successfully.")
+        return self
+
+    def build_model(self, **model_kwargs):
+        """
+        Initializes the PyTorch model using the state extracted from load_train_data.
+        """
+        if self.marker_info is None or self.batch_dict is None:
+            raise RuntimeError("Cannot build model: marker_info and batch_dict are missing. Run .load_train_data() first.")
+        
+        self.model = CyStainerModel(
+            marker_info=self.marker_info,
+            batch_dict=self.batch_dict,
+            **model_kwargs
+        ).to(self.device)
+        print("Model initialized successfully.")
+        return self
+
+    def train(self, dataloader=None, max_epochs=300, lr=0.0001, kl_weight=0.0001, tol=1e-4, patience=3,
+              model_name='cystainer.pt', loss_name='loss.csv', metrics_path='./metrics_and_loss', separated_mask_rate=True):
+        """Trains the model from scratch with early stopping."""
+        if self.model is None:
+            raise RuntimeError("Model is not built. Call .build_model() before training.")
+            
+        train_loader = dataloader or self.train_dataloader
+        if train_loader is None:
+             raise ValueError("No dataloader found. Please provide one or run .load_train_data() first.")
+
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
+        self._run_loop(train_loader, optimizer, max_epochs, kl_weight, tol, patience, model_name, loss_name, metrics_path, separated_mask_rate)
+
+    def finetune(self, ft_loader=None, ft_batch_dict=None, max_epochs=50, lr=0.001, kl_weight=0.0001, tol=1e-4, patience=3,
+                 model_name='cystainer_finetuned.pt', loss_name='loss_finetuned.csv', metrics_path='./metrics_and_loss', separated_mask_rate=True):
+        """Fine-tunes only the batch embeddings for new data with early stopping."""
+        if self.model is None:
+            raise RuntimeError("Model is not trained or loaded.")
+        
+        # Fallback to internal state if arguments aren't provided
+        ft_loader = ft_loader or self.finetune_dataloader
+        ft_batch_dict = ft_batch_dict or self.ft_batch_dict
+
+        if ft_loader is None or ft_batch_dict is None:
+            raise ValueError("Missing fine-tuning data. Run .load_finetune_data() first or pass the dataloader and new_batch_dict explicitly.")
+
+        # Calculate how many new batches were added
+        ft_batches_count = len(ft_batch_dict) - len(self.batch_dict)
+        if ft_batches_count > 0:
+            self.model.add_new_batches(ft_batches_count)
+            self.batch_dict = ft_batch_dict
+
+        # Freeze all parameters except batch_embedding
+        for param in self.model.parameters():
+            param.requires_grad = False
+        self.model.batch_embedding.weight.requires_grad = True
+
+        # Calculate the index where new batches start
+        old_num_embeddings = len(self.batch_dict) - ft_batches_count
+
+        # Register a backward hook to freeze old batch embeddings
+        def freeze_old_embeddings(grad):
+            grad_clone = grad.clone()
+            grad_clone[:old_num_embeddings] = 0.0
+            return grad_clone
+        
+        self.model.batch_embedding.weight.register_hook(freeze_old_embeddings)
+
+        optimizer = torch.optim.AdamW(self.model.batch_embedding.parameters(), lr=lr)
+        self._run_loop(ft_loader, optimizer, max_epochs, kl_weight, tol, patience, model_name, loss_name, metrics_path, separated_mask_rate)
+
+    def predict(self, dataloader=None, output_path='adata_pred.h5ad', keep_markers=None, correct_batch=False, target_batch_name=None, return_pred=False):
+        """Runs inference and saves the predicted AnnData object."""
+        if self.model is None:
+            raise RuntimeError("Model is not built or loaded.")
+
+        # Fallback logic: User provided -> Predict loader -> Finetune loader -> Train loader
+        pred_loader = dataloader or self.predict_dataloader or self.finetune_dataloader or self.train_dataloader
+        
+        if pred_loader is None:
+            raise ValueError("No dataloader found. Pass one explicitly or run .load_predict_data() first.")
+
+        # Block predictions on shuffled dataloaders
+        if isinstance(pred_loader.sampler, RandomSampler):
+            raise RuntimeError(
+                "CRITICAL: You are attempting to predict on a shuffled DataLoader! "
+                "The output rows will NOT correspond to the original order of your cells in the AnnData object. "
+                "You must load the data using .load_predict_data() before running .predict()."
+            )
+
+        target_batch_int = 0
+        if correct_batch and target_batch_name is not None:
+            if target_batch_name not in self.batch_dict:
+                raise ValueError(f"Batch name '{target_batch_name}' not found. Available batches: {list(self.batch_dict.keys())}")
+            target_batch_int = self.batch_dict[target_batch_name]
+
+        self.model.eval()
+        result = []
+        
+        if correct_batch:
+            with torch.no_grad():
+                for batch in tqdm(pred_loader, desc="Predicting (Batch Corrected)"):
+                    batch = {key: value.to(self.device) for key, value in batch.items()}
+                    target_batch = torch.full_like(batch['batch'], fill_value=target_batch_int)
+                    x_pred = self.model.correct_batch(batch['x'], batch['batch'], target_batch, ~batch['pad'].bool())
+                    result.append(x_pred.cpu().numpy())
+        else:
+            with torch.no_grad():
+                for batch in tqdm(pred_loader, desc="Predicting"):
+                    batch = {key: value.to(self.device) for key, value in batch.items()}
+                    x_pred, _, _, _ = self.model.forward(batch['x'], batch['batch'], ~batch['pad'].bool())
+                    result.append(x_pred.cpu().numpy())
+
+        df_pred = pd.DataFrame(np.concatenate(result), columns=self.marker_info['marker_list'])
+        if keep_markers is not None:
+            df_pred = df_pred[keep_markers]
+        if return_pred:
+            return df_pred
+        else:
+            ad.AnnData(df_pred).write(output_path, compression='gzip')
+            print(f"Predictions saved to {output_path}")
+
+    def save(self, losses=None, model_name=None, loss_name=None, metrics_path=None):
+        """Saves the training losses and the entire CyStainer object."""
+        # Save losses if provided
+        if losses is not None:
+            if not os.path.exists(metrics_path):
+                os.makedirs(metrics_path, exist_ok=True)
+            df_loss = pd.DataFrame(losses, columns=['loss'])
+            df_loss.to_csv(f'{metrics_path}/{loss_name}', index=False)
+        
+        torch.save(self, model_name)
+
+    def _run_loop(self, dataloader, optimizer, max_epochs, kl_weight, tol, patience, model_name, loss_name, metrics_path, separated_mask_rate):
+        """Internal training loop with early stopping mechanism."""
+        self.model.train()
+        pbar = tqdm(range(max_epochs))
+        
+        prev_loss = float('inf')
+        patience_counter = 0
+        losses = []
+
+        for epoch in pbar:
+            epoch_losses = []
+            
+            for batch in dataloader:
+                batch = {key: value.to(self.device) for key, value in batch.items()}
+                batch_masked = mask_batch(
+                    batch, 
+                    self.marker_info['marker_list'],
+                    separated_mask_rate=separated_mask_rate,
+                    shared_markers=self.marker_info['shared_markers'],
+                    unique_markers=self.marker_info['unique_markers']
+                )
+                
+                x_pred, mu, var, _ = self.model.forward(batch_masked['x'], batch_masked['batch'], batch_masked['mask'].bool())
+                loss = loss_function(x_pred, batch['x'], batch['pad'].bool(), mu, var, kl_weight)
+                
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                
+                epoch_losses.append(loss.item())
+                
+            # Early Stopping Logic
+            avg_loss = np.mean(epoch_losses)
+            losses.append(avg_loss)
+            pbar.set_description(f'Loss: {avg_loss:.4f}')
+
+            self.save(losses=losses, model_name=model_name, loss_name=loss_name, metrics_path=metrics_path)
+            
+            loss_diff = abs(prev_loss - avg_loss)
+            
+            if loss_diff < tol:
+                patience_counter += 1
+            else:
+                patience_counter = 0
+                
+            if patience_counter >= patience:
+                # We use tqdm.write so it doesn't break the progress bar visual
+                tqdm.write(f"\nEarly stopping triggered at epoch {epoch + 1}. Loss change ({loss_diff:.6f}) is below tolerance ({tol}).")
+                break
+                
+            prev_loss = avg_loss
+
+    @classmethod
+    def load(cls, filepath, device=None):
+        """
+        Safely loads a saved CyStainer object, automatically mapping all tensors
+        and internal states to the requested device (CPU, GPU, or MPS).
+        """
+        # Determine the target device
+        target_device = device or get_default_device()
+        
+        print(f"Loading model from {filepath} onto {target_device}...")
+        
+        # 1. Load the object, intercepting tensors and forcing them to the target device
+        stainer = torch.load(filepath, map_location=target_device, weights_only=False)
+        
+        # 2. Update the internal device attribute so future dataloaders route correctly
+        stainer.device = target_device
+        
+        # 3. Explicitly move the PyTorch model to the new device just to be perfectly safe
+        if stainer.model is not None:
+            stainer.model = stainer.model.to(target_device)
+            
+        print("Model loaded successfully.")
+        return stainer

cystainer-0.1.0.dist-info/METADATA

@@ -0,0 +1,170 @@
+Metadata-Version: 2.4
+Name: cystainer
+Version: 0.1.0
+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 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
+)
+```

cystainer-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+cystainer/__init__.py,sha256=Yv_K9NHhvtxenToMvf3FadTmDS1YSmBwkbjBnkwewJM,301
+cystainer/data.py,sha256=xS5-XX1NguRH4Mxa4otMjwzE00HlhJzwooz43aPD1gw,8321
+cystainer/modules.py,sha256=rKZMW_-7yYJXSEJlSjTVtsY4cyFqT8_qwoXVHVfWrco,11195
+cystainer/runner.py,sha256=CaEsn50m-v3YnyTSgBvHacEG7DbUA5ScsKIXrMkJv7Q,14413
+cystainer-0.1.0.dist-info/licenses/LICENSE,sha256=0qOacCooGqkP7CL7qDsSOyMSmGtEb1s_Q1oSoZSrLxs,1067
+cystainer-0.1.0.dist-info/METADATA,sha256=sSH7f70qUwymDJXee7Bq60cZFETGyfSRE2L2rCPmy3Q,6211
+cystainer-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+cystainer-0.1.0.dist-info/top_level.txt,sha256=8kyy2H_PFxOOr53It06GAWAyUNDGd3Tw1ZThnhaSPcg,10
+cystainer-0.1.0.dist-info/RECORD,,

cystainer-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

cystainer-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Sysgen lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cystainer-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+cystainer