add cpg related index to genome

Author: ruochiz

scprinter/genome.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import os.path
+import pickle
 from pathlib import Path
 
 import gffutils
@@ -278,6 +279,41 @@ def fetch_bias_bw(self, verify=False):
 
         return bias_bw
 
+    def fetch_cpg_index(self):
+        fa_path = str(self.fetch_fa())
+        cpg_index = fa_path + "_cpg_index.pkl"
+        if not os.path.exists(cpg_index):
+            print("Building CpG index...")
+            cpg_index_all, cpg_num, cpg_coord = build_CpG_index(self)
+            with open(cpg_index, "wb") as f:
+                pickle.dump((cpg_index_all, cpg_num, cpg_coord), f)
+        else:
+            with open(cpg_index, "rb") as f:
+                cpg_index_all, cpg_num, cpg_coord = pickle.load(f)
+        return cpg_index_all, cpg_num, cpg_coord
+
+
+def build_CpG_index(genome):
+    fa_path = Path(genome.fetch_fa())
+    fasta = Fasta(fa_path, as_raw=True)
+    cpg_index_all = {}
+    cpg_num = {}
+    chroms_all = set(genome.chrom_sizes.keys())
+    cpg_coord = {}
+    for chrom in chroms_all:
+        cpg_index_all[chrom] = {}
+        cpg_coord[chrom] = []
+        seq = fasta[chrom][:].upper()
+        count = 0
+        for i in trange(len(seq) - 1):
+            if seq[i : i + 2] == "CG":
+                cpg_index_all[chrom][i + 1] = count
+                cpg_coord[chrom].append(i + 1)
+                count += 1
+        cpg_num[chrom] = count
+    fasta.close()
+    return cpg_index_all, cpg_num, cpg_coord
+
 
 def predict_genome_tn5_bias(
     fa_file,

scprinter/seq/scripts/evaluation_model.py

@@ -254,6 +254,7 @@ def main(
             unfinished_id_strs = []
             for id, id_str in zip(ids, id_strs):
                 # id_str = "-".join([str(x) for x in id])
+
                 if os.path.exists(
                     os.path.join(
                         save_dir,
@@ -426,11 +427,15 @@ def main(
             id_str = id_strs[index]
             # id_str = "-".join([str(i) for i in id]) if id[0] is not None else None
             bar.set_description(f"working on {id_str}")
+            filename = f"model_{id_str}." if id[0] is not None else ""
+            filename += f"hypo.{wrapper}.{method}{extra}.{decay}."
+            if (id[0] is None) and ((start != 0) or (end != dataset.summits.shape[0])):
+                filename += f"{start}-{end}."
+            filename += "npz"
             if os.path.exists(
                 os.path.join(
                     save_dir,
-                    (f"model_{id_str}." if id[0] is not None else "")
-                    + f"hypo.{wrapper}.{method}{extra}.{decay}.npz",
+                    filename,
                 )
             ) and (not overwrite):
                 if verbose: