Api tangram

class Tangram(object):
    r"""Tangram spatial deconvolution class for cell type mapping.

    Tangram is a method for integrating single-cell RNA sequencing (scRNA-seq) data
    with spatial transcriptomics data. It enables:
    1. Mapping cell types from scRNA-seq to spatial locations
    2. Deconvolving cell type proportions in spatial spots
    3. Imputing gene expression in spatial data
    4. Analyzing spatial organization of cell types

    The method works by:
    1. Identifying marker genes for each cell type
    2. Training a mapping model using these markers
    3. Projecting cell type annotations to spatial coordinates
    4. Optionally imputing full gene expression profiles

    Attributes:
        adata_sc: AnnData
            Single-cell RNA-seq data with:
            - Gene expression matrix in X
            - Cell type annotations in obs[clusters]
        adata_sp: AnnData
            Spatial transcriptomics data with:
            - Gene expression matrix in X
            - Spatial coordinates in obsm['spatial']
        clusters: str
            Column name in adata_sc.obs containing cell type labels
        markers: list
            Selected marker genes used for mapping
        ad_map: AnnData
            Mapping results after training

    Examples:
        >>> import scanpy as sc
        >>> import omicverse as ov
        >>> # Load data
        >>> adata_sc = sc.read_h5ad("sc_data.h5ad")
        >>> adata_sp = sc.read_visium("spatial_data")
        >>> # Initialize Tangram
        >>> tangram = ov.space.Tangram(
        ...     adata_sc=adata_sc,
        ...     adata_sp=adata_sp,
        ...     clusters='cell_type'
        ... )
        >>> # Train model
        >>> tangram.train(mode='clusters', num_epochs=500)
        >>> # Project cell types
        >>> adata_spatial = tangram.cell2location()
    """
    def check_tangram(self):
        r"""Check if Tangram package is installed.

        This method verifies that the Tangram package is available and prints
        its version number. If not installed, it raises an informative error.

        Raises:
            ImportError: If Tangram package is not installed, with instructions
                for installation.

        Notes:
            - Sets global tg_install flag when successful
            - Required before any Tangram operations
            - Suggests pip installation command if missing
        """
        global tg_install
        try:
            import tangram as tg
            tg_install=True
            print('tangram have been install version:',tg.__version__)
        except ImportError as e:
            raise ImportError(
                'Please install the tangram: `pip install -U tangram-sc`.'
            ) from e

    def __init__(self,
            adata_sc: AnnData,
            adata_sp: AnnData,
            clusters: str = '',
            marker_size: int = 100,
            gene_to_lowercase: bool = False
        ) -> None:
        r"""Initialize Tangram spatial deconvolution object.

        This method sets up the Tangram analysis by:
        1. Checking package installation
        2. Processing input data
        3. Identifying marker genes
        4. Preparing data for mapping

        Arguments:
            adata_sc: AnnData
                Single-cell RNA-seq data containing:
                - Normalized gene expression matrix
                - Cell type annotations in obs[clusters]
            adata_sp: AnnData
                Spatial transcriptomics data containing:
                - Normalized gene expression matrix
                - Spatial coordinates in obsm['spatial']
            clusters: str, optional (default='')
                Column name in adata_sc.obs containing cell type annotations.
            marker_size: int, optional (default=100)
                Number of top marker genes to select per cell type.
                More markers can improve accuracy but increase computation time.
            gene_to_lowercase: bool, optional (default=False)
                Whether to convert gene names to lowercase for matching between
                datasets. Useful when gene naming conventions differ.

        Notes:
            - Automatically filters genes present in at least one cell
            - Identifies marker genes using scanpy's rank_genes_groups
            - Prepares data structures for Tangram mapping
            - Adds reference annotation to both AnnData objects
        """
        self.check_tangram()
        global tg_install
        if tg_install==True:
            global_imports("tangram","tg")
        ad_map_dict={}

        adata_sc.uns['log1p']={}
        adata_sc.uns['log1p']['base']=None
        sc.pp.filter_genes(adata_sc, min_cells=1)
        sc.tl.rank_genes_groups(adata_sc, groupby=clusters,
                                key_added=f'{clusters}_rank_genes_groups',use_raw=False)
        markers_df = pd.DataFrame(adata_sc.uns[f"{clusters}_rank_genes_groups"]["names"]).iloc[0:marker_size, :]

        markers = list(np.unique(markers_df.melt().value.values))
        print('...Calculate The Number of Markers:',len(markers))

        self.adata_sc=adata_sc
        self.adata_sp=adata_sp
        self.clusters=clusters
        self.markers=markers

        tg.pp_adatas(self.adata_sc, self.adata_sp,
                      genes=self.markers,gene_to_lowercase=gene_to_lowercase)

        print('...Model prepared successfully')
        add_reference(self.adata_sc,'tangram','cell type classification with Tangram')
        add_reference(self.adata_sp,'tangram','cell type classification with Tangram')

    def train(self,
            mode: str = "clusters",
            num_epochs: int = 500,
            device: str = "cuda:0",
            **kwargs: Any
        ) -> None:
        r"""Train the Tangram spatial mapping model.

        This method trains a model to map cells or clusters from scRNA-seq data
        to spatial locations. It optimizes the mapping to preserve both gene
        expression patterns and spatial structure.

        Arguments:
            mode: str, optional (default="clusters")
                Mapping mode:
                - "clusters": Map cell type proportions (faster)
                - "cells": Map individual cells (more detailed)
            num_epochs: int, optional (default=500)
                Number of training epochs. More epochs may improve results
                but increase training time.
            device: str, optional (default="cuda:0")
                Computing device to use:
                - "cuda:0" (or other GPU index) for GPU acceleration
                - "cpu" for CPU computation
            **kwargs: Any
                Additional arguments passed to tangram.map_cells_to_space:
                - density_prior: Prior on spatial density
                - lambda_d: Density regularization strength
                - lambda_g1: Gene-expression regularization
                - lambda_g2: Spatial regularization
                - lambda_r: Entropy regularization

        Notes:
            - Automatically stores mapping in self.ad_map
            - Projects cell type annotations to spatial data
            - Adds reference annotation to both AnnData objects
            - Progress is shown during training
            - GPU acceleration recommended for large datasets
        """
        ad_map = tg.map_cells_to_space(self.adata_sc, self.adata_sp,
            mode=mode,
            cluster_label=self.clusters,
            num_epochs=num_epochs,
            device=device,
            **kwargs
        )

        tg.project_cell_annotations(ad_map, self.adata_sp, annotation=self.clusters)
        self.ad_map=ad_map

        print('...Model train successfully')
        add_reference(self.adata_sp,'tangram','cell type classification with Tangram')
        add_reference(self.adata_sc,'tangram','cell type classification with Tangram')

    def cell2location(self,annotation_list=None):
        r"""Project cell type annotations to spatial coordinates.

        This method creates a visualization-ready AnnData object containing the
        predicted cell type proportions for each spatial location.

        Arguments:
            annotation_list: list, optional (default=None)
                List of cell types to include in the projection.
                If None, uses all cell types from training data.

        Returns:
            AnnData
                Modified spatial data containing:
                - Original spatial data
                - Cell type proportions in obsm['tangram_ct_pred']
                - Normalized proportions in obs for each cell type

        Notes:
            - Automatically normalizes cell type proportions
            - Clips extreme values for better visualization
            - Adds reference annotation to both AnnData objects
            - Results can be directly used for spatial plotting
        """
        adata_plot=self.adata_sp.copy()
        if annotation_list is None:
            annotation_list=list(set(self.adata_sc.obs[self.clusters]))

        df = adata_plot.obsm["tangram_ct_pred"][annotation_list]
        construct_obs_plot(df, adata_plot, perc=0)
        add_reference(self.adata_sp,'tangram','cell type classification with Tangram')
        add_reference(self.adata_sc,'tangram','cell type classification with Tangram')
        return adata_plot

    def impute(self,
               ad_map: AnnData = None,
               ad_sc: AnnData = None,
                **kwargs: Any) -> AnnData:
        r"""Impute gene expression in spatial data using trained model.

        This method uses the trained mapping to predict the expression of all genes
        in the spatial locations, including genes not used in the mapping.

        Arguments:
            ad_map: AnnData, optional (default=None)
                Mapping result from train(). If None, uses self.ad_map.
            ad_sc: AnnData, optional (default=None)
                Single-cell reference data. If None, uses self.adata_sc.
            **kwargs: Any
                Additional arguments passed to tangram.project_genes:
                - scale: Whether to scale imputed values
                - filter_genes: Whether to filter genes before imputation
                - filter_threshold: Expression threshold for filtering

        Returns:
            AnnData
                Spatial data with imputed gene expression for all genes
                in the single-cell reference data.

        Notes:
            - Uses mapping weights to transfer expression
            - Can impute genes not used in original mapping
            - Useful for analyzing spatial patterns of any gene
            - Computationally intensive for large gene sets
        """
        if ad_map is None:
            ad_map=self.ad_map
        if ad_sc is None:
            ad_sc=self.adata_sc
        ad_ge = tg.project_genes(adata_map=ad_map, 
                                 adata_sc=ad_sc,**kwargs)
        return ad_ge