Wrote annotate command, storing only unique and univ genes, bitmap boundary bug fix

Author: skovaka

panagram/__main__.py

@@ -91,16 +91,33 @@ def run(self):
 
         idx.close()
 
+@dataclasses.dataclass
+class Annotate:
+    """(Re-)annotate an existing anchored genome using a GFF file """
+
+    index_dir: str = field(positional=True, metavar="index_dir")
+    """Panagram index directory"""
+
+    genome: str = field(positional=True, metavar="genome_name")
+
+    gff_file: str = field(positional=True, metavar="gff_file")
+    
+    def run(self):
+        idx = Index(self.index_dir)
+        idx[self.genome].run_annotate(self.gff_file)
+        idx.close()
+
 @dataclasses.dataclass
 class Main:
     """Alignment-free pan-genome viewer
 
 Subcommands:
     index    Anchor KMC bitvectors to reference FASTA files
     view     Display panagram viewer in a browser window
+    annotate Create or replace GFF annotation for anchored genome
     bitdump  Query pan-kmer bitmap via the commandline"""
 
-    cmd: Union[View, Index, Bitdump]
+    cmd: Union[View, Index, Bitdump, Annotate]
     cprof: str = field(default=None, help=argparse.SUPPRESS)
 
     def run(self):

panagram/index.py

@@ -49,10 +49,10 @@ def init_logger(logfile):
 IDX_SUFFIX = "gzi"
 ANCHOR_DIR = "anchor"
 
-TABIX_COLS = ["chr","start","end","type","attr"]
-#TABIX_COLS = ["chr","start","end","type","name"]
+GFF_COLS = ["chr","start","end","type","attr"]
+TABIX_COLS = ["chr","start","end","type","name"]
 TABIX_TYPES = {"start" : int, "end" : int}
-GENE_TABIX_COLS = TABIX_COLS + ["unique","universal"]
+GENE_COLS = ["chr","start","end","name"]#,"unique","universal"]
 GENE_TABIX_TYPES = {"start" : int, "end" : int, "unique" : int, "universal" : int}
 TABIX_SUFFIX = ".gz"
 
@@ -95,7 +95,7 @@ class Index(Serializable):
 
     gff_gene_types: List[str] = field(default_factory=lambda: ["gene"], help="GFF features to store locations and conservation scores")
     gff_anno_types: List[str] = field(default=None, help="GFF features of which to store locations, but not conservation scores")
-    gff_attrs: List[str] = field(default_factory=lambda: ["ID","Name"], help="GFF attributes to store in annotation indexes")
+    gff_name: List[str] = field(default_factory=lambda: ["Name","ID"], help="GFF attributes to store in annotation indexes")
 
     #Subset of genome IDs to generate anchor genomes for. Will use all genomes as anchors if not specified
     anchor_genomes: List[str] = field(default=None)
@@ -473,12 +473,13 @@ def bitsum_index(self):
 
     @property
     def gene_tabix_cols(self):
-        return TABIX_COLS + list(self.bitsum_index)
+        return GENE_COLS + [1,self.ngenomes]#list(self.bitsum_index)
+        #return TABIX_COLS + list(self.bitsum_index)
 
     @property
     def gene_tabix_types(self):
         r = {"start" : int, "end" : int}
-        for i in self.bitsum_index:
+        for i in [1,self.ngenomes]:#self.bitsum_index:
             r[i] = int
         return r#TABIX_COLS + list(self.bitsum_index)
 
@@ -562,20 +563,21 @@ def seq_len(self, seq_name):
         return self.sizes.loc[seq_name]
 
     def _iter_gff(self):
-        gffattr = lambda df,name: df["attr"].str.extract(f"{name}=([^;]+)", re.IGNORECASE)
+        gffattr = lambda df,name: df["attr"].str.extract(f"{name}=([^;]+)", re.IGNORECASE)[0]
         for df in pd.read_csv(
                 self.gff,
                 sep="\t", comment="#", chunksize=10000,
                 names = ["chr","source","type","start","end","score","strand","phase","attr"],
-                usecols = TABIX_COLS):
-
-            #df["name"] = pd.NA
-            #for attr in df,self.gff_attrs:
-            #    isna = df.index[df["name"].isna()]
-            #    if len(isna) > 0:
-            #        df.loc[isna,"name"] = gffattr(df.loc[isna], attr)
-            #    else:
-            #        break
+                usecols = GFF_COLS):
+
+            df["name"] = pd.NA
+            for attr in self.params["gff_name"]:
+                isna = df.index[df["name"].isna()]
+                if len(isna) > 0:
+                    names = gffattr(df.loc[isna], attr)
+                    df.loc[isna,"name"] = names
+                else:
+                    break
 
             yield df[TABIX_COLS]
 
@@ -595,8 +597,10 @@ def _write_anno_types(self):
             for t in self.gff_anno_types:
                 f.write(f"{t}\n")
 
-    def init_gff(self):
-        if pd.isna(self.gff): return
+    def init_gff(self, filename=None):
+        if filename is None:
+            filename = self.gff
+        if pd.isna(filename): return
 
         genes = list()
         annos = list()
@@ -660,9 +664,9 @@ def query(self, name, start=None, end=None, step=1):
         if end is None:
             end = self.seq_len(name)
 
-        pac = self._query_bytes(name, start, end, step, bstep)
+        pac = self._query_bytes(name, start, end-1, step, bstep)
         bits = self._bytes_to_bits(pac)
-        idx = np.arange(start, end+1, step, dtype=int)#[:len(bits)-1]
+        idx = np.arange(start, end, step, dtype=int)#[:len(bits)-1]
         df = pd.DataFrame(bits, index=idx, columns=self.genome_names)
         return df
 
@@ -720,14 +724,14 @@ def query_genes(self, chrom=None, start=None, end=None):#, attrs=["name","id"]):
         rows = list(rows)
         ret = pd.DataFrame(rows, columns=self.gene_tabix_cols).astype(self.gene_tabix_types)
 
-        if "attr" in ret.columns:
-            attr = lambda a: ret["attr"].str.extract(f"{a}=([^;]+)", re.IGNORECASE)
-            names = attr("Name")
-            ids = attr("ID")
-            names[names.isna()] = ids[names.isna()]
-            ret["name"] = names
-        else:
-            ret["name"] = ""
+        #if "attr" in ret.columns:
+        #    attr = lambda a: ret["attr"].str.extract(f"{a}=([^;]+)", re.IGNORECASE)
+        #    names = attr("Name")
+        #    ids = attr("ID")
+        #    names[names.isna()] = ids[names.isna()]
+        #    ret["name"] = names
+        #else:
+        #    ret["name"] = ""
 
         return ret
 
@@ -789,6 +793,38 @@ def _write_bitmap(self, name, seq):
 
         return self._bytes_to_bits(arrs[1])
 
+    def run_annotate(self, gff_file=None, logfile=None):
+        logging.basicConfig(
+                filename=logfile, level=logging.INFO, 
+                format='[ %(asctime)s %(levelname)7s ] %(message)s',
+                datefmt="%Y-%m-%d %H:%M:%S"
+        )
+
+        gene_df = self.init_gff(gff_file)
+
+        #for chrom,df in gene_df.groupby(level="chr"):
+        for chrom in gene_df.index.unique("chr"):
+            df = gene_df.loc[chrom]
+            st = df.index.get_level_values("start").min()
+            en = min(self.sizes[chrom],df.index.get_level_values("end").max())
+
+            bitsum = self.query(chrom, st, en).sum(axis=1)
+
+            for start,end in df.index:
+                if end <= start or start < 0 or end-st > len(bitsum):
+                    logger.warning(f"Skipping gene at {chrom}:{start}-{end}, coordinates out-of-bounds")
+                    continue
+                occ, counts = np.unique(bitsum[start-st:end-st], return_counts=True)
+                gene_df.loc[(chrom,start,end),occ] += counts
+
+            self.bitsum_genes = gene_df.groupby("chr",sort=False)[self.bitsum_index].sum()#.sort_index()
+            self.bitsum_genes.to_csv(self.chr_genes_fname, sep="\t")
+
+            gene_tabix = gene_df.reset_index()[self.gene_tabix_cols]
+            print(gene_tabix)
+            self._write_tabix(gene_tabix, "gene")
+
+
 
     def run_anchor(self, bitvecs, logfile=None):
         logging.basicConfig(

panagram/view.py

@@ -665,9 +665,9 @@ def plot_interactive(n_skips, bitmap_counts, plot_rep, plot_gene, start_coord, e
         adjusted_bin_size = (bin_size/n_skips)
 
         #Here we are filling in the bins for the main figure.    
-        bin_size_int = n_skips*(int(bin_size) // n_skips)+1
+        bin_size_int = n_skips*(int(bin_size) // n_skips)#+1
 
-        adjusted_bin_size_init = int(adjusted_bin_size) + 1
+        adjusted_bin_size_init = int(adjusted_bin_size) #+ 1
 
         #t2 = time.perf_counter()
         #print(f"\tFigure init {t2 - t_start:0.4f} seconds")
@@ -1142,7 +1142,7 @@ def make_gene_per_genome_fig(anchor_name):
         genes = index.query_genes(anchor_name)
         tib = time.perf_counter()
         sys.stderr.write(f"Queried genes 3 ({tib-tic})\n")
-        genes["name"] = genes["attr"].str.extract("Name=([^;]+)")
+        #genes["name"] = genes["attr"].str.extract("Name=([^;]+)")
         genes["size"] = genes["end"] - genes["start"]
 
         x = [i for i in range(0, len(genes))]
@@ -1569,8 +1569,8 @@ def update_gene_locals(local_gene_list, chrs, start_coord, end_coord, anchor_nam
             printme = ""
         elif len(local_gene_list)==1:
             printme += "Genes: "
-            printme += local_gene_list['name'].iloc[0] + ": "
-            printme += local_gene_list['attr'].iloc[0] #index.query_anno(anchor_name, chrs, start_coord, end_coord)['attr']
+            printme += local_gene_list['name'].iloc[0] #+ ": "
+            #printme += local_gene_list['attr'].iloc[0] #index.query_anno(anchor_name, chrs, start_coord, end_coord)['attr']
         elif len(local_gene_list)<=25:
             printme += "Genes: "
             print(local_gene_list)