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

gsMap/generate_ldscore.py

@@ -10,7 +10,7 @@ from scipy.sparse import csr_matrix
 from tqdm import trange
 
 from gsMap.config import GenerateLDScoreConfig
-from gsMap.utils.generate_r2_matrix import PlinkBEDFileWithR2Cache, getBlockLefts, ID_List_Factory
+from gsMap.utils.generate_r2_matrix import getBlockLefts, load_bfile
 
 warnings.filterwarnings("ignore", category=FutureWarning)
 logger = logging.getLogger(__name__)
@@ -25,34 +25,36 @@ def load_gtf(gtf_file, mk_score, window_size):
     print("Loading gtf data")
     #
     # Load GTF file
-    gtf = pr.read_gtf(gtf_file, )
+    gtf = pr.read_gtf(
+        gtf_file,
+    )
     gtf = gtf.df
     #
     # Select the common genes
-    gtf = gtf[gtf['Feature'] == 'gene']
+    gtf = gtf[gtf["Feature"] == "gene"]
     common_gene = np.intersect1d(mk_score.index, gtf.gene_name)
     #
     gtf = gtf[gtf.gene_name.isin(common_gene)]
     mk_score = mk_score[mk_score.index.isin(common_gene)]
     #
     # Remove duplicated lines
-    gtf = gtf.drop_duplicates(subset='gene_name', keep="first")
+    gtf = gtf.drop_duplicates(subset="gene_name", keep="first")
     #
     # Process the GTF (open 100-KB window: Tss - Ted)
-    gtf_bed = gtf[['Chromosome', 'Start', 'End', 'gene_name', 'Strand']].copy()
-    gtf_bed.loc[:, 'TSS'] = gtf_bed['Start']
-    gtf_bed.loc[:, 'TED'] = gtf_bed['End']
+    gtf_bed = gtf[["Chromosome", "Start", "End", "gene_name", "Strand"]].copy()
+    gtf_bed.loc[:, "TSS"] = gtf_bed["Start"]
+    gtf_bed.loc[:, "TED"] = gtf_bed["End"]
 
-    gtf_bed.loc[:, 'Start'] = gtf_bed['TSS'] - window_size
-    gtf_bed.loc[:, 'End'] = gtf_bed['TED'] + window_size
-    gtf_bed.loc[gtf_bed['Start'] < 0, 'Start'] = 0
+    gtf_bed.loc[:, "Start"] = gtf_bed["TSS"] - window_size
+    gtf_bed.loc[:, "End"] = gtf_bed["TED"] + window_size
+    gtf_bed.loc[gtf_bed["Start"] < 0, "Start"] = 0
     #
     # Correct the negative strand
-    tss_neg = gtf_bed.loc[gtf_bed['Strand'] == '-', 'TSS']
-    ted_neg = gtf_bed.loc[gtf_bed['Strand'] == '-', 'TED']
-    gtf_bed.loc[gtf_bed['Strand'] == '-', 'TSS'] = ted_neg
-    gtf_bed.loc[gtf_bed['Strand'] == '-', 'TED'] = tss_neg
-    gtf_bed = gtf_bed.drop('Strand', axis=1)
+    tss_neg = gtf_bed.loc[gtf_bed["Strand"] == "-", "TSS"]
+    ted_neg = gtf_bed.loc[gtf_bed["Strand"] == "-", "TED"]
+    gtf_bed.loc[gtf_bed["Strand"] == "-", "TSS"] = ted_neg
+    gtf_bed.loc[gtf_bed["Strand"] == "-", "TED"] = tss_neg
+    gtf_bed = gtf_bed.drop("Strand", axis=1)
     #
     # Transform the GTF to PyRanges
     gtf_pr = pr.PyRanges(gtf_bed)
@@ -64,31 +66,28 @@ def load_marker_score(mk_score_file):
     """
     Load marker scores of each cell.
     """
-    mk_score = pd.read_feather(mk_score_file).set_index('HUMAN_GENE_SYM').rename_axis('gene_name')
+    mk_score = pd.read_feather(mk_score_file).set_index("HUMAN_GENE_SYM").rename_axis("gene_name")
     mk_score = mk_score.astype(np.float32, copy=False)
     return mk_score
 
 
-# %%
-# load mkscore get common gene
-# %%
 # load bim
 def load_bim(bfile_root, chrom):
     """
     Load the bim file.
     """
-    bim = pd.read_csv(f'{bfile_root}.{chrom}.bim', sep='\t', header=None)
+    bim = pd.read_csv(f"{bfile_root}.{chrom}.bim", sep="\t", header=None)
     bim.columns = ["CHR", "SNP", "CM", "BP", "A1", "A2"]
     #
     # Transform bim to PyRanges
     bim_pr = bim.copy()
     bim_pr.columns = ["Chromosome", "SNP", "CM", "Start", "A1", "A2"]
 
-    bim_pr['End'] = bim_pr['Start'].copy()
-    bim_pr['Start'] = bim_pr['Start'] - 1  # Due to bim file is 1-based
+    bim_pr["End"] = bim_pr["Start"].copy()
+    bim_pr["Start"] = bim_pr["Start"] - 1  # Due to bim file is 1-based
 
     bim_pr = pr.PyRanges(bim_pr)
-    bim_pr.Chromosome = f'chr{chrom}'
+    bim_pr.Chromosome = f"chr{chrom}"
     return bim, bim_pr
 
 
@@ -100,9 +99,9 @@ def Overlaps_gtf_bim(gtf_pr, bim_pr):
     # Select the overlapped regions (SNPs in gene windows)
     overlaps = gtf_pr.join(bim_pr)
     overlaps = overlaps.df
-    overlaps['Distance'] = np.abs(overlaps['Start_b'] - overlaps['TSS'])
+    overlaps["Distance"] = np.abs(overlaps["Start_b"] - overlaps["TSS"])
     overlaps_small = overlaps.copy()
-    overlaps_small = overlaps_small.loc[overlaps_small.groupby('SNP').Distance.idxmin()]
+    overlaps_small = overlaps_small.loc[overlaps_small.groupby("SNP").Distance.idxmin()]
     return overlaps_small
 
 
@@ -110,7 +109,7 @@ def Overlaps_gtf_bim(gtf_pr, bim_pr):
 def filter_snps_by_keep_snp(bim_df, keep_snp_file):
     # Load the keep_snp file and filter the BIM DataFrame
     keep_snp = pd.read_csv(keep_snp_file, header=None)[0].to_list()
-    filtered_bim_df = bim_df[bim_df['SNP'].isin(keep_snp)]
+    filtered_bim_df = bim_df[bim_df["SNP"].isin(keep_snp)]
     return filtered_bim_df
 
 
@@ -122,7 +121,7 @@ def get_snp_counts(config):
         bim_df, _ = load_bim(config.bfile_root, chrom)
 
         if config.keep_snp_root:
-            keep_snp_file = f'{config.keep_snp_root}.{chrom}.snp'
+            keep_snp_file = f"{config.keep_snp_root}.{chrom}.snp"
             filtered_bim_df = filter_snps_by_keep_snp(bim_df, keep_snp_file)
         else:
             filtered_bim_df = bim_df
@@ -130,11 +129,11 @@ def get_snp_counts(config):
         snp_counts[chrom] = filtered_bim_df.shape[0]
         total_snp += snp_counts[chrom]
 
-    snp_counts['total'] = total_snp
+    snp_counts["total"] = total_snp
 
     chrom_snp_length_array = np.array([snp_counts[chrom] for chrom in range(1, 23)]).cumsum()
 
-    snp_counts['chrom_snp_start_point'] = [0] + chrom_snp_length_array.tolist()
+    snp_counts["chrom_snp_start_point"] = [0] + chrom_snp_length_array.tolist()
 
     return snp_counts
 
@@ -144,56 +143,35 @@ def get_snp_pass_maf(bfile_root, chrom, maf_min=0.05):
     """
     Get the dummy matrix of SNP-gene pairs.
     """
-    # Load the bim file
-    PlinkBIMFile = ID_List_Factory(['CHR', 'SNP', 'CM', 'BP', 'A1', 'A2'], 1, '.bim', usecols=[0, 1, 2, 3, 4, 5])
-    PlinkFAMFile = ID_List_Factory(['IID'], 0, '.fam', usecols=[1])
+    array_snps, array_indivs, geno_array = load_bfile(bfile_chr_prefix=f"{bfile_root}.{chrom}")
 
-    bfile = f'{bfile_root}.{chrom}'
-    snp_file, snp_obj = bfile + '.bim', PlinkBIMFile
-    array_snps = snp_obj(snp_file)
     m = len(array_snps.IDList)
-
-    # Load fam
-    ind_file, ind_obj = bfile + '.fam', PlinkFAMFile
-    array_indivs = ind_obj(ind_file)
     n = len(array_indivs.IDList)
-    array_file, array_obj = bfile + '.bed', PlinkBEDFileWithR2Cache
-    geno_array = array_obj(array_file, n, array_snps, keep_snps=None, keep_indivs=None, mafMin=None)
+    logger.info(
+        f"Loading genotype data for {m} SNPs and {n} individuals from {bfile_root}.{chrom}"
+    )
+
     ii = geno_array.maf > maf_min
     snp_pass_maf = array_snps.IDList[ii]
-    print(f'After filtering SNPs with MAF < {maf_min}, {len(snp_pass_maf)} SNPs remain.')
+    logger.info(f"After filtering SNPs with MAF < {maf_min}, {len(snp_pass_maf)} SNPs remain.")
     return snp_pass_maf.SNP.to_list()
 
 
-def get_ldscore(bfile_root, chrom, annot_matrix, ld_wind, ld_unit='CM'):
-    PlinkBIMFile = ID_List_Factory(['CHR', 'SNP', 'CM', 'BP', 'A1', 'A2'], 1, '.bim', usecols=[0, 1, 2, 3, 4, 5])
-    PlinkFAMFile = ID_List_Factory(['IID'], 0, '.fam', usecols=[1])
+def get_ldscore(bfile_root, chrom, annot_matrix, ld_wind, ld_unit="CM"):
+    array_snps, array_indivs, geno_array = load_bfile(bfile_chr_prefix=f"{bfile_root}.{chrom}")
 
-    bfile = f'{bfile_root}.{chrom}'
-    snp_file, snp_obj = bfile + '.bim', PlinkBIMFile
-    array_snps = snp_obj(snp_file)
-    m = len(array_snps.IDList)
-    print(f'Read list of {m} SNPs from {snp_file}')
-
-    # Load fam
-    ind_file, ind_obj = bfile + '.fam', PlinkFAMFile
-    array_indivs = ind_obj(ind_file)
-    n = len(array_indivs.IDList)
-    print(f'Read list of {n} individuals from {ind_file}')
-    array_file, array_obj = bfile + '.bed', PlinkBEDFileWithR2Cache
-    geno_array = array_obj(array_file, n, array_snps, keep_snps=None, keep_indivs=None, mafMin=None)
     # Load the annotations of the baseline
-    if ld_unit == 'SNP':
+    if ld_unit == "SNP":
         max_dist = ld_wind
         coords = np.array(range(geno_array.m))
-    elif ld_unit == 'KB':
+    elif ld_unit == "KB":
         max_dist = ld_wind * 1000
-        coords = np.array(array_snps.df['BP'])[geno_array.kept_snps]
-    elif ld_unit == 'CM':
+        coords = np.array(array_snps.df["BP"])[geno_array.kept_snps]
+    elif ld_unit == "CM":
         max_dist = ld_wind
-        coords = np.array(array_snps.df['CM'])[geno_array.kept_snps]
+        coords = np.array(array_snps.df["CM"])[geno_array.kept_snps]
     else:
-        raise ValueError(f'Invalid ld_wind_unit: {ld_unit}')
+        raise ValueError(f"Invalid ld_wind_unit: {ld_unit}")
     block_left = getBlockLefts(coords, max_dist)
     # Calculate the LD score
     lN_df = pd.DataFrame(geno_array.ldScoreVarBlocks(block_left, 100, annot=annot_matrix))
@@ -201,25 +179,31 @@ def get_ldscore(bfile_root, chrom, annot_matrix, ld_wind, ld_unit='CM'):
 
 
 # %%
-def calculate_ldscore_from_annotation(SNP_annotation_df, chrom, bfile_root, ld_wind=1, ld_unit='CM'):
+def calculate_ldscore_from_annotation(
+    SNP_annotation_df, chrom, bfile_root, ld_wind=1, ld_unit="CM"
+):
     """
     Calculate the SNP-gene weight matrix.
     """
     # Get the dummy matrix
     # Get the SNP-gene weight matrix
-    snp_gene_weight_matrix = get_ldscore(bfile_root, chrom, SNP_annotation_df.values, ld_wind=ld_wind,
-                                         ld_unit=ld_unit)
+    snp_gene_weight_matrix = get_ldscore(
+        bfile_root, chrom, SNP_annotation_df.values, ld_wind=ld_wind, ld_unit=ld_unit
+    )
     snp_gene_weight_matrix = snp_gene_weight_matrix.astype(np.float32, copy=False)
     snp_gene_weight_matrix.index = SNP_annotation_df.index
     snp_gene_weight_matrix.columns = SNP_annotation_df.columns
     return snp_gene_weight_matrix
 
 
-def calculate_ldscore_from_multiple_annotation(SNP_annotation_df_list, chrom, bfile_root, ld_wind=1, ld_unit='CM'):
+def calculate_ldscore_from_multiple_annotation(
+    SNP_annotation_df_list, chrom, bfile_root, ld_wind=1, ld_unit="CM"
+):
     SNP_annotation_df = pd.concat(SNP_annotation_df_list, axis=1).astype(np.float32, copy=False)
 
-    snp_gene_weight_matrix = get_ldscore(bfile_root, chrom, SNP_annotation_df.values, ld_wind=ld_wind,
-                                         ld_unit=ld_unit)
+    snp_gene_weight_matrix = get_ldscore(
+        bfile_root, chrom, SNP_annotation_df.values, ld_wind=ld_wind, ld_unit=ld_unit
+    )
     snp_gene_weight_matrix = snp_gene_weight_matrix.astype(np.float32, copy=False)
     snp_gene_weight_matrix.index = SNP_annotation_df.index
     snp_gene_weight_matrix.columns = SNP_annotation_df.columns
@@ -229,7 +213,9 @@ def calculate_ldscore_from_multiple_annotation(SNP_annotation_df_list, chrom, bf
     snp_gene_weight_matrix_list = []
     start = 0
     for snp_annotation_len in snp_annotation_len_list:
-        snp_gene_weight_matrix_list.append(snp_gene_weight_matrix.iloc[:, start:start + snp_annotation_len])
+        snp_gene_weight_matrix_list.append(
+            snp_gene_weight_matrix.iloc[:, start : start + snp_annotation_len]
+        )
         start += snp_annotation_len
     return snp_gene_weight_matrix_list
 
@@ -242,21 +228,28 @@ class S_LDSC_Boost:
         self.mk_score = load_marker_score(config.mkscore_feather_path)
 
         # Load GTF and get common markers
-        self.gtf_pr, self.mk_score_common = load_gtf(config.gtf_annotation_file, self.mk_score,
-                                                     window_size=config.gene_window_size)
+        self.gtf_pr, self.mk_score_common = load_gtf(
+            config.gtf_annotation_file, self.mk_score, window_size=config.gene_window_size
+        )
 
         # Load enhancer
         if config.enhancer_annotation_file is not None:
             enhancer_df = pr.read_bed(config.enhancer_annotation_file, as_df=True)
-            enhancer_df.set_index('Name', inplace=True)
-            enhancer_df.index.name = 'gene_name'
+            enhancer_df.set_index("Name", inplace=True)
+            enhancer_df.index.name = "gene_name"
 
             # keep the common genes and add the enhancer score
-            avg_mkscore = pd.DataFrame(self.mk_score_common.mean(axis=1), columns=['avg_mkscore'])
-            enhancer_df = enhancer_df.join(avg_mkscore, how='inner', on='gene_name', )
+            avg_mkscore = pd.DataFrame(self.mk_score_common.mean(axis=1), columns=["avg_mkscore"])
+            enhancer_df = enhancer_df.join(
+                avg_mkscore,
+                how="inner",
+                on="gene_name",
+            )
 
             # add distance to TSS
-            enhancer_df['TSS'] = self.gtf_pr.df.set_index('gene_name').reindex(enhancer_df.index)['TSS']
+            enhancer_df["TSS"] = self.gtf_pr.df.set_index("gene_name").reindex(enhancer_df.index)[
+                "TSS"
+            ]
 
             # convert to pyranges
             self.enhancer_pr = pr.PyRanges(enhancer_df.reset_index())
@@ -265,32 +258,39 @@ class S_LDSC_Boost:
             self.enhancer_pr = None
 
         # create tha zarr file
-        if config.ldscore_save_format == 'zarr':
-
+        if config.ldscore_save_format == "zarr":
             chrom_snp_length_dict = get_snp_counts(config)
-            self.chrom_snp_start_point = chrom_snp_length_dict['chrom_snp_start_point']
+            self.chrom_snp_start_point = chrom_snp_length_dict["chrom_snp_start_point"]
 
-            zarr_path = Path(config.ldscore_save_dir) / f'{config.sample_name}.ldscore.zarr'
+            zarr_path = Path(config.ldscore_save_dir) / f"{config.sample_name}.ldscore.zarr"
             if not zarr_path.exists():
-                self.zarr_file = zarr.open(zarr_path.as_posix(), mode='a', dtype=np.float16,
-                                           chunks=config.zarr_chunk_size,
-                                           shape=(chrom_snp_length_dict['total'], self.mk_score_common.shape[1]))
+                self.zarr_file = zarr.open(
+                    zarr_path.as_posix(),
+                    mode="a",
+                    dtype=np.float16,
+                    chunks=config.zarr_chunk_size,
+                    shape=(chrom_snp_length_dict["total"], self.mk_score_common.shape[1]),
+                )
                 zarr_path.mkdir(parents=True, exist_ok=True)
                 # save spot names
-                self.zarr_file.attrs['spot_names'] = self.mk_score_common.columns.to_list()
+                self.zarr_file.attrs["spot_names"] = self.mk_score_common.columns.to_list()
                 # save chrom_snp_length_dict
-                self.zarr_file.attrs['chrom_snp_start_point'] = self.chrom_snp_start_point
+                self.zarr_file.attrs["chrom_snp_start_point"] = self.chrom_snp_start_point
             else:
-                self.zarr_file = zarr.open(zarr_path.as_posix(), mode='a')
+                self.zarr_file = zarr.open(zarr_path.as_posix(), mode="a")
 
     def process_chromosome(self, chrom: int):
         self.snp_pass_maf = get_snp_pass_maf(self.config.bfile_root, chrom, maf_min=0.05)
 
         # Get SNP-Gene dummy pairs
-        self.snp_gene_pair_dummy = self.get_snp_gene_dummy(chrom, )
+        self.snp_gene_pair_dummy = self.get_snp_gene_dummy(
+            chrom,
+        )
 
         if self.config.keep_snp_root is not None:
-            keep_snp = pd.read_csv(f'{self.config.keep_snp_root}.{chrom}.snp', header=None)[0].to_list()
+            keep_snp = pd.read_csv(f"{self.config.keep_snp_root}.{chrom}.snp", header=None)[
+                0
+            ].to_list()
             self.keep_snp_mask = self.snp_gene_pair_dummy.index.isin(keep_snp)
             # the SNP name of keeped
             self.snp_name = self.snp_gene_pair_dummy.index[self.keep_snp_mask].to_list()
@@ -300,25 +300,37 @@ class S_LDSC_Boost:
 
         if self.config.additional_baseline_annotation is not None:
             additional_baseline_annotation = Path(self.config.additional_baseline_annotation)
-            additional_baseline_annotation_file_path = additional_baseline_annotation / f'baseline.{chrom}.annot.gz'
-            assert additional_baseline_annotation_file_path.exists(), f'additional_baseline_annotation_file_path not exists: {additional_baseline_annotation_file_path}'
-            additional_baseline_annotation_df = pd.read_csv(additional_baseline_annotation_file_path, sep='\t')
-            additional_baseline_annotation_df.set_index('SNP', inplace=True)
+            additional_baseline_annotation_file_path = (
+                additional_baseline_annotation / f"baseline.{chrom}.annot.gz"
+            )
+            assert additional_baseline_annotation_file_path.exists(), (
+                f"additional_baseline_annotation_file_path not exists: {additional_baseline_annotation_file_path}"
+            )
+            additional_baseline_annotation_df = pd.read_csv(
+                additional_baseline_annotation_file_path, sep="\t"
+            )
+            additional_baseline_annotation_df.set_index("SNP", inplace=True)
 
             # drop these columns if exists CHR         BP       CM]
-            additional_baseline_annotation_df.drop(['CHR', 'BP', 'CM'], axis=1, inplace=True, errors='ignore')
+            additional_baseline_annotation_df.drop(
+                ["CHR", "BP", "CM"], axis=1, inplace=True, errors="ignore"
+            )
 
             # reindex, for those SNPs not in additional_baseline_annotation_df, set to 0
-            num_of_not_exist_snp = (~self.snp_gene_pair_dummy.index.isin(additional_baseline_annotation_df.index)).sum()
+            num_of_not_exist_snp = (
+                ~self.snp_gene_pair_dummy.index.isin(additional_baseline_annotation_df.index)
+            ).sum()
             if num_of_not_exist_snp > 0:
                 logger.warning(
-                    f'{num_of_not_exist_snp} SNPs not in additional_baseline_annotation_df but in the reference panel, so the additional baseline annotation of these SNP will set to 0')
+                    f"{num_of_not_exist_snp} SNPs not in additional_baseline_annotation_df but in the reference panel, so the additional baseline annotation of these SNP will set to 0"
+                )
                 additional_baseline_annotation_df = additional_baseline_annotation_df.reindex(
-                    self.snp_gene_pair_dummy.index,
-                    fill_value=0)
+                    self.snp_gene_pair_dummy.index, fill_value=0
+                )
             else:
                 additional_baseline_annotation_df = additional_baseline_annotation_df.reindex(
-                    self.snp_gene_pair_dummy.index)
+                    self.snp_gene_pair_dummy.index
+                )
 
             # do this for saving the cpu time, only calculate r2 once
             self.snp_gene_weight_matrix, additional_baseline_annotation_ldscore = (
@@ -327,56 +339,85 @@ class S_LDSC_Boost:
                     chrom,
                     self.config.bfile_root,
                     ld_wind=self.config.ld_wind,
-                    ld_unit=self.config.ld_unit))
+                    ld_unit=self.config.ld_unit,
+                )
+            )
 
-            additional_baseline_annotation_ldscore = additional_baseline_annotation_ldscore.loc[self.snp_name]
+            additional_baseline_annotation_ldscore = additional_baseline_annotation_ldscore.loc[
+                self.snp_name
+            ]
             # print(additional_baseline_annotation_ldscore.index.to_list()==self.snp_name)
 
-            ld_score_file = f'{self.config.ldscore_save_dir}/additional_baseline/baseline.{chrom}.l2.ldscore.feather'
-            M_file_path = f'{self.config.ldscore_save_dir}/additional_baseline/baseline.{chrom}.l2.M'
-            M_5_file_path = f'{self.config.ldscore_save_dir}/additional_baseline/baseline.{chrom}.l2.M_5_50'
+            ld_score_file = f"{self.config.ldscore_save_dir}/additional_baseline/baseline.{chrom}.l2.ldscore.feather"
+            M_file_path = (
+                f"{self.config.ldscore_save_dir}/additional_baseline/baseline.{chrom}.l2.M"
+            )
+            M_5_file_path = (
+                f"{self.config.ldscore_save_dir}/additional_baseline/baseline.{chrom}.l2.M_5_50"
+            )
 
             # save additional baseline annotation ldscore
-            self.save_ldscore_to_feather(additional_baseline_annotation_ldscore.values,
-                                         column_names=additional_baseline_annotation_ldscore.columns,
-                                         save_file_name=ld_score_file,
-                                         )
+            self.save_ldscore_to_feather(
+                additional_baseline_annotation_ldscore.values,
+                column_names=additional_baseline_annotation_ldscore.columns,
+                save_file_name=ld_score_file,
+            )
 
             # caculate the M and save
             save_dir = Path(M_file_path).parent
             save_dir.mkdir(parents=True, exist_ok=True)
             M_chr_chunk = additional_baseline_annotation_df.values.sum(axis=0, keepdims=True)
-            M_5_chr_chunk = additional_baseline_annotation_df.loc[self.snp_pass_maf].values.sum(axis=0, keepdims=True)
-            np.savetxt(M_file_path, M_chr_chunk, delimiter='\t', )
-            np.savetxt(M_5_file_path, M_5_chr_chunk, delimiter='\t', )
+            M_5_chr_chunk = additional_baseline_annotation_df.loc[self.snp_pass_maf].values.sum(
+                axis=0, keepdims=True
+            )
+            np.savetxt(
+                M_file_path,
+                M_chr_chunk,
+                delimiter="\t",
+            )
+            np.savetxt(
+                M_5_file_path,
+                M_5_chr_chunk,
+                delimiter="\t",
+            )
 
         else:
             # Calculate SNP-Gene weight matrix
-            self.snp_gene_weight_matrix = calculate_ldscore_from_annotation(self.snp_gene_pair_dummy, chrom,
-                                                                            self.config.bfile_root,
-                                                                            ld_wind=self.config.ld_wind,
-                                                                            ld_unit=self.config.ld_unit)
+            self.snp_gene_weight_matrix = calculate_ldscore_from_annotation(
+                self.snp_gene_pair_dummy,
+                chrom,
+                self.config.bfile_root,
+                ld_wind=self.config.ld_wind,
+                ld_unit=self.config.ld_unit,
+            )
         # only keep the snp in keep_snp_root
         if self.keep_snp_mask is not None:
             self.snp_gene_weight_matrix = self.snp_gene_weight_matrix[self.keep_snp_mask]
 
         if self.config.save_pre_calculate_snp_gene_weight_matrix:
-            snp_gene_weight_matrix_save_dir = Path(self.config.ldscore_save_dir) / 'snp_gene_weight_matrix'
+            snp_gene_weight_matrix_save_dir = (
+                Path(self.config.ldscore_save_dir) / "snp_gene_weight_matrix"
+            )
             snp_gene_weight_matrix_save_dir.mkdir(parents=True, exist_ok=True)
-            logger.info(f'Saving snp_gene_weight_matrix for chr{chrom}...')
+            logger.info(f"Saving snp_gene_weight_matrix for chr{chrom}...")
             self.snp_gene_weight_matrix.reset_index().to_feather(
-                snp_gene_weight_matrix_save_dir / f'{chrom}.snp_gene_weight_matrix.feather')
+                snp_gene_weight_matrix_save_dir / f"{chrom}.snp_gene_weight_matrix.feather"
+            )
 
         # convert to sparse
         self.snp_gene_weight_matrix = csr_matrix(self.snp_gene_weight_matrix)
-        logger.info(f'Compute snp_gene_weight_matrix finished. shape: {self.snp_gene_weight_matrix.shape}')
+        logger.info(
+            f"Compute snp_gene_weight_matrix finished. shape: {self.snp_gene_weight_matrix.shape}"
+        )
 
         # calculate baseline ld score
-        logger.info(f'Calculating baseline ld score for chr{chrom}...')
-        self.calculate_ldscore_for_base_line(chrom, self.config.sample_name, self.config.ldscore_save_dir)
+        logger.info(f"Calculating baseline ld score for chr{chrom}...")
+        self.calculate_ldscore_for_base_line(
+            chrom, self.config.sample_name, self.config.ldscore_save_dir
+        )
 
         # calculate ld score for annotation
-        logger.info(f'Calculating ld score for annotation for chr{chrom}...')
+        logger.info(f"Calculating ld score for annotation for chr{chrom}...")
         self.calculate_ldscore_use_SNP_Gene_weight_matrix_by_chr(
             self.mk_score_common.loc[self.snp_gene_pair_dummy.columns[:-1]],
             chrom,
@@ -384,11 +425,11 @@ class S_LDSC_Boost:
             self.config.ldscore_save_dir,
         )
 
-    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,
+    ):
         if drop_dummy_na:
             ldscore_chr_chunk = self.snp_gene_weight_matrix[:, :-1] @ mk_score_chunk
         else:
@@ -407,16 +448,20 @@ class S_LDSC_Boost:
         #     self.keep_snp_mask]
 
         # save for each chunk
-        df = pd.DataFrame(ldscore_chr_chunk,
-                          index=self.snp_name,
-                          columns=column_names,
-                          )
-        df.index.name = 'SNP'
+        df = pd.DataFrame(
+            ldscore_chr_chunk,
+            index=self.snp_name,
+            columns=column_names,
+        )
+        df.index.name = "SNP"
         df.reset_index().to_feather(save_file_name)
 
-    def save_ldscore_chunk_to_zarr(self, ldscore_chr_chunk: np.ndarray,
-                                   chrom: int, start_col_index,
-                                   ):
+    def save_ldscore_chunk_to_zarr(
+        self,
+        ldscore_chr_chunk: np.ndarray,
+        chrom: int,
+        start_col_index,
+    ):
         ldscore_chr_chunk = ldscore_chr_chunk.astype(np.float16, copy=False)
         # avoid overflow of float16, if inf, set to max of float16
         ldscore_chr_chunk[np.isinf(ldscore_chr_chunk)] = np.finfo(np.float16).max
@@ -425,63 +470,90 @@ class S_LDSC_Boost:
         chrom_snp_start_point = self.chrom_snp_start_point[chrom - 1]
         chrom_snp_end_point = self.chrom_snp_start_point[chrom]
 
-        self.zarr_file[chrom_snp_start_point:chrom_snp_end_point,
-        start_col_index:start_col_index + ldscore_chr_chunk.shape[1]] = ldscore_chr_chunk
-
-    def calculate_M_use_SNP_gene_pair_dummy_by_chunk(self,
-                                                     mk_score_chunk,
-                                                     M_file_path, M_5_file_path,
-                                                     drop_dummy_na=True,
-                                                     ):
-        '''
-        calculate M use SNP_gene_pair_dummy_sumed_along_snp_axis and mk_score_chunk
-        '''
-        SNP_gene_pair_dummy_sumed_along_snp_axis = self.snp_gene_pair_dummy.values.sum(axis=0, keepdims=True)
-        SNP_gene_pair_dummy_sumed_along_snp_axis_pass_maf = self.snp_gene_pair_dummy.loc[self.snp_pass_maf].values.sum(
-            axis=0,
-            keepdims=True)
+        self.zarr_file[
+            chrom_snp_start_point:chrom_snp_end_point,
+            start_col_index : start_col_index + ldscore_chr_chunk.shape[1],
+        ] = ldscore_chr_chunk
+
+    def calculate_M_use_SNP_gene_pair_dummy_by_chunk(
+        self,
+        mk_score_chunk,
+        M_file_path,
+        M_5_file_path,
+        drop_dummy_na=True,
+    ):
+        """
+        Calculate M use SNP_gene_pair_dummy_sumed_along_snp_axis and mk_score_chunk
+        """
+        SNP_gene_pair_dummy_sumed_along_snp_axis = self.snp_gene_pair_dummy.values.sum(
+            axis=0, keepdims=True
+        )
+        SNP_gene_pair_dummy_sumed_along_snp_axis_pass_maf = self.snp_gene_pair_dummy.loc[
+            self.snp_pass_maf
+        ].values.sum(axis=0, keepdims=True)
         if drop_dummy_na:
-            SNP_gene_pair_dummy_sumed_along_snp_axis = SNP_gene_pair_dummy_sumed_along_snp_axis[:, :-1]
-            SNP_gene_pair_dummy_sumed_along_snp_axis_pass_maf = SNP_gene_pair_dummy_sumed_along_snp_axis_pass_maf[:,
-                                                                :-1]
+            SNP_gene_pair_dummy_sumed_along_snp_axis = SNP_gene_pair_dummy_sumed_along_snp_axis[
+                :, :-1
+            ]
+            SNP_gene_pair_dummy_sumed_along_snp_axis_pass_maf = (
+                SNP_gene_pair_dummy_sumed_along_snp_axis_pass_maf[:, :-1]
+            )
         save_dir = Path(M_file_path).parent
         save_dir.mkdir(parents=True, exist_ok=True)
         M_chr_chunk = SNP_gene_pair_dummy_sumed_along_snp_axis @ mk_score_chunk
         M_5_chr_chunk = SNP_gene_pair_dummy_sumed_along_snp_axis_pass_maf @ mk_score_chunk
-        np.savetxt(M_file_path, M_chr_chunk, delimiter='\t', )
-        np.savetxt(M_5_file_path, M_5_chr_chunk, delimiter='\t', )
+        np.savetxt(
+            M_file_path,
+            M_chr_chunk,
+            delimiter="\t",
+        )
+        np.savetxt(
+            M_5_file_path,
+            M_5_chr_chunk,
+            delimiter="\t",
+        )
 
-    def calculate_ldscore_use_SNP_Gene_weight_matrix_by_chr(self, mk_score_common, chrom, sample_name, save_dir):
+    def calculate_ldscore_use_SNP_Gene_weight_matrix_by_chr(
+        self, mk_score_common, chrom, sample_name, save_dir
+    ):
         """
         Calculate the LD score using the SNP-gene weight matrix.
         :param sample_name:
         """
         # Calculate the LD score
         chunk_index = 1
-        for i in trange(0, mk_score_common.shape[1], self.config.spots_per_chunk,
-                        desc=f'Calculating LD score by chunk for chr{chrom}'):
-            mk_score_chunk = mk_score_common.iloc[:, i:i + self.config.spots_per_chunk]
-
-            ld_score_file = f'{save_dir}/{sample_name}_chunk{chunk_index}/{sample_name}.{chrom}.l2.ldscore.feather'
-            M_file = f'{save_dir}/{sample_name}_chunk{chunk_index}/{sample_name}.{chrom}.l2.M'
-            M_5_file = f'{save_dir}/{sample_name}_chunk{chunk_index}/{sample_name}.{chrom}.l2.M_5_50'
+        for i in trange(
+            0,
+            mk_score_common.shape[1],
+            self.config.spots_per_chunk,
+            desc=f"Calculating LD score by chunk for chr{chrom}",
+        ):
+            mk_score_chunk = mk_score_common.iloc[:, i : i + self.config.spots_per_chunk]
+
+            ld_score_file = f"{save_dir}/{sample_name}_chunk{chunk_index}/{sample_name}.{chrom}.l2.ldscore.feather"
+            M_file = f"{save_dir}/{sample_name}_chunk{chunk_index}/{sample_name}.{chrom}.l2.M"
+            M_5_file = (
+                f"{save_dir}/{sample_name}_chunk{chunk_index}/{sample_name}.{chrom}.l2.M_5_50"
+            )
 
             ldscore_chr_chunk = self.calculate_ldscore_use_SNP_Gene_weight_matrix_by_chunk(
                 mk_score_chunk,
                 drop_dummy_na=True,
             )
-            if self.config.ldscore_save_format == 'feather':
-                self.save_ldscore_to_feather(ldscore_chr_chunk,
-                                             column_names=mk_score_chunk.columns,
-                                             save_file_name=ld_score_file,
-                                             )
-            elif self.config.ldscore_save_format == 'zarr':
-                self.save_ldscore_chunk_to_zarr(ldscore_chr_chunk,
-                                                chrom=chrom,
-                                                start_col_index=i,
-                                                )
+            if self.config.ldscore_save_format == "feather":
+                self.save_ldscore_to_feather(
+                    ldscore_chr_chunk,
+                    column_names=mk_score_chunk.columns,
+                    save_file_name=ld_score_file,
+                )
+            elif self.config.ldscore_save_format == "zarr":
+                self.save_ldscore_chunk_to_zarr(
+                    ldscore_chr_chunk,
+                    chrom=chrom,
+                    start_col_index=i,
+                )
             else:
-                raise ValueError(f'Invalid ldscore_save_format: {self.config.ldscore_save_format}')
+                raise ValueError(f"Invalid ldscore_save_format: {self.config.ldscore_save_format}")
 
             self.calculate_M_use_SNP_gene_pair_dummy_by_chunk(
                 mk_score_chunk,
@@ -496,21 +568,23 @@ class S_LDSC_Boost:
         # save baseline ld score
         baseline_mk_score = np.ones((self.snp_gene_pair_dummy.shape[1], 2))
         baseline_mk_score[-1, 0] = 0  # all_gene
-        baseline_mk_score_df = pd.DataFrame(baseline_mk_score, index=self.snp_gene_pair_dummy.columns,
-                                            columns=['all_gene', 'base'])
-        ld_score_file = f'{save_dir}/baseline/baseline.{chrom}.l2.ldscore.feather'
-        M_file = f'{save_dir}/baseline/baseline.{chrom}.l2.M'
-        M_5_file = f'{save_dir}/baseline/baseline.{chrom}.l2.M_5_50'
+        baseline_mk_score_df = pd.DataFrame(
+            baseline_mk_score, index=self.snp_gene_pair_dummy.columns, columns=["all_gene", "base"]
+        )
+        ld_score_file = f"{save_dir}/baseline/baseline.{chrom}.l2.ldscore.feather"
+        M_file = f"{save_dir}/baseline/baseline.{chrom}.l2.M"
+        M_5_file = f"{save_dir}/baseline/baseline.{chrom}.l2.M_5_50"
 
         ldscore_chr_chunk = self.calculate_ldscore_use_SNP_Gene_weight_matrix_by_chunk(
             baseline_mk_score_df,
             drop_dummy_na=False,
         )
 
-        self.save_ldscore_to_feather(ldscore_chr_chunk,
-                                     column_names=baseline_mk_score_df.columns,
-                                     save_file_name=ld_score_file,
-                                     )
+        self.save_ldscore_to_feather(
+            ldscore_chr_chunk,
+            column_names=baseline_mk_score_df.columns,
+            save_file_name=ld_score_file,
+        )
         # save baseline M
         self.calculate_M_use_SNP_gene_pair_dummy_by_chunk(
             baseline_mk_score_df,
@@ -519,7 +593,10 @@ class S_LDSC_Boost:
             drop_dummy_na=False,
         )
 
-    def get_snp_gene_dummy(self, chrom, ):
+    def get_snp_gene_dummy(
+        self,
+        chrom,
+    ):
         """
         Get the dummy matrix of SNP-gene pairs.
         """
@@ -527,91 +604,126 @@ class S_LDSC_Boost:
         print("Loading bim data")
         bim, bim_pr = load_bim(self.config.bfile_root, chrom)
 
-        if self.config.gene_window_enhancer_priority in ['gene_window_first', 'enhancer_first']:
-
-            SNP_gene_pair_gtf = self.get_SNP_gene_pair_from_gtf(bim, bim_pr, )
-            SNP_gene_pair_enhancer = self.get_SNP_gene_pair_from_enhancer(bim, bim_pr, )
+        if self.config.gene_window_enhancer_priority in ["gene_window_first", "enhancer_first"]:
+            SNP_gene_pair_gtf = self.get_SNP_gene_pair_from_gtf(
+                bim,
+                bim_pr,
+            )
+            SNP_gene_pair_enhancer = self.get_SNP_gene_pair_from_enhancer(
+                bim,
+                bim_pr,
+            )
             # total_SNP_gene_pair = SNP_gene_pair_gtf.join(SNP_gene_pair_enhancer, how='outer', lsuffix='_gtf', )
 
             mask_of_nan_gtf = SNP_gene_pair_gtf.gene_name.isna()
             mask_of_nan_enhancer = SNP_gene_pair_enhancer.gene_name.isna()
 
-            if self.config.gene_window_enhancer_priority == 'gene_window_first':
+            if self.config.gene_window_enhancer_priority == "gene_window_first":
                 SNP_gene_pair = SNP_gene_pair_gtf
-                SNP_gene_pair.loc[mask_of_nan_gtf, 'gene_name'] = SNP_gene_pair_enhancer.loc[
-                    mask_of_nan_gtf, 'gene_name']
-            elif self.config.gene_window_enhancer_priority == 'enhancer_first':
+                SNP_gene_pair.loc[mask_of_nan_gtf, "gene_name"] = SNP_gene_pair_enhancer.loc[
+                    mask_of_nan_gtf, "gene_name"
+                ]
+            elif self.config.gene_window_enhancer_priority == "enhancer_first":
                 SNP_gene_pair = SNP_gene_pair_enhancer
-                SNP_gene_pair.loc[mask_of_nan_enhancer, 'gene_name'] = SNP_gene_pair_gtf.loc[
-                    mask_of_nan_enhancer, 'gene_name']
+                SNP_gene_pair.loc[mask_of_nan_enhancer, "gene_name"] = SNP_gene_pair_gtf.loc[
+                    mask_of_nan_enhancer, "gene_name"
+                ]
             else:
                 raise ValueError(
-                    f'Invalid self.config.gene_window_enhancer_priority: {self.config.gene_window_enhancer_priority}')
+                    f"Invalid self.config.gene_window_enhancer_priority: {self.config.gene_window_enhancer_priority}"
+                )
 
         elif self.config.gene_window_enhancer_priority is None:  # use gtf only
-            SNP_gene_pair_gtf = self.get_SNP_gene_pair_from_gtf(bim, bim_pr, )
+            SNP_gene_pair_gtf = self.get_SNP_gene_pair_from_gtf(
+                bim,
+                bim_pr,
+            )
             SNP_gene_pair = SNP_gene_pair_gtf
 
-        elif self.config.gene_window_enhancer_priority == 'enhancer_only':
-            SNP_gene_pair_enhancer = self.get_SNP_gene_pair_from_enhancer(bim, bim_pr, )
+        elif self.config.gene_window_enhancer_priority == "enhancer_only":
+            SNP_gene_pair_enhancer = self.get_SNP_gene_pair_from_enhancer(
+                bim,
+                bim_pr,
+            )
             SNP_gene_pair = SNP_gene_pair_enhancer
         else:
-            raise ValueError('gtf_pr and enhancer_pr cannot be None at the same time')
+            raise ValueError("gtf_pr and enhancer_pr cannot be None at the same time")
 
         # save the SNP_gene_pair to feather
-        SNP_gene_pair_save_path = Path(
-            self.config.ldscore_save_dir) / f'SNP_gene_pair/SNP_gene_pair_chr{chrom}.feather'
+        SNP_gene_pair_save_path = (
+            Path(self.config.ldscore_save_dir) / f"SNP_gene_pair/SNP_gene_pair_chr{chrom}.feather"
+        )
         SNP_gene_pair_save_path.parent.mkdir(parents=True, exist_ok=True)
         SNP_gene_pair.reset_index().to_feather(SNP_gene_pair_save_path)
 
         # Get the dummy matrix
-        SNP_gene_pair_dummy = pd.get_dummies(SNP_gene_pair['gene_name'], dummy_na=True)
+        SNP_gene_pair_dummy = pd.get_dummies(SNP_gene_pair["gene_name"], dummy_na=True)
         return SNP_gene_pair_dummy
 
     def get_SNP_gene_pair_from_gtf(self, bim, bim_pr):
         logger.info(
-            "Get SNP-gene pair from gtf, if a SNP is in multiple genes, it will be assigned to the most nearby gene (TSS)")
+            "Get SNP-gene pair from gtf, if a SNP is in multiple genes, it will be assigned to the most nearby gene (TSS)"
+        )
         overlaps_small = Overlaps_gtf_bim(self.gtf_pr, bim_pr)
         # Get the SNP-gene pair
         annot = bim[["CHR", "BP", "SNP", "CM"]]
-        SNP_gene_pair = overlaps_small[['SNP', 'gene_name']].set_index('SNP').join(annot.set_index('SNP'), how='right')
+        SNP_gene_pair = (
+            overlaps_small[["SNP", "gene_name"]]
+            .set_index("SNP")
+            .join(annot.set_index("SNP"), how="right")
+        )
         return SNP_gene_pair
 
-    def get_SNP_gene_pair_from_enhancer(self, bim, bim_pr, ):
+    def get_SNP_gene_pair_from_enhancer(
+        self,
+        bim,
+        bim_pr,
+    ):
         logger.info(
-            "Get SNP-gene pair from enhancer, if a SNP is in multiple genes, it will be assigned to the gene with highest marker score")
+            "Get SNP-gene pair from enhancer, if a SNP is in multiple genes, it will be assigned to the gene with highest marker score"
+        )
         # Get the SNP-gene pair
         overlaps_small = self.enhancer_pr.join(bim_pr).df
         annot = bim[["CHR", "BP", "SNP", "CM"]]
-        if self.config.snp_multiple_enhancer_strategy == 'max_mkscore':
-            logger.debug('select the gene with highest marker score')
-            overlaps_small = overlaps_small.loc[overlaps_small.groupby('SNP').avg_mkscore.idxmax()]
-
-        elif self.config.snp_multiple_enhancer_strategy == 'nearest_TSS':
-            logger.debug('select the gene with nearest TSS')
-            overlaps_small['Distance'] = np.abs(overlaps_small['Start_b'] - overlaps_small['TSS'])
-            overlaps_small = overlaps_small.loc[overlaps_small.groupby('SNP').Distance.idxmin()]
-
-        SNP_gene_pair = overlaps_small[['SNP', 'gene_name']].set_index('SNP').join(annot.set_index('SNP'), how='right')
+        if self.config.snp_multiple_enhancer_strategy == "max_mkscore":
+            logger.debug("select the gene with highest marker score")
+            overlaps_small = overlaps_small.loc[overlaps_small.groupby("SNP").avg_mkscore.idxmax()]
+
+        elif self.config.snp_multiple_enhancer_strategy == "nearest_TSS":
+            logger.debug("select the gene with nearest TSS")
+            overlaps_small["Distance"] = np.abs(overlaps_small["Start_b"] - overlaps_small["TSS"])
+            overlaps_small = overlaps_small.loc[overlaps_small.groupby("SNP").Distance.idxmin()]
+
+        SNP_gene_pair = (
+            overlaps_small[["SNP", "gene_name"]]
+            .set_index("SNP")
+            .join(annot.set_index("SNP"), how="right")
+        )
 
         return SNP_gene_pair
 
 
 def run_generate_ldscore(config: GenerateLDScoreConfig):
-    if config.ldscore_save_format == 'quick_mode':
-        logger.info('Running in quick_mode. Skip the process of generating ldscore. Using the pre-calculated ldscore.')
+    if config.ldscore_save_format == "quick_mode":
+        logger.info(
+            "Running in quick_mode. Skip the process of generating ldscore. Using the pre-calculated ldscore."
+        )
         ldscore_save_dir = config.ldscore_save_dir
 
         # link the baseline annotation
         baseline_annotation_dir = Path(config.baseline_annotation_dir)
-        (ldscore_save_dir / 'baseline').symlink_to(baseline_annotation_dir, target_is_directory=True)
+        (ldscore_save_dir / "baseline").symlink_to(
+            baseline_annotation_dir, target_is_directory=True
+        )
 
         # link the SNP_gene_pair
         SNP_gene_pair_dir = Path(config.SNP_gene_pair_dir)
-        (ldscore_save_dir / 'SNP_gene_pair').symlink_to(SNP_gene_pair_dir, target_is_directory=True)
+        (ldscore_save_dir / "SNP_gene_pair").symlink_to(
+            SNP_gene_pair_dir, target_is_directory=True
+        )
         return
     s_ldsc_boost = S_LDSC_Boost(config)
-    if config.chrom == 'all':
+    if config.chrom == "all":
         for chrom in range(1, 23):
             s_ldsc_boost.process_chromosome(chrom)
     else: