gsMap3D 0.1.0a1__py3-none-any.whl

gsMap/find_latent/st_process.py

@@ -0,0 +1,781 @@
+import logging
+import re
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import scanpy as sc
+import scipy.sparse as sp
+import torch
+from rich.progress import BarColumn, Progress, TaskProgressColumn, TimeRemainingColumn
+from torch.utils.data import DataLoader, TensorDataset
+
+from gsMap.config import FindLatentRepresentationsConfig
+
+from .gnn.gcn import GCN, build_spatial_graph
+
+logger = logging.getLogger(__name__)
+
+
+def convert_to_human_genes(adata, gene_homolog_dict: dict, species: str = None):
+    """
+    Convert gene names in adata to human gene symbols using a pre-computed mapping dictionary.
+
+    Args:
+        adata: AnnData object to convert
+        gene_homolog_dict: Dictionary mapping current gene names to human gene symbols
+        species: Optional species name to use for the original gene column in var
+
+    Returns:
+        adata: Processed AnnData object with human symbols as var_names
+    """
+    # Identify common genes present in mapping
+    common_genes = [g for g in gene_homolog_dict.keys() if g in adata.var_names]
+
+    if len(common_genes) < 300:
+        logger.warning(f"Only {len(common_genes)} genes found in mapping dictionary.")
+
+    # Filter to these genes
+    adata = adata[:, common_genes].copy()
+
+    human_symbols = [gene_homolog_dict[g] for g in adata.var_names]
+
+    # Store original gene names if species is provided
+    if species:
+        species_col_name = f"{species}_GENE_SYM"
+        adata.var[species_col_name] = adata.var_names.values
+
+    # Update var_names to human symbols
+    adata.var_names = human_symbols
+    adata.var.index.name = "HUMAN_GENE_SYM"
+
+    # Remove duplicated genes (keep first occurrence)
+    if adata.var_names.duplicated().any():
+        n_before = adata.n_vars
+        adata = adata[:, ~adata.var_names.duplicated()].copy()
+        logger.info(f"Removed {n_before - adata.n_vars} duplicate human gene symbols.")
+
+    return adata
+
+def find_common_hvg(sample_h5ad_dict, params: FindLatentRepresentationsConfig):
+    """
+    Identifies common highly variable genes (HVGs) across multiple ST datasets and calculates
+    the number of cells to sample from each dataset.
+
+    Args:
+        sample_h5ad_dict (dict): Dictionary mapping sample names to file paths of ST datasets.
+        params (object): Parameter object containing attributes.
+    """
+
+    variances_list = []
+    cell_number = []
+
+    logger.info("Finding highly variable genes (HVGs)...")
+
+    with Progress(
+        BarColumn(),
+        TaskProgressColumn(),
+        TimeRemainingColumn(),
+    ) as progress:
+        task = progress.add_task("Finding HVGs", total=len(sample_h5ad_dict))
+
+        for sample_name, st_file in sample_h5ad_dict.items():
+            adata_temp = sc.read_h5ad(st_file)
+            # sc.pp.filter_genes(adata_temp, min_counts=1)
+
+            # Filter out mitochondrial and hemoglobin genes
+            gene_keep = ~adata_temp.var_names.str.match(re.compile(r'^(HB.-|MT-)', re.IGNORECASE))
+            if removed_genes := adata_temp.n_vars - gene_keep.sum():
+                progress.console.log(f"Removed {removed_genes} mitochondrial and hemoglobin genes in {sample_name}.")
+
+            adata_temp = adata_temp[:,gene_keep].copy()
+
+            # Make gene names unique to avoid issues with HVG calculation
+            adata_temp.var_names_make_unique()
+
+            is_count_data, actual_data_layer = setup_data_layer(adata_temp, params.data_layer, verbose=False)
+            cell_number.append(adata_temp.n_obs)
+
+            try:
+                # Identify highly variable genes
+                if is_count_data:
+                    sc.experimental.pp.highly_variable_genes(
+                        adata_temp, n_top_genes=params.feat_cell, subset=False
+                    )
+                else:
+                    sc.pp.highly_variable_genes(
+                        adata_temp, n_top_genes=params.feat_cell, subset=False, flavor='seurat'
+                    )
+
+                var_df = adata_temp.var.copy()
+                var_df["gene"] = var_df.index.tolist()
+                variances_list.append(var_df)
+            except Exception as e:
+                logger.warning(f"[HVG skipped] {e}")
+                continue
+
+
+            progress.update(task, advance=1)
+
+    # Check if we have any valid variance results
+    if len(variances_list) == 0:
+        raise ValueError("No valid HVG results obtained from any sample. Please check your data and parameters.")
+
+    # Find the common genes across all samples
+    common_genes = np.array(
+        list(set.intersection(
+            *map(set, [st.index.to_list() for st in variances_list])))
+    )
+
+    # Error when gene number is too small
+    if len(common_genes) < 300:
+        raise ValueError(f"Only {len(common_genes)} common genes between all samples.")
+
+    # Aggregate variances and identify HVGs
+    df = pd.concat(variances_list, axis=0)
+    df["highly_variable"] = df["highly_variable"].astype(int)
+
+    if is_count_data:
+        df = df.groupby("gene", observed=True).agg(
+            dict(
+                residual_variances="median",
+                highly_variable="sum",
+            )
+        )
+        df = df.loc[common_genes]
+        df["highly_variable_nbatches"] = df["highly_variable"]
+        df.sort_values(
+            ["highly_variable_nbatches", "residual_variances"],
+            ascending=False,
+            na_position="last",
+            inplace=True,
+        )
+    else:
+        df = df.groupby("gene", observed=True).agg(
+            dict(
+                dispersions_norm="median",
+                highly_variable="sum",
+            )
+        )
+        df = df.loc[common_genes]
+        df["highly_variable_nbatches"] = df["highly_variable"]
+        df.sort_values(
+            ["highly_variable_nbatches", "dispersions_norm"],
+            ascending=False,
+            na_position="last",
+            inplace=True,
+        )
+
+    hvg = df.iloc[: params.feat_cell,].index.tolist()
+
+    # Find the number of sampling cells for each batch
+    total_cell = np.sum(cell_number)
+    total_cell_training = np.minimum(total_cell, params.n_cell_training)
+    cell_proportion = cell_number / total_cell
+    n_cell_used = [
+        int(cell) for cell in (total_cell_training * cell_proportion).tolist()
+    ]
+
+    # Only use the common genes that can be transformed to human genes
+    if params.species is not None:
+        homologs = pd.read_csv(params.homolog_file, sep='\t')
+        if homologs.shape[1] < 2:
+            raise ValueError("Homologs file must have at least two columns: one for the species and one for the human gene symbol.")
+        homologs.columns = [params.species, 'HUMAN_GENE_SYM']
+        homologs.set_index(params.species, inplace=True)
+        common_genes = np.intersect1d(common_genes, homologs.index)
+        gene_homolog_dict = dict(zip(common_genes,homologs.loc[common_genes].HUMAN_GENE_SYM.values, strict=False))
+    else:
+        gene_homolog_dict = dict(zip(common_genes,common_genes, strict=False))
+
+    if len(gene_homolog_dict) < 300:
+        raise ValueError(f"Only {len(gene_homolog_dict)} genes could be mapped to human symbols. Please check the homolog file.")
+
+    return hvg, n_cell_used, gene_homolog_dict
+
+
+def create_subsampled_adata(sample_h5ad_dict, n_cell_used, params: FindLatentRepresentationsConfig):
+    """
+    Create subsampled adata for each sample with sample-specific stratified sampling,
+    add batch and label information, and return concatenated adata.
+
+    Args:
+        sample_h5ad_dict (dict): Dictionary mapping sample names to file paths of ST datasets.
+        n_cell_used (list): Number of cells to sample from each dataset.
+        params (object): Parameter object containing attributes.
+
+    Returns:
+        adata: Concatenated adata object with batch and label information in obs.
+    """
+    subsampled_adatas = []
+
+    logger.info("Creating subsampled adata with stratified sampling...")
+
+    for st_id, (sample_name, st_file) in enumerate(sample_h5ad_dict.items()):
+        logger.info(f"Processing {sample_name}...")
+
+        # Load the data
+        adata = sc.read_h5ad(st_file)
+
+        # Filter out mitochondrial and hemoglobin genes
+        gene_keep = ~adata.var_names.str.match(re.compile(r'^(HB.-|MT-)', re.IGNORECASE))
+        adata = adata[:,gene_keep].copy()
+
+        # Make gene names unique to avoid issues with concatenation
+        adata.var_names_make_unique()
+
+        # Setup data layer consistently
+        is_count_data, actual_data_layer = setup_data_layer(adata, params.data_layer)
+
+        # Filter cells based on annotation if provided
+        if params.annotation is not None:
+            adata = adata[~adata.obs[params.annotation].isnull()]
+
+        # Perform stratified sampling within this sample
+        if params.do_sampling and adata.n_obs > n_cell_used[st_id]:
+            if params.annotation is None:
+                # Simple random sampling
+                num_cell = n_cell_used[st_id]
+                logger.info(f"Downsampling {sample_name} to {num_cell} cells...")
+                random_indices = np.random.choice(adata.n_obs, num_cell, replace=False)
+                adata = adata[random_indices].copy()
+            else:
+                # Stratified sampling based on sample-specific annotation distribution
+                sample_annotation_counts = adata.obs[params.annotation].value_counts()
+                sample_total_cells = len(adata)
+                target_total_cells = n_cell_used[st_id]
+
+                # Calculate sample-specific annotation proportions
+                sample_annotation_proportions = sample_annotation_counts / sample_total_cells
+
+                # Calculate target cells for each annotation in this sample
+                target_cells_per_annotation = (sample_annotation_proportions * target_total_cells).astype(int)
+
+                logger.info(f"Downsampling {sample_name} to {target_total_cells} cells...")
+                logger.debug("---Sample-specific annotation distribution-----")
+                for ann, count in target_cells_per_annotation.items():
+                    logger.debug(f"{ann}: {count} cells")
+
+                # Perform stratified sampling
+                sampled_cells = (
+                    adata.obs.groupby(params.annotation, group_keys=False, observed=True)
+                    .apply(
+                        lambda x: x.sample(
+                            max(min(target_cells_per_annotation.get(x.name, 0), len(x)), 1),
+                            replace=False,
+                        ),
+                        include_groups=False
+                    )
+                    .index
+                )
+
+                # Filter adata to sampled cells
+                adata = adata[sampled_cells].copy()
+
+        # Add batch information to obs
+        adata.obs['batch_id'] = f"S{st_id}"
+        adata.obs['sample_name'] = sample_name
+
+        # Add label information to obs (ensure it's properly set)
+        if params.annotation is not None:
+            # The annotation column already exists, just ensure it's called 'label'
+            if 'label' not in adata.obs.columns:
+                adata.obs['label'] = adata.obs[params.annotation]
+        else:
+            # Create dummy labels
+            adata.obs['label'] = 'unknown'
+
+        subsampled_adatas.append(adata)
+        logger.info(f"Processed {sample_name}: {adata.n_obs} cells, {adata.n_vars} genes")
+
+    # Concatenate all samples
+    logger.info("Concatenating all processed data...")
+    concatenated_adata = sc.concat(subsampled_adatas, axis=0, join='inner',
+                                   index_unique='_', fill_value=0)
+
+    logger.info(f"Final concatenated adata: {concatenated_adata.n_obs} cells, {concatenated_adata.n_vars} genes")
+
+    return concatenated_adata
+
+def _looks_like_count_matrix(X, max_check=100, tol=1e-8):
+    """
+    Heuristically check whether X contains integer values.
+    """
+    if X is None:
+        return False
+
+    if sp.issparse(X):
+        data = X.data
+        if data.size == 0:
+            return False
+        if data.size > max_check:
+            data = data[:max_check]
+    else:
+        flat = np.asarray(X).ravel()
+        if flat.size == 0:
+            return False
+        if flat.size > max_check:
+            flat = flat[:max_check]
+        data = flat
+
+    nz = data[data != 0]
+    if nz.size == 0:
+        return True  # all zero, good enough
+    return np.allclose(nz, np.round(nz), atol=tol)
+
+
+def setup_data_layer(adata, data_layer, verbose=True):
+    """
+    Setup and validate data layer for adata object.
+
+    Args:
+        adata: AnnData object
+        data_layer: Requested data layer name
+        verbose: Whether to log information messages (default: True)
+
+    Returns:
+        tuple: (is_count_data, actual_data_layer)
+            - is_count_data: Boolean indicating if data is count data
+            - actual_data_layer: The actual layer name to use ("X" if using adata.X)
+    """
+    # Check if the requested data layer exists in layers
+    if  data_layer in adata.layers:
+        # Use the specified layer
+        adata.X = adata.layers[data_layer]
+        actual_data_layer = data_layer
+        if verbose:
+            logger.info(f"Using data layer: {data_layer}")
+    elif data_layer == "X":
+        actual_data_layer = "X"
+    else:
+        raise ValueError(f"Data layer '{data_layer}' not found in adata.layers.")
+
+    # Determine if this is count data
+    is_count_data = actual_data_layer in ["count", "counts", "raw_counts","impute_count"] or \
+                   (adata.X is not None and np.issubdtype(adata.X.dtype, np.integer)) or \
+                     _looks_like_count_matrix(adata.X)
+
+    if verbose:
+        if is_count_data:
+            logger.info("Detected count data")
+        else:
+            logger.info("Data appears to be normalized/log-transformed")
+
+    return is_count_data, actual_data_layer
+
+
+def normalize_for_analysis(adata, is_count_data, preserve_raw=True):
+    """
+    Normalize data for DEG analysis or module scoring if needed.
+
+    Args:
+        adata: AnnData object
+        is_count_data: Boolean indicating if data is count data
+        preserve_raw: Whether to preserve raw data in adata.raw
+
+    Returns:
+        adata: AnnData object with normalized data
+    """
+    if is_count_data:
+        logger.info("Normalizing and log-transforming count data...")
+
+        # Store raw counts if requested and not already stored
+        if preserve_raw and adata.raw is None:
+            adata.raw = adata
+
+        # Normalize to 10,000 reads per cell
+        sc.pp.normalize_total(adata, target_sum=1e4)
+
+        # Log-transform (log1p)
+        sc.pp.log1p(adata)
+
+        logger.info("Data normalized and log-transformed")
+    else:
+        logger.info("Using data as-is (already normalized/log-transformed)")
+
+    return adata
+
+
+def filter_significant_degs(deg_results, annotation, adata=None, pval_threshold=0.05, lfc_threshold=0.25, max_genes=100):
+    """
+    Filter DEGs based on statistical significance and fold change criteria.
+
+    Args:
+        deg_results: DEG results from scanpy rank_genes_groups
+        annotation: Annotation label to get DEGs for
+        adata: Optional adata to check gene existence
+        pval_threshold: P-value threshold (default: 0.05)
+        lfc_threshold: Log fold change threshold (default: 0.25)
+        max_genes: Maximum number of genes to return (default: 100)
+
+    Returns:
+        list: Filtered list of significant DEG gene names
+    """
+    gene_names = deg_results['names'][annotation]
+    pvals_adj = deg_results['pvals_adj'][annotation]
+    logfoldchanges = deg_results['logfoldchanges'][annotation]
+
+    # Filter genes based on significance and fold change
+    annotation_genes = []
+    for gene, pval, lfc in zip(gene_names, pvals_adj, logfoldchanges, strict=False):
+        if gene is not None and pval < pval_threshold and abs(lfc) > lfc_threshold:
+            annotation_genes.append(gene)
+
+    # Limit to max genes if we have more
+    if len(annotation_genes) > max_genes:
+        annotation_genes = annotation_genes[:max_genes]
+
+    # Ensure genes exist in adata if provided
+    if adata is not None:
+        annotation_genes = [gene for gene in annotation_genes if gene in adata.var_names]
+
+    return annotation_genes
+
+
+def calculate_module_scores_from_degs(adata, deg_results, annotation_key):
+    """
+    Calculate module scores using existing DEG results.
+
+    Args:
+        adata: AnnData object to calculate scores for
+        deg_results: DEG results from scanpy rank_genes_groups
+        annotation_key: Column name in obs containing annotation labels
+
+    Returns:
+        adata: Updated adata with module score columns
+    """
+
+    logger.info("Calculating module scores using existing DEG results...")
+
+    # Detect count data and normalize if needed
+    is_count_data = _looks_like_count_matrix(adata.X)
+    adata = normalize_for_analysis(adata, is_count_data, preserve_raw=False)
+
+    # Ensure annotation is categorical if it exists
+    if annotation_key in adata.obs.columns:
+        adata.obs[annotation_key] = adata.obs[annotation_key].astype('category')
+        available_annotations = adata.obs[annotation_key].cat.categories
+    else:
+        # If annotation doesn't exist, use all annotations from DEG results
+        available_annotations = list(deg_results['names'].dtype.names)
+
+    # Calculate module score for each annotation
+    for annotation in available_annotations:
+        if annotation in deg_results['names'].dtype.names:
+            # Get significant DEGs for this annotation
+            annotation_genes = filter_significant_degs(deg_results, annotation, adata)
+
+            if len(annotation_genes) > 0:
+                logger.info(f"Calculating module score for {annotation} using {len(annotation_genes)} genes")
+
+                # Calculate module score
+                sc.tl.score_genes(
+                    adata,
+                    gene_list=annotation_genes,
+                    score_name=f"{annotation}_module_score",
+                    use_raw=False
+                )
+            else:
+                logger.warning(f"No valid DEGs found for {annotation}")
+                adata.obs[f"{annotation}_module_score"] = 0.0
+        else:
+            logger.warning(f"Annotation {annotation} not found in DEG results")
+            adata.obs[f"{annotation}_module_score"] = 0.0
+
+    return adata
+
+
+def calculate_module_score(training_adata, annotation_key):
+    """
+    Perform DEG analysis for each annotation and calculate module scores.
+
+    Args:
+        training_adata: Concatenated training adata with annotation information
+        annotation_key: Column name in obs containing annotation labels
+
+    Returns:
+        training_adata: Updated adata with module scores for each annotation
+    """
+
+    logger.info("Performing DEG analysis for each annotation...")
+
+    # Make a copy to avoid modifying the original
+    adata = training_adata.copy()
+
+    # Ensure annotation is categorical
+    adata.obs[annotation_key] = adata.obs[annotation_key].astype('category')
+
+    # Detect count data and normalize if needed
+    is_count_data = _looks_like_count_matrix(adata.X)
+    adata = normalize_for_analysis(adata, is_count_data, preserve_raw=True)
+
+    # Perform DEG analysis with DataFrame fragmentation warnings suppressed
+    # These warnings are harmless and come from Scanpy's internal implementation
+    import warnings
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore', category=pd.errors.PerformanceWarning,
+                              message='DataFrame is highly fragmented')
+        sc.tl.rank_genes_groups(
+            adata,
+            groupby=annotation_key,
+            method='wilcoxon',
+            use_raw=False,
+            n_genes=100
+        )
+
+    logger.info("Calculating module scores for each annotation...")
+
+    # Get DEG results
+    deg_results = adata.uns['rank_genes_groups']
+
+    # Calculate module score for each annotation
+    for annotation in adata.obs[annotation_key].cat.categories:
+
+        # Get significant DEGs for this annotation
+        annotation_genes = filter_significant_degs(deg_results, annotation, adata)
+
+        if len(annotation_genes) > 0:
+            logger.info(f"Calculating module score for {annotation} using {len(annotation_genes)} genes")
+
+            # Calculate module score
+            sc.tl.score_genes(
+                adata,
+                gene_list=annotation_genes,
+                score_name=f"{annotation}_module_score",
+                use_raw=False
+            )
+        else:
+            logger.warning(f"No valid DEGs found for {annotation}")
+            adata.obs[f"{annotation}_module_score"] = 0.0
+
+    logger.info("Module score calculation completed")
+    return adata
+
+
+def apply_module_score_qc(adata, annotation_key, module_score_threshold_dict):
+    """
+    Apply quality control based on module scores.
+
+    Args:
+        adata: AnnData object to apply QC to
+        annotation_key: Column name in obs containing annotation labels
+        module_score_threshold_dict: Dictionary mapping annotation to threshold values
+
+    Returns:
+        adata: Updated adata with QC information
+    """
+    logger.info("Applying module score-based quality control...")
+
+    # Initialize High_quality column as boolean (True = high quality)
+    adata.obs['High_quality'] = True
+
+    # Check if we have the annotation key
+    if annotation_key not in adata.obs.columns:
+        logger.warning(f"Annotation key '{annotation_key}' not found in adata.obs. Skipping module score QC.")
+        return adata
+
+    # Apply QC for each annotation
+    for annotation, threshold in module_score_threshold_dict.items():
+        module_score_col = f"{annotation}_module_score"
+
+        if module_score_col in adata.obs.columns:
+            # Find cells of this annotation with low module scores
+            annotation_mask = adata.obs[annotation_key] == annotation
+            low_score_mask = adata.obs[module_score_col] < threshold
+
+            # Set High_quality to False for cells that match both conditions (low quality)
+            low_quality_mask = annotation_mask & low_score_mask
+            adata.obs.loc[low_quality_mask, 'High_quality'] = False
+
+            n_low_quality = low_quality_mask.sum()
+            n_annotation_cells = annotation_mask.sum()
+
+            logger.info(f"{annotation}: {n_low_quality}/{n_annotation_cells} cells marked as low quality "
+                       f"(threshold: {threshold:.3f})")
+        else:
+            logger.warning(f"Module score column '{module_score_col}' not found in adata.obs")
+
+    total_low_quality = (~adata.obs['High_quality']).sum()
+    logger.info(f"Total low quality cells: {total_low_quality}/{adata.n_obs}")
+
+    return adata
+
+
+# prepare the trainning data
+class TrainingData:
+    """
+    Managing and processing training data for graph-based models.
+
+    Attributes:
+        params (dict): A dictionary of parameters used for data processing and training.
+    """
+
+    def __init__(self, params):
+        self.params = params
+        self.gcov = GCN(self.params.K)
+        self.expression_merge = None
+        self.expression_gcn_merge = None
+        self.label_merge = None
+        self.batch_merge = None
+        self.batch_size = None
+        self.label_name = None
+
+
+    def prepare(self, concatenated_adata, hvg):
+        logger.info("Processing concatenated subsampled data...")
+
+        # Get labels from obs
+        if self.params.annotation is not None:
+            label = concatenated_adata.obs['label'].values
+        else:
+            label = np.zeros(concatenated_adata.n_obs)
+
+        # Get batch information from obs
+        batch_labels = concatenated_adata.obs['batch_id'].values
+
+        # Get expression array for HVG genes
+        expression_array = torch.Tensor(concatenated_adata[:, hvg].X.toarray())
+        logger.info(f"Expression array shape: {expression_array.shape}")
+
+        # Process each batch separately for GCN (since spatial graphs are sample-specific)
+        expression_array_gcn_list = []
+
+        for batch_id in concatenated_adata.obs['batch_id'].unique():
+            batch_mask = concatenated_adata.obs['batch_id'] == batch_id
+            batch_adata = concatenated_adata[batch_mask]
+            batch_expression = expression_array[batch_mask.values]
+
+            logger.info(f"Processing batch {batch_id} with {batch_adata.n_obs} cells...")
+
+            # Build spatial graph for this batch
+            edge = build_spatial_graph(
+                coords=np.array(batch_adata.obsm[self.params.spatial_key]),
+                n_neighbors=self.params.n_neighbors,
+            )
+            edge = torch.from_numpy(edge.T).long()
+
+            # Apply GCN to this batch
+            batch_expression_gcn = self.gcov(batch_expression, edge)
+            expression_array_gcn_list.append(batch_expression_gcn)
+
+            logger.info(f"Graph for {batch_id} has {edge.size(1)} edges, {batch_adata.n_obs} cells.")
+
+        # Concatenate GCN results in the same order as the original data
+        expression_array_gcn = torch.cat(expression_array_gcn_list, dim=0)
+
+        # Convert batch labels to numeric codes
+        batch_codes = pd.Categorical(batch_labels).codes
+
+        # Convert labels to categorical codes
+        cat_labels = pd.Categorical(label)
+        label_codes = cat_labels.codes
+
+        # Store results
+        self.expression_merge = expression_array
+        self.expression_gcn_merge = expression_array_gcn
+        self.batch_merge = torch.Tensor(batch_codes)
+        self.label_merge = torch.Tensor(label_codes).long()
+
+        if self.params.annotation is not None:
+            self.label_name = cat_labels.categories.take(np.unique(cat_labels.codes)).to_list()
+        else:
+            self.label_name = None
+
+        # Set batch size
+        self.batch_size = len(torch.unique(self.batch_merge))
+
+
+# Inference for each ST dataset
+class InferenceData:
+    """
+    Infer cell embeddings for each spatial transcriptomics (ST) dataset.
+    Attributes:
+        hvg: List of highly variable genes.
+        batch_size: Integer defining the batch size for inference.
+        model: Model to be used for inference.
+        params: Dictionary containing additional parameters for inference.
+    """
+
+    def __init__(self, hvg, batch_size, model, label_name, params):
+        self.params = params
+        self.gcov = GCN(self.params.K)
+        self.hvg = hvg
+        self.batch_size = batch_size
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = model.to(self.device)
+        self.label_name = label_name
+        self.processed_list_path = self.params.latent_dir / 'processed.list'
+
+
+    def infer_embedding_single(self, st_id, st_file) -> Path:
+        st_name = (Path(st_file).name).split(".h5ad")[0]
+        logger.info(f"Infering cell embeddings for {st_name}...")
+
+        # Load the ST data
+        adata = sc.read_h5ad(st_file)
+        # sc.pp.filter_genes(adata, min_counts=1)
+
+        # Make gene names unique to avoid reindexing issues
+        adata.var_names_make_unique()
+
+        # Setup data layer consistently
+        is_count_data, actual_data_layer = setup_data_layer(adata, self.params.data_layer)
+
+        # print(adata.shape)
+        # Convert expression data to torch.Tensor
+        expression_array = torch.Tensor(adata[:, self.hvg].X.toarray())
+
+        # Graph convolution of expression array
+        edge = build_spatial_graph(
+            coords=np.array(adata.obsm[self.params.spatial_key]),
+            n_neighbors=self.params.n_neighbors,
+        )
+        edge = torch.from_numpy(edge.T).long()
+        expression_array_gcn = self.gcov(expression_array, edge)
+
+        # Build batch vector as one-hot encoding
+        n_cell = adata.n_obs
+        batch_indices = torch.full((n_cell,), st_id, dtype=torch.long)
+
+        # Prepare the evaluation DataLoader
+        dataset = TensorDataset(expression_array_gcn,expression_array, batch_indices)
+        Inference_loader = DataLoader(dataset=dataset, batch_size=512, shuffle=False)
+
+        # Inference process
+        emb_cell, emb_niche, class_prob = [], [], []
+
+        for (
+            expression_gcn_focal,
+            expression_focal,
+            batch_indices_fcocal,
+        ) in Inference_loader:
+            expression_gcn_focal = expression_gcn_focal.to(self.device)
+            expression_focal = expression_focal.to(self.device)
+            batch_indices_fcocal = batch_indices_fcocal.to(self.device)
+
+            self.model.eval()
+            with torch.no_grad():
+                mu_focal = self.model.encode(
+                    [expression_focal, expression_gcn_focal], batch_indices_fcocal
+                )
+                _,x_class, _, _ = self.model(
+                    [expression_focal, expression_gcn_focal], batch_indices_fcocal
+                )
+
+                class_prob.append(x_class.cpu().numpy())
+                emb_cell.append(mu_focal[0].cpu().numpy())
+                emb_niche.append(mu_focal[1].cpu().numpy())
+
+        # Concatenate results and store embeddings in adata
+        emb_cell = np.concatenate(emb_cell, axis=0)
+        emb_niche = np.concatenate(emb_niche, axis=0)
+        class_prob = np.concatenate(class_prob, axis=0)
+
+        # if self.label_name is not None:
+        #     class_prob = pd.DataFrame(softmax(class_prob,axis=1), columns=self.label_name,index=adata.obs_names)
+        #     adata.obsm["class_prob"] = class_prob
+
+        adata.obsm[self.params.latent_representation_cell] = emb_cell
+        adata.obsm[self.params.latent_representation_niche] = emb_niche
+
+        return adata