gsMap 1.66__py3-none-any.whl → 1.70__py3-none-any.whl

gsMap/find_latent_representation.py

@@ -1,145 +1,133 @@
-import logging
-import random
-
-import numpy as np
-import pandas as pd
-import scanpy as sc
-import torch
-from sklearn import preprocessing
-
-from gsMap.GNN_VAE.adjacency_matrix import Construct_Adjacency_Matrix
-from gsMap.GNN_VAE.train import Model_Train
-from gsMap.config import FindLatentRepresentationsConfig
-
-logger = logging.getLogger(__name__)
-
-def set_seed(seed_value):
-    """
-    Set seed for reproducibility in PyTorch.
-    """
-    torch.manual_seed(seed_value)  # Set the seed for PyTorch
-    np.random.seed(seed_value)  # Set the seed for NumPy
-    random.seed(seed_value)  # Set the seed for Python random module
-    if torch.cuda.is_available():
-        logger.info('Running use GPU')
-        torch.cuda.manual_seed(seed_value)  # Set seed for all CUDA devices
-        torch.cuda.manual_seed_all(seed_value)  # Set seed for all CUDA devices
-    else:
-        logger.info('Running use CPU')
-
-
-# The class for finding latent representations
-class Latent_Representation_Finder:
-
-    def __init__(self, adata, args:FindLatentRepresentationsConfig):
-        self.adata = adata.copy()
-        self.Params = args
-
-        # Standard process
-        if self.Params.data_layer == 'count' or self.Params.data_layer == 'counts':
-            self.adata.X = self.adata.layers[self.Params.data_layer]
-            sc.pp.highly_variable_genes(self.adata, flavor="seurat_v3", n_top_genes=self.Params.feat_cell)
-            sc.pp.normalize_total(self.adata, target_sum=1e4)
-            sc.pp.log1p(self.adata)
-            sc.pp.scale(self.adata)
-        else:
-            if self.Params.data_layer != 'X':
-                self.adata.X = self.adata.layers[self.Params.data_layer]
-            sc.pp.highly_variable_genes(self.adata, n_top_genes=self.Params.feat_cell)
-
-    def Run_GNN_VAE(self, label, verbose='whole ST data'):
-
-        # Construct the neighbouring graph
-        graph_dict = Construct_Adjacency_Matrix(self.adata, self.Params)
-
-        # Process the feature matrix
-        node_X = self.adata[:, self.adata.var.highly_variable].X
-        logger.info(f'The shape of feature matrix is {node_X.shape}.')
-        if self.Params.input_pca:
-            node_X = sc.pp.pca(node_X, n_comps=self.Params.n_comps)
-
-        # Update the input shape
-        self.Params.n_nodes = node_X.shape[0]
-        self.Params.feat_cell = node_X.shape[1]
-
-        # Run GNN-VAE
-        logger.info(f'------Finding latent representations for {verbose}...')
-        gvae = Model_Train(node_X, graph_dict, self.Params, label)
-        gvae.run_train()
-
-        return gvae.get_latent()
-
-    def Run_PCA(self):
-        sc.tl.pca(self.adata)
-        return self.adata.obsm['X_pca'][:, 0:self.Params.n_comps]
-
-
-def run_find_latent_representation(args:FindLatentRepresentationsConfig):
-    set_seed(2024)
-    num_features = args.feat_cell
-    args.hdf5_with_latent_path.parent.mkdir(parents=True, exist_ok=True,mode=0o755)
-    # Load the ST data
-    logger.info(f'------Loading ST data of {args.sample_name}...')
-    adata = sc.read_h5ad(f'{args.input_hdf5_path}')
-    adata.var_names_make_unique()
-    adata.X = adata.layers[args.data_layer] if args.data_layer in adata.layers.keys() else adata.X
-    logger.info('The ST data contains %d cells, %d genes.' % (adata.shape[0], adata.shape[1]))
-    # Load the cell type annotation
-    if not args.annotation is None:
-        # remove cells without enough annotations
-        adata = adata[~pd.isnull(adata.obs[args.annotation]), :]
-        num = adata.obs[args.annotation].value_counts()
-        adata = adata[adata.obs[args.annotation].isin(num[num >= 30].index.to_list())]
-
-        le = preprocessing.LabelEncoder()
-        le.fit(adata.obs[args.annotation])
-        adata.obs['categorical_label'] = le.transform(adata.obs[args.annotation])
-        label = adata.obs['categorical_label'].to_list()
-    else:
-        label = None
-    # Find latent representations
-    latent_rep = Latent_Representation_Finder(adata, args)
-    latent_GVAE = latent_rep.Run_GNN_VAE(label)
-    latent_PCA = latent_rep.Run_PCA()
-    # Add latent representations to the spe data
-    logger.info(f'------Adding latent representations...')
-    adata.obsm["latent_GVAE"] = latent_GVAE
-    adata.obsm["latent_PCA"] = latent_PCA
-    # Run umap based on latent representations
-    for name in ['latent_GVAE', 'latent_PCA']:
-        sc.pp.neighbors(adata, n_neighbors=10, use_rep=name)
-        sc.tl.umap(adata)
-        adata.obsm['X_umap_' + name] = adata.obsm['X_umap']
-
-        # Find the latent representations hierarchically (optionally)
-    if not args.annotation is None and args.hierarchically:
-        logger.info(f'------Finding latent representations hierarchically...')
-        PCA_all = pd.DataFrame()
-        GVAE_all = pd.DataFrame()
-
-        for ct in adata.obs[args.annotation].unique():
-            adata_part = adata[adata.obs[args.annotation] == ct, :]
-            logger.info(adata_part.shape)
-
-            # Find latent representations for the selected ct
-            latent_rep = Latent_Representation_Finder(adata_part, args)
-
-            latent_PCA_part = pd.DataFrame(latent_rep.Run_PCA())
-            if adata_part.shape[0] <= args.n_comps:
-                latent_GVAE_part = latent_PCA_part
-            else:
-                latent_GVAE_part = pd.DataFrame(latent_rep.Run_GNN_VAE(label=None, verbose=ct))
-
-            latent_GVAE_part.index = adata_part.obs_names
-            latent_PCA_part.index = adata_part.obs_names
-
-            GVAE_all = pd.concat((GVAE_all, latent_GVAE_part), axis=0)
-            PCA_all = pd.concat((PCA_all, latent_PCA_part), axis=0)
-
-            args.feat_cell = num_features
-
-            adata.obsm["latent_GVAE_hierarchy"] = np.array(GVAE_all.loc[adata.obs_names,])
-            adata.obsm["latent_PCA_hierarchy"] = np.array(PCA_all.loc[adata.obs_names,])
-    logger.info(f'------Saving ST data...')
-    adata.write(args.hdf5_with_latent_path)
-
+import logging
+import random
+import numpy as np
+import scanpy as sc
+import torch
+from sklearn.decomposition import PCA
+from sklearn.preprocessing import LabelEncoder
+from gsMap.GNN.adjacency_matrix import construct_adjacency_matrix
+from gsMap.GNN.train import ModelTrainer
+from gsMap.config import FindLatentRepresentationsConfig
+
+logger = logging.getLogger(__name__)
+
+
+def set_seed(seed_value):
+    """
+    Set seed for reproducibility in PyTorch and other libraries.
+    """
+    torch.manual_seed(seed_value)
+    np.random.seed(seed_value)
+    random.seed(seed_value)
+    if torch.cuda.is_available():
+        logger.info('Using GPU for computations.')
+        torch.cuda.manual_seed(seed_value)
+        torch.cuda.manual_seed_all(seed_value)
+    else:
+        logger.info('Using CPU for computations.')
+
+def preprocess_data(adata, params):
+    """
+    Preprocess the AnnData
+    """
+    logger.info('Preprocessing data...')
+    adata.var_names_make_unique()
+
+    sc.pp.filter_genes(adata, min_cells=30)
+    if params.data_layer in adata.layers.keys():
+        logger.info(f'Using data layer: {params.data_layer}...')
+        adata.X = adata.layers[params.data_layer]
+    else:
+        raise ValueError(f'Invalid data layer: {params.data_layer}, please check the input data.')
+
+    if params.data_layer in ['count', 'counts']:
+        # HVGs based on count
+        sc.pp.highly_variable_genes(adata,flavor="seurat_v3",n_top_genes=params.feat_cell)
+        # Normalize the data
+        sc.pp.normalize_total(adata, target_sum=1e4)
+        sc.pp.log1p(adata)
+
+    elif params.data_layer in adata.layers.keys():
+        sc.pp.highly_variable_genes(adata,flavor="seurat",n_top_genes=params.feat_cell)
+
+    return adata
+
+
+class LatentRepresentationFinder:
+    def __init__(self, adata, args: FindLatentRepresentationsConfig):
+        self.params = args
+
+        self.expression_array = adata[:, adata.var.highly_variable].X.copy()
+        self.expression_array = sc.pp.scale(self.expression_array, max_value=10)
+
+        # Construct the neighboring graph
+        self.graph_dict = construct_adjacency_matrix(adata, self.params)
+
+    def compute_pca(self):
+        self.latent_pca = PCA(n_components=self.params.feat_cell).fit_transform(self.expression_array)
+        return self.latent_pca
+
+    def run_gnn_vae(self, label, verbose='whole ST data'):
+
+        # Use PCA if specified
+        if self.params.input_pca:
+            node_X = self.compute_pca()
+        else:
+            node_X = self.expression_array
+
+        # Update the input shape
+        self.params.n_nodes = node_X.shape[0]
+        self.params.feat_cell = node_X.shape[1]
+
+        # Run GNN
+        logger.info(f'Finding latent representations for {verbose}...')
+        gvae = ModelTrainer(node_X, self.graph_dict, self.params, label)
+        gvae.run_train()
+
+        del self.graph_dict
+
+        return gvae.get_latent()
+
+
+def run_find_latent_representation(args: FindLatentRepresentationsConfig):
+    set_seed(2024)
+
+    # Load the ST data
+    logger.info(f'Loading ST data of {args.sample_name}...')
+    adata = sc.read_h5ad(args.input_hdf5_path)
+    logger.info(f'The ST data contains {adata.shape[0]} cells, {adata.shape[1]} genes.')
+
+    # Load the cell type annotation
+    if args.annotation is not None:
+        # Remove cells without enough annotations
+        adata = adata[~adata.obs[args.annotation].isnull()]
+        num = adata.obs[args.annotation].value_counts()
+        valid_annotations = num[num >= 30].index.to_list()
+        adata = adata[adata.obs[args.annotation].isin(valid_annotations)]
+
+        le = LabelEncoder()
+        label =  le.fit_transform(adata.obs[args.annotation])
+    else:
+        label = None
+
+    # Preprocess data
+    adata = preprocess_data(adata, args)
+
+    latent_rep = LatentRepresentationFinder(adata, args)
+    latent_gvae = latent_rep.run_gnn_vae(label)
+    latent_pca = latent_rep.latent_pca
+
+    # Add latent representations to the AnnData object
+    logger.info('Adding latent representations...')
+    adata.obsm["latent_GVAE"] = latent_gvae
+    adata.obsm["latent_PCA"] = latent_pca
+
+    # Run UMAP based on latent representations
+    #for name in ['latent_GVAE', 'latent_PCA']:
+    #    sc.pp.neighbors(adata, n_neighbors=10, use_rep=name)
+    #    sc.tl.umap(adata)
+    #    adata.obsm['X_umap_' + name] = adata.obsm['X_umap']
+
+    # Save the AnnData object
+    logger.info('Saving ST data...')
+    adata.write(args.hdf5_with_latent_path)