gsMap 1.62__py3-none-any.whl → 1.64__py3-none-any.whl

gsMap/diagnosis.py

@@ -0,0 +1,273 @@
+import logging
+import warnings
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import scanpy as sc
+from scipy.stats import norm
+
+from gsMap.config import DiagnosisConfig
+from gsMap.utils.manhattan_plot import ManhattanPlot
+from gsMap.visualize import draw_scatter, load_st_coord, estimate_point_size_for_plot
+
+
+warnings.filterwarnings("ignore", category=FutureWarning)
+logger = logging.getLogger(__name__)
+
+
+def convert_z_to_p(gwas_data):
+    """Convert Z-scores to P-values."""
+    gwas_data['P'] = norm.sf(abs(gwas_data['Z'])) * 2
+    min_p_value = 1e-300
+    gwas_data['P'] = gwas_data['P'].clip(lower=min_p_value)
+    return gwas_data
+
+
+def load_ldsc(ldsc_input_file):
+    """Load LDSC data and calculate logp."""
+    ldsc = pd.read_csv(ldsc_input_file, compression='gzip')
+    ldsc['spot'] = ldsc['spot'].astype(str).replace('\.0', '', regex=True)
+    ldsc.set_index('spot', inplace=True)
+    ldsc['logp'] = -np.log10(ldsc['p'])
+    return ldsc
+
+
+def load_gene_diagnostic_info(config:DiagnosisConfig):
+    """Load or compute gene diagnostic info."""
+    gene_diagnostic_info_save_path = config.get_gene_diagnostic_info_save_path(config.trait_name)
+    if gene_diagnostic_info_save_path.exists():
+        logger.info(f'Loading gene diagnostic information from {gene_diagnostic_info_save_path}...')
+        return pd.read_csv(gene_diagnostic_info_save_path)
+    else:
+        logger.info('Gene diagnostic information not found. Calculating gene diagnostic information...')
+        return compute_gene_diagnostic_info(config)
+
+
+def compute_gene_diagnostic_info(config: DiagnosisConfig):
+    """Calculate gene diagnostic info and save it to adata."""
+    logger.info('Loading ST data and LDSC results...')
+    # adata = sc.read_h5ad(config.hdf5_with_latent_path, backed='r')
+    mk_score = pd.read_feather(config.mkscore_feather_path)
+    mk_score.set_index('HUMAN_GENE_SYM', inplace=True)
+    mk_score = mk_score.T
+    trait_ldsc_result = load_ldsc(config.get_ldsc_result_file(config.trait_name))
+
+    # Align marker scores with trait LDSC results
+    mk_score = mk_score.loc[trait_ldsc_result.index]
+    mk_score = mk_score.loc[:, mk_score.sum(axis=0) != 0]
+
+    logger.info('Calculating correlation between gene marker scores and trait logp-values...')
+    corr = mk_score.corrwith(trait_ldsc_result['logp'])
+    corr.name = 'PCC'
+
+    grouped_mk_score = mk_score.groupby(adata.obs[config.annotation]).median()
+    max_annotations = grouped_mk_score.idxmax()
+
+    high_GSS_Gene_annotation_pair = pd.DataFrame({
+        'Gene': max_annotations.index,
+        'Annotation': max_annotations.values,
+        'Median_GSS': grouped_mk_score.max().values
+    })
+
+    # Filter based on median GSS score
+    high_GSS_Gene_annotation_pair = high_GSS_Gene_annotation_pair[high_GSS_Gene_annotation_pair['Median_GSS'] >= 1.0]
+    high_GSS_Gene_annotation_pair = high_GSS_Gene_annotation_pair.merge(corr, left_on='Gene', right_index=True)
+
+    # Prepare the final gene diagnostic info dataframe
+    gene_diagnostic_info_cols = ['Gene', 'Annotation', 'Median_GSS', 'PCC']
+    gene_diagnostic_info = high_GSS_Gene_annotation_pair[gene_diagnostic_info_cols].drop_duplicates().dropna(
+        subset=['Gene'])
+    gene_diagnostic_info.sort_values('PCC', ascending=False, inplace=True)
+
+    # Save gene diagnostic info to a file
+    gene_diagnostic_info_save_path = config.get_gene_diagnostic_info_save_path(config.trait_name)
+    gene_diagnostic_info.to_csv(gene_diagnostic_info_save_path, index=False)
+    logger.info(f'Gene diagnostic information saved to {gene_diagnostic_info_save_path}.')
+
+    # Save to adata.var with the trait_name prefix
+    logger.info('Saving gene diagnostic info to adata.var...')
+    gene_diagnostic_info.set_index('Gene', inplace=True)  # Use 'Gene' as the index to align with adata.var
+    adata.var[f'{config.trait_name}_Annotation'] = gene_diagnostic_info['Annotation']
+    adata.var[f'{config.trait_name}_Median_GSS'] = gene_diagnostic_info['Median_GSS']
+    adata.var[f'{config.trait_name}_PCC'] = gene_diagnostic_info['PCC']
+
+    # Save trait_ldsc_result to adata.obs
+    logger.info(f'Saving trait LDSC results to adata.obs as gsMap_{config.trait_name}_p_value...')
+    adata.obs[f'gsMap_{config.trait_name}_p_value'] = trait_ldsc_result['p']
+    adata.write(config.hdf5_with_latent_path, )
+
+    return gene_diagnostic_info.reset_index()
+
+
+def load_gwas_data(config:DiagnosisConfig):
+    """Load and process GWAS data."""
+    logger.info('Loading and processing GWAS data...')
+    gwas_data = pd.read_csv(config.sumstats_file, compression='gzip', sep='\t')
+    return convert_z_to_p(gwas_data)
+
+
+def load_snp_gene_pairs(config:DiagnosisConfig):
+    """Load SNP-gene pairs from multiple chromosomes."""
+    ldscore_save_dir = Path(config.ldscore_save_dir)
+    return pd.concat([
+        pd.read_feather(ldscore_save_dir / f'SNP_gene_pair/SNP_gene_pair_chr{chrom}.feather')
+        for chrom in range(1, 23)
+    ])
+
+
+def filter_snps(gwas_data_with_gene_annotation_sort, SUBSAMPLE_SNP_NUMBER):
+    """Filter the SNPs based on significance levels."""
+    pass_suggestive_line_mask = gwas_data_with_gene_annotation_sort['P'] < 1e-5
+    pass_suggestive_line_number = pass_suggestive_line_mask.sum()
+
+    if pass_suggestive_line_number > SUBSAMPLE_SNP_NUMBER:
+        snps2plot = gwas_data_with_gene_annotation_sort[pass_suggestive_line_mask].SNP
+        logger.info(f'To reduce the number of SNPs to plot, only {snps2plot.shape[0]} SNPs with P < 1e-5 are plotted.')
+    else:
+        snps2plot = gwas_data_with_gene_annotation_sort.head(SUBSAMPLE_SNP_NUMBER).SNP
+        logger.info(
+            f'To reduce the number of SNPs to plot, only {SUBSAMPLE_SNP_NUMBER} SNPs with the smallest P-values are plotted.')
+
+    return snps2plot
+
+
+def generate_manhattan_plot(config: DiagnosisConfig):
+    """Generate Manhattan plot."""
+    report_save_dir = config.get_report_dir(config.trait_name)
+    gwas_data = load_gwas_data(config)
+    snp_gene_pair = load_snp_gene_pairs(config)
+    gwas_data_with_gene = snp_gene_pair.merge(gwas_data, on='SNP', how='inner').rename(columns={'gene_name': 'GENE'})
+    gene_diagnostic_info = load_gene_diagnostic_info(config)
+    gwas_data_with_gene_annotation = gwas_data_with_gene.merge(gene_diagnostic_info, left_on='GENE', right_on='Gene',
+                                                               how='left')
+
+    gwas_data_with_gene_annotation = gwas_data_with_gene_annotation[
+        ~gwas_data_with_gene_annotation['Annotation'].isna()]
+    gwas_data_with_gene_annotation_sort = gwas_data_with_gene_annotation.sort_values('P')
+
+    snps2plot = filter_snps(gwas_data_with_gene_annotation_sort, SUBSAMPLE_SNP_NUMBER=100_000)
+    gwas_data_to_plot = gwas_data_with_gene_annotation[
+        gwas_data_with_gene_annotation['SNP'].isin(snps2plot)].reset_index(drop=True)
+    gwas_data_to_plot['Annotation_text'] = 'PCC: ' + gwas_data_to_plot['PCC'].round(2).astype(
+        str) + '<br>' + 'Annotation: ' + gwas_data_to_plot['Annotation'].astype(str)
+
+    fig = ManhattanPlot(
+        dataframe=gwas_data_to_plot,
+        title='gsMap Diagnosis Manhattan Plot',
+        point_size=3,
+        highlight_gene_list=config.selected_genes or gene_diagnostic_info.Gene.iloc[:config.top_corr_genes].tolist(),
+        suggestiveline_value=-np.log10(1e-5),
+        annotation='Annotation_text',
+    )
+
+    save_manhattan_plot_path = config.get_manhattan_html_plot_path(config.trait_name)
+    fig.write_html(save_manhattan_plot_path)
+    logger.info(f'Diagnostic Manhattan Plot saved to {save_manhattan_plot_path}.')
+
+
+def generate_GSS_distribution(config: DiagnosisConfig):
+    """Generate GSS distribution plots."""
+    # logger.info('Loading ST data...')
+    # adata = sc.read_h5ad(config.hdf5_with_latent_path)
+    mk_score = pd.read_feather(config.mkscore_feather_path).set_index('HUMAN_GENE_SYM').T
+
+    plot_genes = config.selected_genes or load_gene_diagnostic_info(config).Gene.iloc[:config.top_corr_genes].tolist()
+    if config.selected_genes is not None:
+        logger.info(f'Generating GSS & Expression distribution plot for selected genes: {plot_genes}...')
+    else:
+        logger.info(f'Generating GSS & Expression distribution plot for top {config.top_corr_genes} correlated genes...')
+
+    if config.customize_fig:
+        pixel_width, pixel_height, point_size = config.fig_width, config.fig_height, config.point_size
+    else:
+        (pixel_width, pixel_height), point_size = estimate_point_size_for_plot(adata.obsm['spatial'])
+    sub_fig_save_dir = config.get_GSS_plot_dir(config.trait_name)
+
+    # save plot gene list
+    config.get_GSS_plot_select_gene_file(config.trait_name).write_text('\n'.join(plot_genes))
+
+    for selected_gene in plot_genes:
+        expression_series = pd.Series(adata[:, selected_gene].X.toarray().flatten(), index=adata.obs.index,name='Expression')
+        threshold = np.quantile(expression_series,0.9999)
+        expression_series[expression_series > threshold] = threshold
+        generate_and_save_plots(adata, mk_score, expression_series, selected_gene, point_size, pixel_width,
+                                pixel_height, sub_fig_save_dir, config.sample_name, config.annotation)
+
+
+def generate_and_save_plots(adata, mk_score, expression_series, selected_gene, point_size, pixel_width, pixel_height,
+                            sub_fig_save_dir, sample_name, annotation):
+    """Generate and save the plots."""
+    select_gene_expression_with_space_coord = load_st_coord(adata, expression_series, annotation)
+    sub_fig_1 = draw_scatter(select_gene_expression_with_space_coord, title=f'{selected_gene} (Expression)',
+                             annotation='annotation', color_by='Expression', point_size=point_size, width=pixel_width,
+                             height=pixel_height)
+    save_plot(sub_fig_1, sub_fig_save_dir, sample_name, selected_gene, 'Expression')
+
+    select_gene_GSS_with_space_coord = load_st_coord(adata, mk_score[selected_gene].rename('GSS'), annotation)
+    sub_fig_2 = draw_scatter(select_gene_GSS_with_space_coord, title=f'{selected_gene} (GSS)', annotation='annotation',
+                             color_by='GSS', point_size=point_size, width=pixel_width, height=pixel_height)
+    save_plot(sub_fig_2, sub_fig_save_dir, sample_name, selected_gene, 'GSS')
+
+    # combined_fig = make_subplots(rows=1, cols=2,
+    #                              subplot_titles=(f'{selected_gene} (Expression)', f'{selected_gene} (GSS)'))
+    # for trace in sub_fig_1.data:
+    #     combined_fig.add_trace(trace, row=1, col=1)
+    # for trace in sub_fig_2.data:
+    #     combined_fig.add_trace(trace, row=1, col=2)
+    #
+
+def save_plot(sub_fig, sub_fig_save_dir, sample_name, selected_gene, plot_type):
+    """Save the plot to HTML and PNG."""
+    save_sub_fig_path = sub_fig_save_dir / f'{sample_name}_{selected_gene}_{plot_type}_Distribution.html'
+    # sub_fig.write_html(str(save_sub_fig_path))
+    sub_fig.update_layout(showlegend=False)
+    sub_fig.write_image(str(save_sub_fig_path).replace('.html', '.png'))
+
+
+def generate_gsMap_plot(config: DiagnosisConfig):
+    """Generate gsMap plot."""
+    logger.info('Creating gsMap plot...')
+
+    trait_ldsc_result = load_ldsc(config.get_ldsc_result_file(config.trait_name))
+    space_coord_concat = load_st_coord(adata, trait_ldsc_result, annotation=config.annotation)
+
+    if config.customize_fig:
+        pixel_width, pixel_height, point_size = config.fig_width, config.fig_height, config.point_size
+    else:
+        (pixel_width, pixel_height), point_size = estimate_point_size_for_plot(adata.obsm['spatial'])
+    fig = draw_scatter(space_coord_concat,
+                       title=f'{config.trait_name} (gsMap)',
+                       point_size=point_size,
+                       width=pixel_width,
+                       height=pixel_height,
+                       annotation=config.annotation
+                       )
+
+    output_dir = config.get_gsMap_plot_save_dir(config.trait_name)
+    output_file_html = config.get_gsMap_html_plot_save_path(config.trait_name)
+    output_file_png = output_file_html.with_suffix('.png')
+    output_file_csv = output_file_html.with_suffix('.csv')
+
+    fig.write_html(output_file_html)
+    fig.write_image(output_file_png)
+    space_coord_concat.to_csv(output_file_csv)
+
+    logger.info(f'gsMap plot created and saved in {output_dir}.')
+
+
+def run_Diagnosis(config: DiagnosisConfig):
+    """Main function to run the diagnostic plot generation."""
+    global adata
+    adata = sc.read_h5ad(config.hdf5_with_latent_path)
+    if 'log1p' not in adata.uns.keys() and adata.X.max() > 14: 
+        sc.pp.normalize_total(adata, target_sum=1e4)
+        sc.pp.log1p(adata)
+
+    if config.plot_type in ['manhattan', 'all']:
+        generate_manhattan_plot(config)
+    if config.plot_type in ['GSS', 'all']:
+        generate_GSS_distribution(config)
+    if config.plot_type in ['gsMap', 'all']:
+        generate_gsMap_plot(config)
+

gsMap/find_latent_representation.py

@@ -1,9 +1,5 @@
-import argparse
 import logging
-import pprint
 import random
-import time
-from pathlib import Path
 
 import numpy as np
 import pandas as pd
@@ -13,17 +9,9 @@ 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 add_find_latent_representations_args, FindLatentRepresentationsConfig
-
-# seed all
+from gsMap.config import FindLatentRepresentationsConfig
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-handler = logging.StreamHandler()
-handler.setFormatter(logging.Formatter(
-    '[{asctime}] {levelname:8s} {filename} {message}', style='{'))
-logger.addHandler(handler)
-
 
 def set_seed(seed_value):
     """
@@ -33,30 +21,30 @@ def set_seed(seed_value):
     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():
-        print('Running use GPU')
+        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:
-        print('Running use CPU')
+        logger.info('Running use CPU')
 
-set_seed(2024)
 
 # The class for finding latent representations
 class Latent_Representation_Finder:
 
-    def __init__(self, adata, Params):
+    def __init__(self, adata, args:FindLatentRepresentationsConfig):
         self.adata = adata.copy()
-        self.Params = Params
+        self.Params = args
 
         # Standard process
-        if self.Params.type == 'count' or self.Params.type == 'counts':
-            self.adata.X = self.adata.layers[self.Params.type]
+        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:
-            self.adata.X = self.adata.layers[self.Params.type]
+            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'):
@@ -66,7 +54,7 @@ class Latent_Representation_Finder:
 
         # Process the feature matrix
         node_X = self.adata[:, self.adata.var.highly_variable].X
-        print(f'The shape of feature matrix is {node_X.shape}.')
+        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)
 
@@ -75,7 +63,7 @@ class Latent_Representation_Finder:
         self.Params.feat_cell = node_X.shape[1]
 
         # Run GNN-VAE
-        print(f'------Finding latent representations for {verbose}...')
+        logger.info(f'------Finding latent representations for {verbose}...')
         gvae = Model_Train(node_X, graph_dict, self.Params, label)
         gvae.run_train()
 
@@ -87,20 +75,21 @@ class Latent_Representation_Finder:
 
 
 def run_find_latent_representation(args:FindLatentRepresentationsConfig):
+    set_seed(2024)
     num_features = args.feat_cell
-    args.output_dir = Path(args.output_hdf5_path).parent
-    args.output_dir.mkdir(parents=True, exist_ok=True,mode=0o755)
+    args.hdf5_with_latent_path.parent.mkdir(parents=True, exist_ok=True,mode=0o755)
     # Load the ST data
-    print(f'------Loading ST data of {args.sample_name}...')
+    logger.info(f'------Loading ST data of {args.sample_name}...')
     adata = sc.read_h5ad(f'{args.input_hdf5_path}')
     adata.var_names_make_unique()
-    print('The ST data contains %d cells, %d genes.' % (adata.shape[0], adata.shape[1]))
+    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()),]
+        adata = adata[adata.obs[args.annotation].isin(num[num >= 30].index.to_list())]
 
         le = preprocessing.LabelEncoder()
         le.fit(adata.obs[args.annotation])
@@ -113,7 +102,7 @@ def run_find_latent_representation(args:FindLatentRepresentationsConfig):
     latent_GVAE = latent_rep.Run_GNN_VAE(label)
     latent_PCA = latent_rep.Run_PCA()
     # Add latent representations to the spe data
-    print(f'------Adding latent representations...')
+    logger.info(f'------Adding latent representations...')
     adata.obsm["latent_GVAE"] = latent_GVAE
     adata.obsm["latent_PCA"] = latent_PCA
     # Run umap based on latent representations
@@ -124,13 +113,13 @@ def run_find_latent_representation(args:FindLatentRepresentationsConfig):
 
         # Find the latent representations hierarchically (optionally)
     if not args.annotation is None and args.hierarchically:
-        print(f'------Finding latent representations 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, :]
-            print(adata_part.shape)
+            logger.info(adata_part.shape)
 
             # Find latent representations for the selected ct
             latent_rep = Latent_Representation_Finder(adata_part, args)
@@ -151,59 +140,6 @@ def run_find_latent_representation(args:FindLatentRepresentationsConfig):
 
             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,])
-    print(f'------Saving ST data...')
-    adata.write(args.output_hdf5_path)
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description="This script is designed to find latent representations in spatial transcriptomics data using a Graph Neural Network Variational Autoencoder (GNN-VAE). It processes input data, constructs a neighboring graph, and runs GNN-VAE to output latent representations.")
-    add_find_latent_representations_args(parser)
-    TEST=True
-    if TEST:
-        test_dir = '/storage/yangjianLab/chenwenhao/projects/202312_gsMap/data/gsMap_test/Nature_Neuroscience_2021'
-        name = 'Cortex_151507'
+    logger.info(f'------Saving ST data...')
+    adata.write(args.hdf5_with_latent_path)
 
-
-        args = parser.parse_args(
-            [
-                '--input_hdf5_path','/storage/yangjianLab/songliyang/SpatialData/Data/Brain/Human/Nature_Neuroscience_2021/processed/h5ad/Cortex_151507.h5ad',
-                '--output_hdf5_path',f'{test_dir}/{name}/hdf5/{name}_add_latent.h5ad',
-                '--sample_name', name,
-                '--annotation','layer_guess',
-                '--type','count',
-            ]
-
-        )
-
-    else:
-        args = parser.parse_args()
-    config=FindLatentRepresentationsConfig(**{'annotation': 'SubClass',
- 'convergence_threshold': 0.0001,
- 'epochs': 300,
- 'feat_cell': 3000,
- 'feat_hidden1': 256,
- 'feat_hidden2': 128,
- 'gcn_decay': 0.01,
- 'gcn_hidden1': 64,
- 'gcn_hidden2': 30,
- 'gcn_lr': 0.001,
- 'hierarchically': False,
- 'input_hdf5_path': '/storage/yangjianLab/songliyang/SpatialData/Data/Brain/macaque/Cell/processed/h5ad/T862_macaque3.h5ad',
- 'label_w': 1.0,
- 'n_comps': 300,
- 'n_neighbors': 11,
- 'nheads': 3,
- 'output_hdf5_path': 'T862_macaque3/find_latent_representations/T862_macaque3_add_latent.h5ad',
- 'p_drop': 0.1,
- 'rec_w': 1.0,
- 'sample_name': 'T862_macaque3',
- 'type': 'SCT',
- 'var': False,
- 'weighted_adj': False})
-    # config=FindLatentRepresentationsConfig(**vars(args))
-    start_time = time.time()
-    logger.info(f'Find latent representations for {config.sample_name}...')
-    pprint.pprint(config)
-    run_find_latent_representation(config)
-    end_time = time.time()
-    logger.info(f'Find latent representations for {config.sample_name} finished. Time spent: {(end_time - start_time) / 60:.2f} min.')