gsMap 1.71__py3-none-any.whl → 1.71.2__py3-none-any.whl

gsMap/GNN/model.py

@@ -85,5 +85,6 @@ class GATModel(nn.Module):
         mu, logvar = self.encode(x, edge_index)
         z = self.reparameterize(mu, logvar)
         x_reconstructed = self.decoder(z)
-        pred_label = F.softmax(self.cluster(z), dim=1)
+        # pred_label = F.softmax(self.cluster(z), dim=1)
+        pred_label = self.cluster(z)
         return pred_label, x_reconstructed, z, mu, logvar

gsMap/__init__.py

@@ -2,4 +2,4 @@
 Genetics-informed pathogenic spatial mapping
 '''
 
-__version__ = '1.71'
+__version__ = '1.71.2'

gsMap/config.py

@@ -412,6 +412,7 @@ class LatentToGeneConfig(ConfigWithAutoPaths):
     homolog_file: str = None
     gM_slices: str = None
     annotation: str = None
+    species: str = None
 
     def __post_init__(self):
         if self.homolog_file is not None:

gsMap/diagnosis.py

@@ -85,17 +85,18 @@ def compute_gene_diagnostic_info(config: DiagnosisConfig):
     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, )
+    # TODO: A new script is needed to save the gene diagnostic info to adata.var and trait_ldsc_result to adata.obs when running multiple traits
+    # # 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()
 

gsMap/find_latent_representation.py

@@ -33,15 +33,21 @@ def preprocess_data(adata, params):
     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]
+        sc.pp.filter_genes(adata, min_cells=30)
+    elif params.data_layer == 'X':
+        logger.info(f'Using data layer: {params.data_layer}...')
+        if adata.X.dtype == 'float32' or adata.X.dtype == 'float64':
+            logger.warning(f'The data layer should be raw count data')
+        sc.pp.filter_genes(adata, min_cells=30)
     else:
         raise ValueError(f'Invalid data layer: {params.data_layer}, please check the input data.')
 
-    if params.data_layer in ['count', 'counts']:
+    if params.data_layer in ['count', 'counts', 'X']:
         # HVGs based on count
+        logger.info('Dealing with count data...')
         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)

gsMap/latent_to_gene.py

@@ -4,6 +4,7 @@ from pathlib import Path
 import numpy as np
 import pandas as pd
 import scanpy as sc
+import scipy
 from scipy.stats import gmean
 from scipy.stats import rankdata
 from sklearn.metrics.pairwise import cosine_similarity
@@ -62,7 +63,7 @@ def build_spatial_net(adata, annotation, num_neighbour):
         logger.info(f'Cell annotations are not provided...')
         spatial_net = find_neighbors(coor, num_neighbour)
 
-    return spatial_net
+    return spatial_net.groupby('Cell1')['Cell2'].apply(np.array).to_dict()
 
 
 def find_neighbors_regional(cell_pos, spatial_net_dict, coor_latent, config, cell_annotations):
@@ -128,18 +129,20 @@ def compute_regional_mkscore(cell_pos, spatial_net_dict, coor_latent, config, ce
 def run_latent_to_gene(config: LatentToGeneConfig):
     logger.info('------Loading the spatial data...')
     adata = sc.read_h5ad(config.hdf5_with_latent_path)
+    logger.info(f'Loaded spatial data with {adata.n_obs} cells and {adata.n_vars} genes.')
 
     if config.annotation is not None:
         logger.info(f'------Cell annotations are provided as {config.annotation}...')
+        initial_cell_count = adata.n_obs
         adata = adata[~pd.isnull(adata.obs[config.annotation]), :]
+        logger.info(f'Removed null annotations. Cells retained: {adata.n_obs} (initial: {initial_cell_count}).')
 
     # Homologs transformation
     if config.homolog_file is not None:
         logger.info(f'------Transforming the {config.species} to HUMAN_GENE_SYM...')
         homologs = pd.read_csv(config.homolog_file, sep='\t')
         if homologs.shape[1] != 2:
-            raise ValueError(
-                "Homologs file must have two columns: one for the species and one for the human gene symbol.")
+            raise ValueError("Homologs file must have two columns: one for the species and one for the human gene symbol.")
 
         homologs.columns = [config.species, 'HUMAN_GENE_SYM']
         homologs.set_index(config.species, inplace=True)
@@ -156,35 +159,30 @@ def run_latent_to_gene(config: LatentToGeneConfig):
 
     if config.annotation is not None:
         cell_annotations = adata.obs[config.annotation].values
+        logger.info(f'Using cell annotations for {len(cell_annotations)} cells.')
     else:
         cell_annotations = None
 
     # Build the spatial graph
-    spatial_net = build_spatial_net(adata, config.annotation, config.num_neighbour_spatial)
-    spatial_net_dict = spatial_net.groupby('Cell1')['Cell2'].apply(np.array).to_dict()
+    logger.info('------Building the spatial graph...')
+    spatial_net_dict = build_spatial_net(adata, config.annotation, config.num_neighbour_spatial)
+    logger.info('Spatial graph built successfully.')
 
     # Extract the latent representation
+    logger.info('------Extracting the latent representation...')
     coor_latent = adata.obsm[config.latent_representation]
     coor_latent = coor_latent.astype(np.float32)
-
-    # Compute ranks
-    logger.info('------Ranking the spatial data...')
-    adata_X = adata.X.tocsr()
-    ranks = np.zeros((n_cells, n_genes), dtype=np.float32)
-
-    for i in tqdm(range(n_cells), desc="Computing ranks per cell"):
-        data = adata_X[i, :].toarray().flatten()
-        ranks[i, :] = rankdata(data, method='average')
+    logger.info('Latent representation extracted.')
 
     # Geometric mean across slices
+    gM = None
     if config.gM_slices is not None:
         logger.info('Geometrical mean across multiple slices is provided.')
         gM_df = pd.read_parquet(config.gM_slices)
         if config.species is not None:
-            homologs = pd.read_csv(config.homolog_file, sep='\t', header=None)
+            homologs = pd.read_csv(config.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.")
+                raise ValueError("Homologs file must have at least two columns: one for the species and one for the human gene symbol.")
             homologs.columns = [config.species, 'HUMAN_GENE_SYM']
             homologs.set_index(config.species, inplace=True)
             gM_df = gM_df.loc[gM_df.index.isin(homologs.index)]
@@ -193,20 +191,36 @@ def run_latent_to_gene(config: LatentToGeneConfig):
         gM_df = gM_df.loc[common_genes]
         gM = gM_df['G_Mean'].values
         adata = adata[:, common_genes]
-        ranks = ranks[:, np.isin(adata.var_names, common_genes)]
+        logger.info(f'{len(common_genes)} common genes retained after loading the cross slice geometric mean.')
+
+    # Compute ranks after taking common genes with gM_slices
+    logger.info('------Ranking the spatial data...')
+    if not scipy.sparse.issparse(adata.X):
+        adata_X = scipy.sparse.csr_matrix(adata.X)
+    elif isinstance(adata.X, scipy.sparse.csr_matrix):
+        adata_X = adata.X  # Avoid copying if already CSR
     else:
+        adata_X = adata.X.tocsr()
+
+    ranks = np.zeros((n_cells, adata.n_vars), dtype=np.float32)
+
+    for i in tqdm(range(n_cells), desc="Computing ranks per cell"):
+        data = adata_X[i, :].toarray().flatten()
+        ranks[i, :] = rankdata(data, method='average')
+
+    if gM is None:
         gM = gmean(ranks, axis=0)
 
     # Compute the fraction of each gene across cells
     adata_X_bool = adata_X.astype(bool)
     frac_whole = np.asarray(adata_X_bool.sum(axis=0)).flatten() / n_cells
+    logger.info('Gene expression proportion of each gene across cells computed.')
 
     # Normalize the ranks
-    ranks = ranks / gM
+    ranks /= gM
 
     # Compute marker scores in parallel
     logger.info('------Computing marker scores...')
-
     def compute_mk_score_wrapper(cell_pos):
         return compute_regional_mkscore(
             cell_pos, spatial_net_dict, coor_latent, config, cell_annotations, ranks, frac_whole, adata_X_bool
@@ -214,12 +228,14 @@ def run_latent_to_gene(config: LatentToGeneConfig):
 
     mk_scores = [compute_mk_score_wrapper(cell_pos) for cell_pos in tqdm(range(n_cells), desc="Calculating marker scores")]
     mk_score = np.vstack(mk_scores).T
+    logger.info('Marker scores computed.')
 
     # Remove mitochondrial genes
     gene_names = adata.var_names.values.astype(str)
     mt_gene_mask = ~(np.char.startswith(gene_names, 'MT-') | np.char.startswith(gene_names, 'mt-'))
     mk_score = mk_score[mt_gene_mask, :]
     gene_names = gene_names[mt_gene_mask]
+    logger.info(f'Removed mitochondrial genes. Remaining genes: {len(gene_names)}.')
 
     # Save the marker scores
     logger.info(f'------Saving marker scores ...')
@@ -229,6 +245,8 @@ def run_latent_to_gene(config: LatentToGeneConfig):
     mk_score_df.reset_index(inplace=True)
     mk_score_df.rename(columns={'index': 'HUMAN_GENE_SYM'}, inplace=True)
     mk_score_df.to_feather(output_file_path)
+    logger.info(f'Marker scores saved to {output_file_path}.')
 
     # Save the modified adata object to disk
     adata.write(config.hdf5_with_latent_path)
+    logger.info(f'Modified adata object saved to {config.hdf5_with_latent_path}.')

gsMap/spatial_ldsc_multiple_sumstats.py

@@ -2,6 +2,7 @@ import gc
 import logging
 import os
 from collections import defaultdict
+from functools import partial
 from pathlib import Path
 
 import anndata as ad
@@ -18,14 +19,15 @@ from gsMap.utils.regression_read import _read_sumstats, _read_w_ld, _read_ref_ld
 logger = logging.getLogger('gsMap.spatial_ldsc')
 
 
-# %%
-def _coef_new(jknife):
+def _coef_new(jknife, Nbar):
+    """Calculate coefficients adjusted by Nbar."""
     est_ = jknife.jknife_est[0, 0] / Nbar
     se_ = jknife.jknife_se[0, 0] / Nbar
     return est_, se_
 
 
 def append_intercept(x):
+    """Append an intercept term to the design matrix."""
     n_row = x.shape[0]
     intercept = np.ones((n_row, 1))
     x_new = np.concatenate((x, intercept), axis=1)
@@ -33,6 +35,7 @@ def append_intercept(x):
 
 
 def filter_sumstats_by_chisq(sumstats, chisq_max):
+    """Filter summary statistics based on chi-squared threshold."""
     before_len = len(sumstats)
     if chisq_max is None:
         chisq_max = max(0.001 * sumstats.N.max(), 80)
@@ -48,12 +51,14 @@ def filter_sumstats_by_chisq(sumstats, chisq_max):
 
 
 def aggregate(y, x, N, M, intercept=1):
+    """Aggregate function used in weight calculation."""
     num = M * (np.mean(y) - intercept)
     denom = np.mean(np.multiply(x, N))
     return num / denom
 
 
 def weights(ld, w_ld, N, M, hsq, intercept=1):
+    """Calculate weights for regression."""
     M = float(M)
     hsq = np.clip(hsq, 0.0, 1.0)
     ld = np.maximum(ld, 1.0)
@@ -65,178 +70,132 @@ def weights(ld, w_ld, N, M, hsq, intercept=1):
     return w
 
 
-def jackknife_for_processmap(spot_id):
-    # calculate the initial weight for each spot
+def get_weight_optimized(sumstats, x_tot_precomputed, M_tot, w_ld, intercept=1):
+    """Optimized function to calculate initial weights."""
+    tot_agg = aggregate(sumstats.chisq, x_tot_precomputed, sumstats.N, M_tot, intercept)
+    initial_w = weights(x_tot_precomputed, w_ld.LD_weights.values, sumstats.N.values, M_tot, tot_agg, intercept)
+    initial_w = np.sqrt(initial_w)
+    return initial_w
+
+
+def jackknife_for_processmap(spot_id, spatial_annotation, ref_ld_baseline_column_sum, sumstats, baseline_annotation, w_ld_common_snp, Nbar, n_blocks):
+    """Perform jackknife resampling for a given spot."""
     spot_spatial_annotation = spatial_annotation[:, spot_id]
     spot_x_tot_precomputed = spot_spatial_annotation + ref_ld_baseline_column_sum
-    initial_w = (
-        get_weight_optimized(sumstats, x_tot_precomputed=spot_x_tot_precomputed,
-                             M_tot=10000, w_ld=w_ld_common_snp, intercept=1)
-        .astype(np.float32)
-        .reshape((-1, 1)))
-
-    # apply the weight to baseline annotation, spatial annotation and CHISQ
+    initial_w = get_weight_optimized(sumstats, x_tot_precomputed=spot_x_tot_precomputed,
+                                     M_tot=10000, w_ld=w_ld_common_snp, intercept=1).astype(np.float32).reshape((-1, 1))
     initial_w_scaled = initial_w / np.sum(initial_w)
     baseline_annotation_spot = baseline_annotation * initial_w_scaled
     spatial_annotation_spot = spot_spatial_annotation.reshape((-1, 1)) * initial_w_scaled
     CHISQ = sumstats.chisq.values.reshape((-1, 1))
     y = CHISQ * initial_w_scaled
-
-    # run the jackknife
-    x_focal = np.concatenate((spatial_annotation_spot,
-                              baseline_annotation_spot), axis=1)
+    x_focal = np.concatenate((spatial_annotation_spot, baseline_annotation_spot), axis=1)
     try:
         jknife = jk.LstsqJackknifeFast(x_focal, y, n_blocks)
-        # LinAlgError
     except np.linalg.LinAlgError as e:
         logger.warning(f'LinAlgError: {e}')
         return np.nan, np.nan
-    return _coef_new(jknife)
-
-
-# Updated function
-def get_weight_optimized(sumstats, x_tot_precomputed, M_tot, w_ld, intercept=1):
-    tot_agg = aggregate(sumstats.chisq, x_tot_precomputed, sumstats.N, M_tot, intercept)
-    initial_w = weights(x_tot_precomputed, w_ld.LD_weights.values, sumstats.N.values, M_tot, tot_agg, intercept)
-    initial_w = np.sqrt(initial_w)
-    return initial_w
+    return _coef_new(jknife, Nbar)
 
 
 def _preprocess_sumstats(trait_name, sumstat_file_path, baseline_and_w_ld_common_snp: pd.Index, chisq_max=None):
-    # Load the gwas summary statistics
+    """Preprocess summary statistics."""
     sumstats = _read_sumstats(fh=sumstat_file_path, alleles=False, dropna=False)
     sumstats.set_index('SNP', inplace=True)
     sumstats = sumstats.astype(np.float32)
     sumstats = filter_sumstats_by_chisq(sumstats, chisq_max)
-
-    # NB: The intersection order is essential for keeping the same order of SNPs by its BP location
     common_snp = baseline_and_w_ld_common_snp.intersection(sumstats.index)
     if len(common_snp) < 200000:
         logger.warning(f'WARNING: number of SNPs less than 200k; for {trait_name} this is almost always bad.')
-
     sumstats = sumstats.loc[common_snp]
-
-    # get the common index position of baseline_and_w_ld_common_snp for quick access
     sumstats['common_index_pos'] = pd.Index(baseline_and_w_ld_common_snp).get_indexer(sumstats.index)
     return sumstats
 
 
-def _get_sumstats_from_sumstats_dict(sumstats_config_dict: dict, baseline_and_w_ld_common_snp: pd.Index,
-                                     chisq_max=None):
-    # first validate if all sumstats file exists
+def _get_sumstats_with_common_snp_from_sumstats_dict(sumstats_config_dict: dict, baseline_and_w_ld_common_snp: pd.Index, chisq_max=None):
+    """Get summary statistics with common SNPs among all traits."""
     logger.info('Validating sumstats files...')
     for trait_name, sumstat_file_path in sumstats_config_dict.items():
         if not os.path.exists(sumstat_file_path):
             raise FileNotFoundError(f'{sumstat_file_path} not found')
-    # then load all sumstats
     sumstats_cleaned_dict = {}
     for trait_name, sumstat_file_path in sumstats_config_dict.items():
-        sumstats_cleaned_dict[trait_name] = _preprocess_sumstats(trait_name, sumstat_file_path,
-                                                                 baseline_and_w_ld_common_snp, chisq_max)
-    logger.info('cleaned sumstats loaded')
-    return sumstats_cleaned_dict
+        sumstats_cleaned_dict[trait_name] = _preprocess_sumstats(trait_name, sumstat_file_path, baseline_and_w_ld_common_snp, chisq_max)
+    common_snp_among_all_sumstats = None
+    for trait_name, sumstats in sumstats_cleaned_dict.items():
+        if common_snp_among_all_sumstats is None:
+            common_snp_among_all_sumstats = sumstats.index
+        else:
+            common_snp_among_all_sumstats = common_snp_among_all_sumstats.intersection(sumstats.index)
+    for trait_name, sumstats in sumstats_cleaned_dict.items():
+        sumstats_cleaned_dict[trait_name] = sumstats.loc[common_snp_among_all_sumstats]
+    logger.info(f'Common SNPs among all sumstats: {len(common_snp_among_all_sumstats)}')
+    return sumstats_cleaned_dict, common_snp_among_all_sumstats
 
 
 class S_LDSC_Boost_with_pre_calculate_SNP_Gene_weight_matrix:
+    """Class to handle pre-calculated SNP-Gene weight matrix for quick mode."""
     def __init__(self, config: SpatialLDSCConfig, common_snp_among_all_sumstats_pos):
         self.config = config
         mk_score = pd.read_feather(config.mkscore_feather_path).set_index('HUMAN_GENE_SYM')
         mk_score_genes = mk_score.index
-
         snp_gene_weight_adata = ad.read_h5ad(config.snp_gene_weight_adata_path)
         common_genes = mk_score_genes.intersection(snp_gene_weight_adata.var.index)
         common_snps = snp_gene_weight_adata.obs.index
-        # self.snp_gene_weight_adata = snp_gene_weight_adata[common_snp_among_all_sumstats:, common_genes.to_list()]
         self.snp_gene_weight_matrix = snp_gene_weight_adata[common_snp_among_all_sumstats_pos, common_genes.to_list()].X
         self.mk_score_common = mk_score.loc[common_genes]
-
-        # calculate the chunk number
         self.chunk_starts = list(range(0, self.mk_score_common.shape[1], self.config.spots_per_chunk_quick_mode))
 
     def fetch_ldscore_by_chunk(self, chunk_index):
+        """Fetch LD score by chunk."""
         chunk_start = self.chunk_starts[chunk_index]
-        mk_score_chunk = self.mk_score_common.iloc[:,
-                         chunk_start:chunk_start + self.config.spots_per_chunk_quick_mode]
-        ldscore_chunk = self.calculate_ldscore_use_SNP_Gene_weight_matrix_by_chunk(
-            mk_score_chunk,
-            drop_dummy_na=False,
-        )
-
+        mk_score_chunk = self.mk_score_common.iloc[:, chunk_start:chunk_start + self.config.spots_per_chunk_quick_mode]
+        ldscore_chunk = self.calculate_ldscore_use_SNP_Gene_weight_matrix_by_chunk(mk_score_chunk, drop_dummy_na=False)
         spots_name = self.mk_score_common.columns[chunk_start:chunk_start + self.config.spots_per_chunk_quick_mode]
         return ldscore_chunk, spots_name
 
-    def calculate_ldscore_use_SNP_Gene_weight_matrix_by_chunk(self,
-                                                              mk_score_chunk,
-                                                              drop_dummy_na=True,
-                                                              ):
-
+    def calculate_ldscore_use_SNP_Gene_weight_matrix_by_chunk(self, mk_score_chunk, drop_dummy_na=True):
+        """Calculate LD score using SNP-Gene weight matrix by chunk."""
         if drop_dummy_na:
             ldscore_chr_chunk = self.snp_gene_weight_matrix[:, :-1] @ mk_score_chunk
         else:
             ldscore_chr_chunk = self.snp_gene_weight_matrix @ mk_score_chunk
-
         return ldscore_chr_chunk
 
 
-def _get_sumstats_with_common_snp_from_sumstats_dict(sumstats_config_dict: dict, baseline_and_w_ld_common_snp: pd.Index,
-                                                     chisq_max=None):
-    # first validate if all sumstats file exists
-    logger.info('Validating sumstats files...')
-    for trait_name, sumstat_file_path in sumstats_config_dict.items():
-        if not os.path.exists(sumstat_file_path):
-            raise FileNotFoundError(f'{sumstat_file_path} not found')
-    # then load all sumstats
-    sumstats_cleaned_dict = {}
-    for trait_name, sumstat_file_path in sumstats_config_dict.items():
-        sumstats_cleaned_dict[trait_name] = _preprocess_sumstats(trait_name, sumstat_file_path,
-                                                                 baseline_and_w_ld_common_snp, chisq_max)
-    # get the common snps among all sumstats
-    common_snp_among_all_sumstats = None
-    for trait_name, sumstats in sumstats_cleaned_dict.items():
-        if common_snp_among_all_sumstats is None:
-            common_snp_among_all_sumstats = sumstats.index
-        else:
-            common_snp_among_all_sumstats = common_snp_among_all_sumstats.intersection(sumstats.index)
-
-    # filter the common snps among all sumstats
-    for trait_name, sumstats in sumstats_cleaned_dict.items():
-        sumstats_cleaned_dict[trait_name] = sumstats.loc[common_snp_among_all_sumstats]
-
-    logger.info(f'Common SNPs among all sumstats: {len(common_snp_among_all_sumstats)}')
-    return sumstats_cleaned_dict, common_snp_among_all_sumstats
+def load_ldscore_chunk_from_feather(chunk_index, common_snp_among_all_sumstats_pos, config):
+    """Load LD score chunk from feather format."""
+    sample_name = config.sample_name
+    ld_file_spatial = f'{config.ldscore_save_dir}/{sample_name}_chunk{chunk_index}/{sample_name}.'
+    ref_ld_spatial = _read_ref_ld_v2(ld_file_spatial)
+    ref_ld_spatial = ref_ld_spatial.iloc[common_snp_among_all_sumstats_pos]
+    ref_ld_spatial = ref_ld_spatial.astype(np.float32, copy=False)
+    spatial_annotation_cnames = ref_ld_spatial.columns
+    return ref_ld_spatial.values, spatial_annotation_cnames
 
 
 def run_spatial_ldsc(config: SpatialLDSCConfig):
-    global spatial_annotation, baseline_annotation, n_blocks, Nbar, sumstats, ref_ld_baseline_column_sum, w_ld_common_snp
-    # config
+    """Run spatial LDSC analysis."""
+    logger.info(f'------Running Spatial LDSC for {config.sample_name}...')
     n_blocks = config.n_blocks
     sample_name = config.sample_name
 
-    print(f'------Running Spatial LDSC for {sample_name}...')
-    # Load the regression weights
+    # Load regression weights
     w_ld = _read_w_ld(config.w_file)
     w_ld_cname = w_ld.columns[1]
     w_ld.set_index('SNP', inplace=True)
 
     ld_file_baseline = f'{config.ldscore_save_dir}/baseline/baseline.'
-
     ref_ld_baseline = _read_ref_ld_v2(ld_file_baseline)
-    # n_annot_baseline = len(ref_ld_baseline.columns)
-    # M_annot_baseline = _read_M_v2(ld_file_baseline, n_annot_baseline, config.not_M_5_50)
-
-    # common snp between baseline and w_ld
     baseline_and_w_ld_common_snp = ref_ld_baseline.index.intersection(w_ld.index)
     baseline_and_w_ld_common_snp_pos = pd.Index(ref_ld_baseline.index).get_indexer(baseline_and_w_ld_common_snp)
 
-    # Clean the sumstats
     sumstats_cleaned_dict, common_snp_among_all_sumstats = _get_sumstats_with_common_snp_from_sumstats_dict(
-        config.sumstats_config_dict, baseline_and_w_ld_common_snp,
-        chisq_max=config.chisq_max)
+        config.sumstats_config_dict, baseline_and_w_ld_common_snp, chisq_max=config.chisq_max)
     common_snp_among_all_sumstats_pos = ref_ld_baseline.index.get_indexer(common_snp_among_all_sumstats)
 
-    # insure the order is monotonic
-    assert pd.Series(
-        common_snp_among_all_sumstats_pos).is_monotonic_increasing, 'common_snp_among_all_sumstats_pos is not monotonic increasing'
+    if not pd.Series(common_snp_among_all_sumstats_pos).is_monotonic_increasing:
+        raise ValueError('common_snp_among_all_sumstats_pos is not monotonic increasing')
 
     if len(common_snp_among_all_sumstats) < 200000:
         logger.warning(
@@ -245,136 +204,157 @@ def run_spatial_ldsc(config: SpatialLDSCConfig):
     ref_ld_baseline = ref_ld_baseline.loc[common_snp_among_all_sumstats]
     w_ld = w_ld.loc[common_snp_among_all_sumstats]
 
-    # load additional baseline annotations
+    # Load additional baseline annotations if needed
     if config.use_additional_baseline_annotation:
-        print('Using additional baseline annotations')
+        logger.info('Using additional baseline annotations')
         ld_file_baseline_additional = f'{config.ldscore_save_dir}/additional_baseline/baseline.'
         ref_ld_baseline_additional = _read_ref_ld_v2(ld_file_baseline_additional)
-        n_annot_baseline_additional = len(ref_ld_baseline_additional.columns)
-        logger.info(f'{len(ref_ld_baseline_additional.columns)} additional baseline annotations loaded')
-        # M_annot_baseline_additional = _read_M_v2(ld_file_baseline_additional, n_annot_baseline_additional,
-        #                                             config.not_M_5_50)
         ref_ld_baseline_additional = ref_ld_baseline_additional.loc[common_snp_among_all_sumstats]
         ref_ld_baseline = pd.concat([ref_ld_baseline, ref_ld_baseline_additional], axis=1)
         del ref_ld_baseline_additional
 
-    # Detect available chunk files
+    # Initialize s_ldsc once if quick_mode
+    s_ldsc = None
     if config.ldscore_save_format == 'quick_mode':
         s_ldsc = S_LDSC_Boost_with_pre_calculate_SNP_Gene_weight_matrix(config, common_snp_among_all_sumstats_pos)
         total_chunk_number_found = len(s_ldsc.chunk_starts)
-        print(f'Split data into {total_chunk_number_found} chunks')
+        logger.info(f'Split data into {total_chunk_number_found} chunks')
     else:
-        all_file = os.listdir(config.ldscore_save_dir)
-        total_chunk_number_found = sum('chunk' in name for name in all_file)
-        print(f'Find {total_chunk_number_found} chunked files in {config.ldscore_save_dir}')
-
-    if config.all_chunk is None:
-        if config.chunk_range is not None:
-            assert config.chunk_range[0] >= 1 and config.chunk_range[
-                1] <= total_chunk_number_found, 'Chunk range out of bound. It should be in [1, all_chunk]'
-            print(
-                f'chunk range provided, using chunked files from {config.chunk_range[0]} to {config.chunk_range[1]}')
-            start_chunk, end_chunk = config.chunk_range
-        else:
-            start_chunk, end_chunk = 1, total_chunk_number_found
-    else:
-        all_chunk = config.all_chunk
-        print(f'using {all_chunk} chunked files by provided argument')
-        print(f'\t')
-        print(f'Input {all_chunk} chunked files')
-        start_chunk, end_chunk = 1, all_chunk
+        total_chunk_number_found = determine_total_chunks(config)
 
+    start_chunk, end_chunk = determine_chunk_range(config, total_chunk_number_found)
     running_chunk_number = end_chunk - start_chunk + 1
 
-    # Process each chunk
-    output_dict = defaultdict(list)
-    zarr_path = Path(config.ldscore_save_dir) / f'{config.sample_name}.ldscore.zarr'
+    # Load zarr file if needed
+    zarr_file, spots_name = None, None
     if config.ldscore_save_format == 'zarr':
-        assert zarr_path.exists(), f'{zarr_path} not found, which is required for zarr format'
+        zarr_path = Path(config.ldscore_save_dir) / f'{config.sample_name}.ldscore.zarr'
+        if not zarr_path.exists():
+            raise FileNotFoundError(f'{zarr_path} not found, which is required for zarr format')
         zarr_file = zarr.open(str(zarr_path))
         spots_name = zarr_file.attrs['spot_names']
 
+    output_dict = defaultdict(list)
     for chunk_index in range(start_chunk, end_chunk + 1):
-        if config.ldscore_save_format == 'feather':
-            ref_ld_spatial, spatial_annotation_cnames = load_ldscore_chunk_from_feather(chunk_index,
-                                                                                        common_snp_among_all_sumstats_pos,
-                                                                                        config,
-                                                                                        )
-        elif config.ldscore_save_format == 'zarr':
-            ref_ld_spatial = zarr_file.blocks[:, chunk_index - 1][common_snp_among_all_sumstats_pos]
-            start_spot = (chunk_index - 1) * zarr_file.chunks[1]
-            ref_ld_spatial = ref_ld_spatial.astype(np.float32, copy=False)
-            spatial_annotation_cnames = spots_name[start_spot:start_spot + zarr_file.chunks[1]]
-        elif config.ldscore_save_format == 'quick_mode':
-            ref_ld_spatial, spatial_annotation_cnames = s_ldsc.fetch_ldscore_by_chunk(chunk_index - 1)
-        else:
-            raise ValueError(f'Invalid ld score save format: {config.ldscore_save_format}')
-
-        # get the x_tot_precomputed matrix by adding baseline and spatial annotation
+        ref_ld_spatial, spatial_annotation_cnames = load_ldscore_chunk(
+            chunk_index,
+            common_snp_among_all_sumstats_pos,
+            config,
+            zarr_file,
+            spots_name,
+            s_ldsc  # Pass s_ldsc to the function
+        )
         ref_ld_baseline_column_sum = ref_ld_baseline.sum(axis=1).values
-        # x_tot_precomputed = ref_ld_spatial + ref_ld_baseline_column_sum
 
         for trait_name, sumstats in sumstats_cleaned_dict.items():
-
             spatial_annotation = ref_ld_spatial.astype(np.float32, copy=False)
             baseline_annotation = ref_ld_baseline.copy().astype(np.float32, copy=False)
             w_ld_common_snp = w_ld.astype(np.float32, copy=False)
 
-            # weight the baseline annotation by N
             baseline_annotation = baseline_annotation * sumstats.N.values.reshape((-1, 1)) / sumstats.N.mean()
-            # append intercept
             baseline_annotation = append_intercept(baseline_annotation)
 
-            # Run the jackknife
             Nbar = sumstats.N.mean()
             chunk_size = spatial_annotation.shape[1]
-            out_chunk = thread_map(jackknife_for_processmap, range(chunk_size),
-                                   max_workers=config.num_processes,
-                                   chunksize=10,
-                                   desc=f'Chunk-{chunk_index}/Total-chunk-{running_chunk_number} for {trait_name}',
-                                   )
-
-            # cache the results
-            out_chunk = pd.DataFrame.from_records(out_chunk,
-                                                  columns=['beta', 'se', ],
-                                                  index=spatial_annotation_cnames)
-            # get the spots with nan
+
+            jackknife_func = partial(
+                jackknife_for_processmap,
+                spatial_annotation=spatial_annotation,
+                ref_ld_baseline_column_sum=ref_ld_baseline_column_sum,
+                sumstats=sumstats,
+                baseline_annotation=baseline_annotation,
+                w_ld_common_snp=w_ld_common_snp,
+                Nbar=Nbar,
+                n_blocks=n_blocks
+            )
+
+            out_chunk = thread_map(
+                jackknife_func,
+                range(chunk_size),
+                max_workers=config.num_processes,
+                chunksize=10,
+                desc=f'Chunk-{chunk_index}/Total-chunk-{running_chunk_number} for {trait_name}',
+            )
+
+            out_chunk = pd.DataFrame.from_records(out_chunk, columns=['beta', 'se'], index=spatial_annotation_cnames)
             nan_spots = out_chunk[out_chunk.isna().any(axis=1)].index
             if len(nan_spots) > 0:
                 logger.info(f'Nan spots: {nan_spots} in chunk-{chunk_index} for {trait_name}. They are removed.')
-            # drop the nan
             out_chunk = out_chunk.dropna()
-
             out_chunk['z'] = out_chunk.beta / out_chunk.se
             out_chunk['p'] = norm.sf(out_chunk['z'])
             output_dict[trait_name].append(out_chunk)
 
-        del ref_ld_spatial, spatial_annotation, baseline_annotation, w_ld_common_snp
-        gc.collect()
+            del spatial_annotation, baseline_annotation, w_ld_common_snp
+            gc.collect()
+
+    save_results(output_dict, config, running_chunk_number, start_chunk, end_chunk)
+    logger.info(f'------Spatial LDSC for {sample_name} finished!')
+
 
-    # Save the results
+def determine_total_chunks(config):
+    """Determine total number of chunks based on the ldscore save format."""
+    if config.ldscore_save_format == 'quick_mode':
+        s_ldsc = S_LDSC_Boost_with_pre_calculate_SNP_Gene_weight_matrix(config, [])
+        total_chunk_number_found = len(s_ldsc.chunk_starts)
+        logger.info(f'Split data into {total_chunk_number_found} chunks')
+    else:
+        all_file = os.listdir(config.ldscore_save_dir)
+        total_chunk_number_found = sum('chunk' in name for name in all_file)
+        logger.info(f'Find {total_chunk_number_found} chunked files in {config.ldscore_save_dir}')
+    return total_chunk_number_found
+
+
+def determine_chunk_range(config, total_chunk_number_found):
+    """Determine the range of chunks to process."""
+    if config.all_chunk is None:
+        if config.chunk_range is not None:
+            if not (1 <= config.chunk_range[0] <= total_chunk_number_found) or not (1 <= config.chunk_range[1] <= total_chunk_number_found):
+                raise ValueError('Chunk range out of bound. It should be in [1, all_chunk]')
+            start_chunk, end_chunk = config.chunk_range
+            logger.info(f'Chunk range provided, using chunked files from {start_chunk} to {end_chunk}')
+        else:
+            start_chunk, end_chunk = 1, total_chunk_number_found
+    else:
+        all_chunk = config.all_chunk
+        logger.info(f'Using {all_chunk} chunked files by provided argument')
+        start_chunk, end_chunk = 1, all_chunk
+    return start_chunk, end_chunk
+
+
+def load_ldscore_chunk(chunk_index, common_snp_among_all_sumstats_pos, config, zarr_file=None, spots_name=None, s_ldsc=None):
+    """Load LD score chunk based on save format."""
+    if config.ldscore_save_format == 'feather':
+        return load_ldscore_chunk_from_feather(chunk_index, common_snp_among_all_sumstats_pos, config)
+    elif config.ldscore_save_format == 'zarr':
+        ref_ld_spatial = zarr_file.blocks[:, chunk_index - 1][common_snp_among_all_sumstats_pos]
+        start_spot = (chunk_index - 1) * zarr_file.chunks[1]
+        ref_ld_spatial = ref_ld_spatial.astype(np.float32, copy=False)
+        spatial_annotation_cnames = spots_name[start_spot:start_spot + zarr_file.chunks[1]]
+        return ref_ld_spatial, spatial_annotation_cnames
+    elif config.ldscore_save_format == 'quick_mode':
+        # Use the pre-initialized s_ldsc
+        if s_ldsc is None:
+            raise ValueError("s_ldsc must be provided in quick_mode")
+        return s_ldsc.fetch_ldscore_by_chunk(chunk_index - 1)
+    else:
+        raise ValueError(f'Invalid ld score save format: {config.ldscore_save_format}')
+
+
+def save_results(output_dict, config, running_chunk_number, start_chunk, end_chunk):
+    """Save the results to the specified directory."""
     out_dir = config.ldsc_save_dir
     for trait_name, out_chunk_list in output_dict.items():
         out_all = pd.concat(out_chunk_list, axis=0)
-        if running_chunk_number == total_chunk_number_found:
+        sample_name = config.sample_name
+        if running_chunk_number == end_chunk - start_chunk + 1:
             out_file_name = out_dir / f'{sample_name}_{trait_name}.csv.gz'
         else:
             out_file_name = out_dir / f'{sample_name}_{trait_name}_chunk{start_chunk}-{end_chunk}.csv.gz'
         out_all['spot'] = out_all.index
         out_all = out_all[['spot', 'beta', 'se', 'z', 'p']]
-        out_all.to_csv(out_file_name, compression='gzip', index=False)
-
-        logger.info(f'Output saved to {out_file_name} for {trait_name}')
-    logger.info(f'------Spatial LDSC for {sample_name} finished!')
-
-
-def load_ldscore_chunk_from_feather(chunk_index, common_snp_among_all_sumstats_pos, config, ):
-    # Load the spatial annotations for this chunk
-    sample_name = config.sample_name
-    ld_file_spatial = f'{config.ldscore_save_dir}/{sample_name}_chunk{chunk_index}/{sample_name}.'
-    ref_ld_spatial = _read_ref_ld_v2(ld_file_spatial)
-    ref_ld_spatial = ref_ld_spatial.iloc[common_snp_among_all_sumstats_pos]
-    ref_ld_spatial = ref_ld_spatial.astype(np.float32, copy=False)
 
-    spatial_annotation_cnames = ref_ld_spatial.columns
-    return ref_ld_spatial.values, spatial_annotation_cnames
+        # clip the p-values
+        out_all['p'] = out_all['p'].clip(1e-300, 1)
+        out_all.to_csv(out_file_name, compression='gzip', index=False)
+        logger.info(f'Output saved to {out_file_name} for {trait_name}')

{gsmap-1.71.dist-info → gsmap-1.71.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: gsMap
-Version: 1.71
+Version: 1.71.2
 Summary: Genetics-informed pathogenic spatial mapping
 Author-email: liyang <songliyang@westlake.edu.cn>, wenhao <chenwenhao@westlake.edu.cn>
 Requires-Python: >=3.8
@@ -11,7 +11,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Operating System :: POSIX :: Linux
-Requires-Dist: numpy
+Requires-Dist: numpy < 2.0.0
 Requires-Dist: pandas
 Requires-Dist: scipy
 Requires-Dist: scikit-learn
@@ -28,7 +28,7 @@ Requires-Dist: kaleido
 Requires-Dist: jinja2
 Requires-Dist: scanpy >=1.8.0
 Requires-Dist: zarr
-Requires-Dist: bitarray
+Requires-Dist: bitarray >=2.9.2, <3.0.0
 Requires-Dist: pyarrow
 Requires-Dist: scikit-misc
 Requires-Dist: sphinx ; extra == "doc"

{gsmap-1.71.dist-info → gsmap-1.71.2.dist-info}/RECORD

@@ -1,21 +1,21 @@
-gsMap/__init__.py,sha256=xwuyVm3wc8MEgSfC3D0LNmecSOC2bkVXYWf0cpwX3I0,74
+gsMap/__init__.py,sha256=TKJMAFwuDpLdq1rs6HhXdBwKqbdPXrH9doabIfnGCRg,76
 gsMap/__main__.py,sha256=Ih3PPSJezd9UpmkdfH1PXYYecCpk6hkQF039BQU1ocs,60
 gsMap/cauchy_combination_test.py,sha256=jMXNKWajI5NSmKeKu5PVJMYMtiApL6ZeYrSOSqC-ESU,5553
-gsMap/config.py,sha256=2U_KNlrENMh6XFRJWn88CCOe2unctWYehyUshHt9LU0,40362
-gsMap/diagnosis.py,sha256=SUbEXqEPfziSrbytxNnH7hiNypocyCOAYPeUYDwgPHw,12947
-gsMap/find_latent_representation.py,sha256=h3douGRcrkyNrWEieNnwzc4dXacySx_VMAPZrn2wKyg,4460
+gsMap/config.py,sha256=eddRqfld9vgkjrXII2wKVbzAKJuiZzxvp6e4HqYb0cw,40386
+gsMap/diagnosis.py,sha256=MMOrKn_TV1fWeSbgXQ2gar9uYHelJUzzEK4WTId5ei0,13117
+gsMap/find_latent_representation.py,sha256=ZarK-rjjBu3UUJqz0GxYsxcxE5A6RdZdeSC5bSP4x48,4807
 gsMap/format_sumstats.py,sha256=vPJD5qdqd_CjHw5MkihYpSMLcDx7wsNySsmoXRElaW0,13439
 gsMap/generate_ldscore.py,sha256=s0evdoGJ2QKZYMxHGm_z01yf7yE6RWa1iUOJzSh-OZY,28471
-gsMap/latent_to_gene.py,sha256=W7icT8o7GjUwrnHyNgdPFTCr03T6w8fB8vbApg1i-So,9452
+gsMap/latent_to_gene.py,sha256=HtO3RSHQ_yeicqGC8iJt65OIq_akoKJhUfIdfb3k3kk,10622
 gsMap/main.py,sha256=b3pLWvc70UlO24QXW9Zl7Nd7rlErxlQSM5Y_RqwcNNM,1254
 gsMap/report.py,sha256=pr7K8zEyunun5LHzRtLBaoLjsw8-PfdbJWKMvsTxoj0,6803
 gsMap/run_all_mode.py,sha256=fXO5QNZbkQpQjkZTznqq5Ct1OihIRiFb4iwGHtnEryA,8194
 gsMap/setup.py,sha256=eOoPalGAHTY06_7a35nvbOaKQhq0SBE5GK4pc4bp3wc,102
-gsMap/spatial_ldsc_multiple_sumstats.py,sha256=Pc4MtsLilZZVtCVhQWaxePEyVj6tXva1VWQ2oRZOFRA,18011
+gsMap/spatial_ldsc_multiple_sumstats.py,sha256=YrqTWC7FHmimGhhh-aoxIUzwouPFI02fbimBByu3gVw,17431
 gsMap/visualize.py,sha256=pLp8iyh0L7UU405LfpjEgbGuKbf19_4YyMHZWo6wE7M,7255
 gsMap/GNN/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 gsMap/GNN/adjacency_matrix.py,sha256=nlzIW4b6JSf77SCsjSk75eSEGiwBZLEQQ0D0mpIy4go,3258
-gsMap/GNN/model.py,sha256=Lbf1QBSQByYVdBJFJhh6H0-GeFxEwlXtEkRpqQdhSsM,2870
+gsMap/GNN/model.py,sha256=W4jPnaSDUUEacvGH5sNACNlbbO79-6hwLtSudyXIYgM,2909
 gsMap/GNN/train.py,sha256=s9F3_eeOlTzqOra2jHvtV3SDp6fKkt_AVTnssNueruQ,3094
 gsMap/templates/report_template.html,sha256=QODZEbVxpW1xsLz7lDrD_DyUfzYoi9E17o2tLJlf8OQ,8016
 gsMap/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -24,8 +24,8 @@ gsMap/utils/jackknife.py,sha256=d3OVLAoU8W0MYdZBT4-8GPAdZDN9aNHnDHwSh4AkxP8,1773
 gsMap/utils/make_annotations.py,sha256=X9vxrH8A0oqrGJ-C82H3yactMDLpSM-mYA_8DurxPzE,21679
 gsMap/utils/manhattan_plot.py,sha256=vS8Avgn9kbYa6HerRKpRNPYgJyyh5mE2f9p1BuHwwS8,25596
 gsMap/utils/regression_read.py,sha256=oivyNicRdLLQvBv36GoR0AweUlvuCmBduTVRPpGhNRU,8474
-gsmap-1.71.dist-info/entry_points.txt,sha256=s_P2Za22O077tc1FPLKMinbdRVXaN_HTcDBgWMYpqA4,41
-gsmap-1.71.dist-info/LICENSE,sha256=dCWx-ENipnYph2UTEA9wJaEZ_tkjNZ_tog6XRd3nd2k,1073
-gsmap-1.71.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-gsmap-1.71.dist-info/METADATA,sha256=JqVynQWFd04_vu8qXlL7QDtjwGeEh7YLNGcBVNWOjTY,3691
-gsmap-1.71.dist-info/RECORD,,
+gsmap-1.71.2.dist-info/entry_points.txt,sha256=s_P2Za22O077tc1FPLKMinbdRVXaN_HTcDBgWMYpqA4,41
+gsmap-1.71.2.dist-info/LICENSE,sha256=dCWx-ENipnYph2UTEA9wJaEZ_tkjNZ_tog6XRd3nd2k,1073
+gsmap-1.71.2.dist-info/WHEEL,sha256=-ta_u8-23-bwm4JkHhR_rJkpb_X1cxtvIsLH0KOv1sU,82
+gsmap-1.71.2.dist-info/METADATA,sha256=sObxX5A_Pt26uqYQ8G4gpkrthY81qvkC4p1pon4j-sM,3717
+gsmap-1.71.2.dist-info/RECORD,,

{gsmap-1.71.dist-info → gsmap-1.71.2.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: flit 3.9.0
+Generator: flit 3.10.0
 Root-Is-Purelib: true
 Tag: py3-none-any