Feature/np core 20241215

Dockerfile

@@ -1,7 +1,7 @@
 FROM ubuntu:22.04
 
 LABEL maintainer="cschu1981@gmail.com"
-LABEL version="2.18.5"
+LABEL version="2.20.0"
 LABEL description="gffquant - functional profiling of metagenomic/transcriptomic wgs samples"
 
 

gffquant/alignment/__init__.py

@@ -6,6 +6,7 @@
 
 from .aln_group import AlignmentGroup
 from .pysam_alignment_processor import AlignmentProcessor
+from .reference_hit import ReferenceHit
 from .samflags import SamFlags
 from .cigarops import CigarOps
 

gffquant/alignment/aln_group.py

@@ -79,6 +79,9 @@ def get_all_hits(self, as_ambiguous=False):
                     except TypeError as err:
                         raise TypeError(f"Cannot derive sequencing library from tags: {aln.tags}") from err
 
+                # in region mode, there can be more hits
+                # (if the alignment overlaps multiple features of the target sequence)
+                # in gene mode, each alignment is a hit, i.e. there is at most 1 hit / alignment
                 yield aln.hits, n_aln
 
     def get_ambig_align_counts(self):

gffquant/alignment/reference_hit.py

@@ -0,0 +1,38 @@
+# pylint: disable=R0902
+
+""" module docstring """
+
+from dataclasses import dataclass, asdict
+
+
+@dataclass(slots=True)
+class ReferenceHit:
+    rid: int = None
+    start: int = None
+    end: int = None
+    rev_strand: bool = None
+    cov_start: int = None
+    cov_end: int = None
+    has_annotation: bool = None
+    n_aln: int = None
+    is_ambiguous: bool = None
+    library_mod: int = None
+    mate_id: int = None
+
+    def __hash__(self):
+        return hash(tuple(asdict(self).values()))
+
+    def __eq__(self, other):
+        return all(
+            item[0][1] == item[1][1]
+            for item in zip(
+                sorted(asdict(self).items()),
+                sorted(asdict(other).items())
+            )
+        )
+
+    def __str__(self):
+        return "\t".join(map(str, asdict(self).values()))
+
+    def __repr__(self):
+        return str(self)

gffquant/annotation/__init__.py

@@ -2,5 +2,8 @@
 
 """ module docstring """
 
-from .count_annotator import GeneCountAnnotator, RegionCountAnnotator
+# from .count_annotator import GeneCountAnnotator, RegionCountAnnotator
+from .count_annotator import CountAnnotator
 from .count_writer import CountWriter
+from .genecount_annotator import GeneCountAnnotator
+from .regioncount_annotator import RegionCountAnnotator

gffquant/annotation/count_annotator.py

@@ -98,6 +98,7 @@ def calc_scaling_factor(raw, normed, default=0):
             return (raw / normed) if normed else default
 
         total_uniq, total_uniq_normed, total_ambi, total_ambi_normed = self.total_counts
+        # total_uniq, total_ambi, total_uniq_normed, total_ambi_normed = self.total_counts
         logger.info(
             "TOTAL COUNTS: uraw=%s unorm=%s araw=%s anorm=%s",
             total_uniq, total_uniq_normed, total_ambi, total_ambi_normed
@@ -111,19 +112,20 @@ def calc_scaling_factor(raw, normed, default=0):
             total_ambi, total_ambi_normed, default_scaling_factor
         )
 
-        total_uniq, total_uniq_normed, total_ambi, total_ambi_normed = self.total_gene_counts
-        logger.info(
-            "TOTAL GENE COUNTS: uraw=%s unorm=%s araw=%s anorm=%s",
-            total_uniq, total_uniq_normed, total_ambi, total_ambi_normed
-        )
+        # total_uniq, total_uniq_normed, total_ambi, total_ambi_normed = self.total_gene_counts
+        # total_uniq, total_ambi, total_uniq_normed, total_ambi_normed = self.total_gene_counts
+        # logger.info(
+        #     "TOTAL GENE COUNTS: uraw=%s unorm=%s araw=%s anorm=%s",
+        #     total_uniq, total_uniq_normed, total_ambi, total_ambi_normed
+        # )
 
-        self.scaling_factors["total_gene_uniq"] = calc_scaling_factor(
-            total_uniq, total_uniq_normed, default_scaling_factor
-        )
+        # self.scaling_factors["total_gene_uniq"] = calc_scaling_factor(
+        #     total_uniq, total_uniq_normed, default_scaling_factor
+        # )
 
-        self.scaling_factors["total_gene_ambi"] = calc_scaling_factor(
-            total_ambi, total_ambi_normed, default_scaling_factor
-        )
+        # self.scaling_factors["total_gene_ambi"] = calc_scaling_factor(
+        #     total_ambi, total_ambi_normed, default_scaling_factor
+        # )
 
         fc_items = self.feature_count_sums.items()
         for category, (

gffquant/annotation/count_writer.py

@@ -1,4 +1,4 @@
-# pylint: disable=C0103,W1514,R0913,R0917
+# pylint: disable=C0103,W1514,R0913,R0917,R0914
 
 """ module docstring """
 
@@ -8,6 +8,9 @@
 
 import numpy as np
 
+from ..counters import AlignmentCounter
+from ..counters.count_matrix import CountMatrix
+
 
 logger = logging.getLogger(__name__)
 
@@ -74,6 +77,7 @@ def compile_block(raw, lnorm, scaling_factors):
 
         p, row = 0, []
         rpkm_factor = 1e9 / self.filtered_readcount
+
         # unique counts
         row += compile_block(*counts[p:p + 2], (scaling_factor, rpkm_factor,))
         p += 2
@@ -107,69 +111,92 @@ def compile_block(raw, lnorm, scaling_factors):
     def write_row(header, data, stream=sys.stdout):
         print(header, *(f"{c:.5f}" for c in data), flush=True, sep="\t", file=stream)
 
-    # pylint: disable=R0914
-    def write_feature_counts(self, db, featcounts, unannotated_reads=None, report_unseen=True):
-        for category_id, counts in sorted(featcounts.items()):
-            scaling_factor, ambig_scaling_factor = featcounts.scaling_factors[
-                category_id
-            ]
-            category = db.query_category(category_id).name
-            if "scaled" in self.publish_reports:
-                logger.info(
-                    "SCALING FACTORS %s %s %s",
-                    category, scaling_factor, ambig_scaling_factor
+    def write_category(
+        self,
+        category_id,
+        category_name,
+        category_sum,
+        counts,
+        # feature_names,
+        features,
+        unannotated_reads=None,
+        report_unseen=True,
+    ):
+        with gzip.open(f"{self.out_prefix}.{category_name}.txt.gz", "wt") as feat_out:
+            header = self.get_header()
+            print("feature", *header, sep="\t", file=feat_out)
+
+            if unannotated_reads is not None:
+                print("unannotated", unannotated_reads, sep="\t", file=feat_out)
+
+            if "total_readcount" in self.publish_reports:
+                CountWriter.write_row(
+                    "total_reads",
+                    np.zeros(len(header)) + self.total_readcount,
+                    stream=feat_out,
                 )
-            with gzip.open(f"{self.out_prefix}.{category}.txt.gz", "wt") as feat_out:
-                header = self.get_header()
-                print("feature", *header, sep="\t", file=feat_out)
-
-                if unannotated_reads is not None:
-                    print("unannotated", unannotated_reads, sep="\t", file=feat_out)
-
-                if "total_readcount" in self.publish_reports:
-                    CountWriter.write_row(
-                        "total_reads",
-                        np.zeros(len(header)) + self.total_readcount,
-                        stream=feat_out,
-                    )
 
-                if "filtered_readcount" in self.publish_reports:
-                    CountWriter.write_row(
-                        "filtered_reads",
-                        np.zeros(len(header)) + self.filtered_readcount,
-                        stream=feat_out,
-                    )
+            if "filtered_readcount" in self.publish_reports:
+                CountWriter.write_row(
+                    "filtered_reads",
+                    np.zeros(len(header)) + self.filtered_readcount,
+                    stream=feat_out,
+                )
 
-                if "category" in self.publish_reports:
-                    cat_counts = counts.get(f"cat:::{category_id}")
-                    if cat_counts is not None:
-                        cat_row = self.compile_output_row(
-                            cat_counts,
-                            scaling_factor=featcounts.scaling_factors["total_uniq"],
-                            ambig_scaling_factor=featcounts.scaling_factors["total_ambi"],
-                        )
-                        CountWriter.write_row("category", cat_row, stream=feat_out)
-
-                for feature in db.get_features(category_id):
-                    f_counts = counts.get(str(feature.id), np.zeros(len(header)))
-                    if report_unseen or f_counts.sum():
-                        out_row = self.compile_output_row(
-                            f_counts,
-                            scaling_factor=scaling_factor,
-                            ambig_scaling_factor=ambig_scaling_factor,
-                        )
-                        CountWriter.write_row(feature.name, out_row, stream=feat_out)
-
-    def write_gene_counts(self, gene_counts, uniq_scaling_factor, ambig_scaling_factor):
-        if "scaled" in self.publish_reports:
-            logger.info("SCALING_FACTORS %s %s", uniq_scaling_factor, ambig_scaling_factor)
+            if "category" in self.publish_reports:
+                # cat_counts = counts[0]
+                cat_counts = category_sum
+                logger.info("CAT %s: %s", category_name, str(cat_counts))
+                if cat_counts is not None:
+                    CountWriter.write_row("category", category_sum, stream=feat_out)
+
+            # for item in counts:
+            #     if not isinstance(item[0], tuple):
+            #         logger.info("ITEM: %s", str(item))
+            #         raise TypeError(f"Weird key: {str(item)}")
+            #     (cid, fid), fcounts = item
+            #     if (report_unseen or fcounts.sum()) and cid == category_id:
+            #         CountWriter.write_row(feature_names[fid], fcounts, stream=feat_out,)
+
+            empty_row = np.zeros(6, dtype=CountMatrix.NUMPY_DTYPE)
+            for feature in features:
+                key = (category_id, feature.id)
+                if counts.has_record(key):
+                    row = counts[key]
+                else:
+                    row = empty_row
+                if (report_unseen or row.sum()):
+                    CountWriter.write_row(feature.name, row, stream=feat_out,)
+
+
+            # for (cid, fid), fcounts in counts:
+            #     if (report_unseen or fcounts.sum()) and cid == category_id:
+            #         CountWriter.write_row(feature_names[fid], fcounts, stream=feat_out,)
+
+    def write_gene_counts(
+        self,
+        gene_counts: AlignmentCounter,
+        refmgr,
+        gene_group_db=False,
+    ):
         with gzip.open(f"{self.out_prefix}.gene_counts.txt.gz", "wt") as gene_out:
             print("gene", *self.get_header(), sep="\t", file=gene_out, flush=True)
 
-            for gene, g_counts in sorted(gene_counts.items()):
-                out_row = self.compile_output_row(
-                    g_counts,
-                    scaling_factor=uniq_scaling_factor,
-                    ambig_scaling_factor=ambig_scaling_factor
+            ref_stream = (
+                (
+                    refmgr.get(rid[0] if isinstance(rid, tuple) else rid)[0],
+                    rid,
                 )
-                CountWriter.write_row(gene, out_row, stream=gene_out)
+                for rid, _ in gene_counts
+            )
+
+            for ref, rid in sorted(ref_stream):
+                counts = gene_counts[rid]
+                # if gene_group_db:
+                #     ref_tokens = ref.split(".")
+                #     gene_id, _ = ".".join(ref_tokens[:-1]), ref_tokens[-1]
+                # else:
+                #     gene_id = ref
+                gene_id = ref
+
+                CountWriter.write_row(gene_id, counts, stream=gene_out,)

gffquant/annotation/genecount_annotator.py

@@ -0,0 +1,79 @@
+# pylint: disable=R0914
+
+""" module docstring """
+import logging
+
+import numpy as np
+
+from .count_annotator import CountAnnotator
+from ..counters import AlignmentCounter
+from ..counters.count_matrix import CountMatrix
+from ..db.annotation_db import AnnotationDatabaseManager
+
+
+logger = logging.getLogger(__name__)
+
+
+class GeneCountAnnotator(CountAnnotator):
+    """ CountAnnotator subclass for gene-based counting. """
+
+    def __init__(self, strand_specific, report_scaling_factors=True):
+        """ __init__() """
+        CountAnnotator.__init__(self, strand_specific, report_scaling_factors=report_scaling_factors)
+
+    def annotate_gene_counts(
+        self,
+        refmgr,
+        db: AnnotationDatabaseManager,
+        counter: AlignmentCounter,
+        gene_group_db=False
+    ):
+        categories = list(db.get_categories())
+        category_sums = np.zeros((len(categories), 6))
+        functional_counts = CountMatrix(6)
+
+        # for category in categories:
+        #     features = ((feature.name, feature) for feature in db.get_features(category.id))
+        #     for _, feature in sorted(features, key=lambda x: x[0]):
+        #         _ = functional_counts[(category.id, feature.id)]
+
+        for rid, counts in counter:
+            # counts = counter[rid]
+            if gene_group_db:
+                ggroup_id = rid
+                region_annotation = None
+                if ggroup_id != "0":
+                    region_annotation = db.query_sequence(int(ggroup_id, 16), grouped_db=True,)
+            else:
+                ref, _ = refmgr.get(rid[0] if isinstance(rid, tuple) else rid)
+                region_annotation = db.query_sequence(ref)
+
+            if region_annotation is not None:
+                _, _, region_annotation = region_annotation
+                for category_id, features in region_annotation:
+                    category_id = int(category_id)
+                    category_sums[category_id] += counts
+                    for feature_id in features:
+                        feature_id = int(feature_id)
+                        functional_counts[(category_id, feature_id)] += counts
+
+        functional_counts.drop_unindexed()
+
+        for i, category in enumerate(categories):
+            u_sf, c_sf = (
+                CountMatrix.calculate_scaling_factor(*category_sums[i][0:2]),
+                CountMatrix.calculate_scaling_factor(*category_sums[i][3:5]),
+            )
+
+            rows = tuple(
+                key[0] == category.id
+                for key, _ in functional_counts
+            )
+
+            functional_counts.scale_column(1, u_sf, rows=rows)
+            functional_counts.scale_column(4, c_sf, rows=rows)
+
+            category_sums[i, 2] = category_sums[i, 1] * u_sf
+            category_sums[i, 5] = category_sums[i, 4] * c_sf
+
+        return functional_counts, category_sums