Refactor sequence normalization and statistics functions: streamline code for improved readability and performance, and replace custom implementations with core functions.

Author: imjiaoyuan

src/jsrc/analyze/core.py

@@ -2,15 +2,13 @@
 from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
 
+_NORM_TABLE = str.maketrans("U", "T", "".join(
+    c for c in map(chr, range(256)) if c not in "ACGTN"
+))
+
 
 def normalize_sequence(seq: str) -> str:
-    cleaned = []
-    for ch in seq.upper():
-        if ch == "U":
-            ch = "T"
-        if ch in {"A", "C", "G", "T", "N"}:
-            cleaned.append(ch)
-    return "".join(cleaned)
+    return seq.upper().translate(_NORM_TABLE)
 
 
 def pad_alignment(records: list[SeqRecord]) -> MultipleSeqAlignment:

src/jsrc/analyze/qc.py

@@ -1,28 +1,11 @@
-import gzip
 import json
 import re
 from argparse import Namespace
 from typing import Any
 
 from Bio import SeqIO
 
-
-def _open_text(path: str):
-    if path.endswith(".gz"):
-        return gzip.open(path, "rt", encoding="utf-8")
-    return open(path, encoding="utf-8")
-
-
-def _nxx(lengths: list[int], pct: float) -> int:
-    if not lengths:
-        return 0
-    target = sum(lengths) * pct
-    acc = 0
-    for v in sorted(lengths, reverse=True):
-        acc += v
-        if acc >= target:
-            return v
-    return 0
+from jsrc.core import nxx, open_text
 
 
 def _assembly_stats(fasta_path: str) -> dict[str, float | int]:
@@ -42,8 +25,8 @@ def _assembly_stats(fasta_path: str) -> dict[str, float | int]:
         "total_bases": total_len,
         "max_contig": max(lengths) if lengths else 0,
         "min_contig": min(lengths) if lengths else 0,
-        "n50": _nxx(lengths, 0.5),
-        "n90": _nxx(lengths, 0.9),
+        "n50": nxx(lengths, 0.5),
+        "n90": nxx(lengths, 0.9),
         "gc_percent": (gc / acgt * 100.0) if acgt else 0.0,
         "n_percent": (n_bases / total_len * 100.0) if total_len else 0.0,
     }
@@ -65,7 +48,7 @@ def _sam_stats(sam_path: str) -> dict[str, float | int]:
     mapped = 0
     covered_ref_bases = 0
     ref_len = 0
-    with _open_text(sam_path) as f:
+    with open_text(sam_path) as f:
         for line in f:
             if line.startswith("@SQ"):
                 fields = line.rstrip("\n").split("\t")
@@ -100,7 +83,7 @@ def _fastq_stats(paths: list[str], genome_size: int | None) -> dict[str, float |
     reads = 0
     bases = 0
     for path in paths:
-        with _open_text(path) as f:
+        with open_text(path) as f:
             for i, line in enumerate(f, start=1):
                 if i % 4 == 2:
                     reads += 1
@@ -120,7 +103,7 @@ def _vcf_stats(vcf_path: str) -> dict[str, int]:
     snp = 0
     indel = 0
     other = 0
-    with _open_text(vcf_path) as f:
+    with open_text(vcf_path) as f:
         for line in f:
             if line.startswith("#"):
                 continue

src/jsrc/core.py

@@ -1,7 +1,26 @@
+import gzip
 import sys
 import time
 from collections.abc import Generator, Iterable
-from typing import Any
+from typing import IO, Any
+
+
+def open_text(path: str) -> IO[str]:
+    if path.endswith(".gz"):
+        return gzip.open(path, "rt", encoding="utf-8")
+    return open(path, encoding="utf-8")
+
+
+def nxx(lengths: list[int], pct: float) -> int:
+    if not lengths:
+        return 0
+    target = sum(lengths) * pct
+    acc = 0
+    for v in sorted(lengths, reverse=True):
+        acc += v
+        if acc >= target:
+            return v
+    return 0
 
 
 def _fmt_duration(seconds: float) -> str:

src/jsrc/seq/extract.py

@@ -80,9 +80,10 @@ def cmd(args: Namespace) -> None:
                 if chrom not in genome:
                     continue
                 chrom_seq = genome[chrom].seq
-                seq = Seq("")
+                seq_parts = []
                 for start, end in regions:
-                    seq += chrom_seq[start:end]
+                    seq_parts.append(str(chrom_seq[start:end]))
+                seq = Seq("".join(seq_parts))
                 if strand == "-":
                     seq = seq.reverse_complement()
                 desc = f"feature={args.feature};match={args.match};locus={chrom};strand={strand}"

src/jsrc/seq/kmer.py

@@ -9,6 +9,9 @@
 from jsrc.core import progressbar
 
 
+_VALID_BASES = frozenset("ACGT")
+
+
 def _kmer_counter(path: str, k: int) -> Counter:
     c: Counter[str] = Counter()
     records = list(SeqIO.parse(path, "fasta"))
@@ -17,7 +20,7 @@ def _kmer_counter(path: str, k: int) -> Counter:
         seq = str(rec.seq).upper().replace("U", "T")
         for i in range(0, len(seq) - k + 1):
             kmer = seq[i : i + k]
-            if set(kmer) <= {"A", "C", "G", "T"}:
+            if all(b in _VALID_BASES for b in kmer):
                 c[kmer] += 1
         bar.update()
     bar.finish()

src/jsrc/seq/qc.py

@@ -1,34 +1,15 @@
-import gzip
 import json
 import logging
 from argparse import Namespace
-from typing import IO, Any
+from typing import Any
 
 from Bio import SeqIO
 
-from jsrc.core import progressbar
+from jsrc.core import open_text, nxx, progressbar
 
 logger = logging.getLogger(__name__)
 
 
-def _open_text(path: str) -> IO[str]:
-    if path.endswith(".gz"):
-        return gzip.open(path, "rt", encoding="utf-8")
-    return open(path, encoding="utf-8")
-
-
-def _nxx(lengths: list[int], pct: float) -> int:
-    if not lengths:
-        return 0
-    target = sum(lengths) * pct
-    acc = 0
-    for x in sorted(lengths, reverse=True):
-        acc += x
-        if acc >= target:
-            return x
-    return 0
-
-
 def _fasta_stats(path: str) -> dict[str, float | int]:
     lengths = []
     gc = 0
@@ -46,8 +27,8 @@ def _fasta_stats(path: str) -> dict[str, float | int]:
         "total_bases": total,
         "max_len": max(lengths) if lengths else 0,
         "min_len": min(lengths) if lengths else 0,
-        "n50": _nxx(lengths, 0.5),
-        "n90": _nxx(lengths, 0.9),
+        "n50": nxx(lengths, 0.5),
+        "n90": nxx(lengths, 0.9),
         "gc_percent": (gc / acgt * 100.0) if acgt else 0.0,
         "n_percent": (n_bases / total * 100.0) if total else 0.0,
     }
@@ -58,7 +39,7 @@ def _fastq_stats(paths: list[str], genome_size: int | None) -> dict[str, float |
     bases = 0
     bar = progressbar(total=0, desc="FASTQ reads")
     for path in paths:
-        with _open_text(path) as f:
+        with open_text(path) as f:
             for line_no, line in enumerate(f, start=1):
                 if line_no % 4 == 2:
                     reads += 1

src/jsrc/vision/efd.py

@@ -36,27 +36,23 @@ def calculate(
             return np.zeros((order, 4), dtype=float)
 
         phi = (2.0 * np.pi * t) / total_len
-        coeffs = np.zeros((order, 4), dtype=float)
-        dphi_cos = {}
-        dphi_sin = {}
-
-        for n in range(1, order + 1):
-            phi_n = phi * n
-            dphi_cos[n] = np.cos(phi_n[1:]) - np.cos(phi_n[:-1])
-            dphi_sin[n] = np.sin(phi_n[1:]) - np.sin(phi_n[:-1])
+        ns = np.arange(1, order + 1)[:, None]
+        phi_n = ns * phi[None, :]
+        dphi_cos = np.cos(phi_n[:, 1:]) - np.cos(phi_n[:, :-1])
+        dphi_sin = np.sin(phi_n[:, 1:]) - np.sin(phi_n[:, :-1])
 
         dx_over_dt = np.zeros_like(dx, dtype=float)
         dy_over_dt = np.zeros_like(dy, dtype=float)
         dx_over_dt[valid] = dx[valid] / dt[valid]
         dy_over_dt[valid] = dy[valid] / dt[valid]
 
-        for n in range(1, order + 1):
-            const = total_len / (2.0 * (n * np.pi) ** 2)
-            an = const * np.sum(dx_over_dt * dphi_cos[n])
-            bn = const * np.sum(dx_over_dt * dphi_sin[n])
-            cn = const * np.sum(dy_over_dt * dphi_cos[n])
-            dn = const * np.sum(dy_over_dt * dphi_sin[n])
-            coeffs[n - 1] = [an, bn, cn, dn]
+        consts = total_len / (2.0 * (ns[:, 0] * np.pi) ** 2)
+        coeffs = np.column_stack([
+            consts * (dphi_cos @ dx_over_dt),
+            consts * (dphi_sin @ dx_over_dt),
+            consts * (dphi_cos @ dy_over_dt),
+            consts * (dphi_sin @ dy_over_dt),
+        ])
 
         if normalize:
             coeffs = EllipticFourier.normalize(coeffs)
@@ -108,11 +104,12 @@ def normalize(coeffs: np.ndarray) -> np.ndarray:
     @staticmethod
     def reconstruct(coeffs: np.ndarray, num_points: int = 300) -> np.ndarray:
         t = np.linspace(0.0, 1.0, num_points)
-        xt = np.zeros(num_points)
-        yt = np.zeros(num_points)
-        for n, (an, bn, cn, dn) in enumerate(coeffs, start=1):
-            xt += an * np.cos(2.0 * np.pi * n * t) + bn * np.sin(2.0 * np.pi * n * t)
-            yt += cn * np.cos(2.0 * np.pi * n * t) + dn * np.sin(2.0 * np.pi * n * t)
+        ns = np.arange(1, len(coeffs) + 1)[:, None]
+        angles = 2.0 * np.pi * ns * t[None, :]
+        ab = coeffs[:, :2]
+        cd = coeffs[:, 2:]
+        xt = (ab[:, 0:1] * np.cos(angles) + ab[:, 1:2] * np.sin(angles)).sum(axis=0)
+        yt = (cd[:, 0:1] * np.cos(angles) + cd[:, 1:2] * np.sin(angles)).sum(axis=0)
         return np.stack([xt, yt], axis=1)