codonyat 1.3.0__tar.gz → 1.3.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. {codonyat-1.3.0 → codonyat-1.3.2}/PKG-INFO +2 -9
  2. {codonyat-1.3.0 → codonyat-1.3.2}/README.md +1 -7
  3. codonyat-1.3.2/aa_caller/aavf.py +169 -0
  4. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/app.py +1 -0
  5. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/cli.py +1 -2
  6. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/cli_options.py +10 -1
  7. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/constants.py +1 -1
  8. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/container.py +211 -65
  9. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/reference.py +43 -21
  10. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/runner.py +3 -1
  11. codonyat-1.3.2/aa_caller/sam.py +219 -0
  12. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/stats_writer.py +17 -9
  13. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/validators.py +11 -4
  14. {codonyat-1.3.0 → codonyat-1.3.2}/codonyat.egg-info/PKG-INFO +2 -9
  15. {codonyat-1.3.0 → codonyat-1.3.2}/codonyat.egg-info/SOURCES.txt +3 -1
  16. {codonyat-1.3.0 → codonyat-1.3.2}/pyproject.toml +1 -2
  17. codonyat-1.3.2/tests/test_non_hiv.py +128 -0
  18. codonyat-1.3.2/tests/test_v1_3_1_features.py +163 -0
  19. codonyat-1.3.0/aa_caller/aavf.py +0 -163
  20. codonyat-1.3.0/aa_caller/sam.py +0 -129
  21. {codonyat-1.3.0 → codonyat-1.3.2}/LICENSE +0 -0
  22. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/__init__.py +0 -0
  23. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/__main__.py +0 -0
  24. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/genetic_code.py +0 -0
  25. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/metrics.py +0 -0
  26. {codonyat-1.3.0 → codonyat-1.3.2}/aa_caller/models.py +0 -0
  27. {codonyat-1.3.0 → codonyat-1.3.2}/codonyat.egg-info/dependency_links.txt +0 -0
  28. {codonyat-1.3.0 → codonyat-1.3.2}/codonyat.egg-info/entry_points.txt +0 -0
  29. {codonyat-1.3.0 → codonyat-1.3.2}/codonyat.egg-info/requires.txt +0 -0
  30. {codonyat-1.3.0 → codonyat-1.3.2}/codonyat.egg-info/top_level.txt +0 -0
  31. {codonyat-1.3.0 → codonyat-1.3.2}/setup.cfg +0 -0
  32. {codonyat-1.3.0 → codonyat-1.3.2}/tests/test_app.py +0 -0
  33. {codonyat-1.3.0 → codonyat-1.3.2}/tests/test_integration.py +0 -0
  34. {codonyat-1.3.0 → codonyat-1.3.2}/tests/test_metrics.py +0 -0
  35. {codonyat-1.3.0 → codonyat-1.3.2}/tests/test_protein_selection.py +0 -0
  36. {codonyat-1.3.0 → codonyat-1.3.2}/tests/test_v1_2_0_features.py +0 -0
  37. {codonyat-1.3.0 → codonyat-1.3.2}/tests/test_v1_3_0_features.py +0 -0
@@ -1,13 +1,12 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: codonyat
3
- Version: 1.3.0
3
+ Version: 1.3.2
4
4
  Summary: codon-yat — A codon-aware (Yet Another) amino acid variant typer from SAM alignments of viral NGS data.
5
5
  Author-email: Marc Noguera Julian <info@treetopunder.com>
6
6
  License-Expression: MIT
7
7
  Project-URL: Source, https://github.com/mnoguera/codonyat
8
8
  Project-URL: Documentation, https://mnoguera.github.io/codonyat/
9
9
  Project-URL: Issues, https://github.com/mnoguera/codonyat/issues
10
- Project-URL: Nextflow pipeline (nf-codonyat), https://github.com/sjoclaudi/nf-codonyat
11
10
  Classifier: Programming Language :: Python :: 3
12
11
  Classifier: Programming Language :: Python :: 3.10
13
12
  Classifier: Programming Language :: Python :: 3.11
@@ -54,13 +53,6 @@ pip install codonyat
54
53
 
55
54
  Requires Python 3.10 or newer. The only runtime dependency is Biopython.
56
55
 
57
- ## Running it in a workflow
58
-
59
- This repository contains only the Python package. To run codonyat on many samples
60
- (from FASTQ or SAM, with read QC, trimming, alignment and a merged report), use the
61
- separate Nextflow pipeline **[nf-codonyat](https://github.com/sjoclaudi/nf-codonyat)**,
62
- which installs this package from PyPI.
63
-
64
56
  ## CLI usage
65
57
 
66
58
  Once installed, the `codonyat` entry point is available:
@@ -90,6 +82,7 @@ codonyat sample.sam reference.fasta amplicons.tsv --annotation annotations.tsv -
90
82
  | `--min-freq` | `0.01` | Frequency floor (fraction) for `lowfreq` / AAVF FILTER |
91
83
  | `--min-coverage` | `10` | Coverage floor for `lowcov` / AAVF FILTER |
92
84
  | `--sb-alpha` | `0.05` | Fisher p-value cutoff for the `sb` flag |
85
+ | `--collapsed-read-counts` | off | Treat a read name ending in `_<digits>` as a collapsed duplicate count. Off by default: each SAM row counts once |
93
86
 
94
87
  ### Output
95
88
 
@@ -25,13 +25,6 @@ pip install codonyat
25
25
 
26
26
  Requires Python 3.10 or newer. The only runtime dependency is Biopython.
27
27
 
28
- ## Running it in a workflow
29
-
30
- This repository contains only the Python package. To run codonyat on many samples
31
- (from FASTQ or SAM, with read QC, trimming, alignment and a merged report), use the
32
- separate Nextflow pipeline **[nf-codonyat](https://github.com/sjoclaudi/nf-codonyat)**,
33
- which installs this package from PyPI.
34
-
35
28
  ## CLI usage
36
29
 
37
30
  Once installed, the `codonyat` entry point is available:
@@ -61,6 +54,7 @@ codonyat sample.sam reference.fasta amplicons.tsv --annotation annotations.tsv -
61
54
  | `--min-freq` | `0.01` | Frequency floor (fraction) for `lowfreq` / AAVF FILTER |
62
55
  | `--min-coverage` | `10` | Coverage floor for `lowcov` / AAVF FILTER |
63
56
  | `--sb-alpha` | `0.05` | Fisher p-value cutoff for the `sb` flag |
57
+ | `--collapsed-read-counts` | off | Treat a read name ending in `_<digits>` as a collapsed duplicate count. Off by default: each SAM row counts once |
64
58
 
65
59
  ### Output
66
60
 
@@ -0,0 +1,169 @@
1
+ """AAVF v1.0 writer. Research use only. Never emits CAT or SRVL."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from collections import defaultdict
6
+ from collections.abc import Sequence
7
+ from datetime import date
8
+ from pathlib import Path
9
+
10
+ from .constants import PACKAGE_VERSION
11
+ from .genetic_code import codon_to_aminoacid
12
+ from .metrics import classify_flags, is_acgt_codon
13
+ from .models import Variant
14
+ from .reference import FullReference
15
+
16
+ _AAVF_AA = frozenset("ACDEFGHIKLMNPQRSTVWYX*")
17
+
18
+
19
+ def aavf_amino_acid(codon: str) -> str:
20
+ """Map a codon to an AAVF amino-acid letter (``*`` = stop, ``X`` = other)."""
21
+ aa = codon_to_aminoacid(codon)
22
+ if aa == "_":
23
+ return "*"
24
+ if aa in _AAVF_AA:
25
+ return aa
26
+ return "X"
27
+
28
+
29
+ def mark_changed_bases(ref_codon: str, alt_codon: str) -> str:
30
+ """Lowercase matches to the reference codon; uppercase substitutions.
31
+
32
+ Matches the AAVF v1.0 examples (``RC=aaa;AC=aaC``).
33
+ """
34
+ ref = ref_codon.lower()
35
+ alt = alt_codon.lower()
36
+ marked = []
37
+ for i, base in enumerate(alt):
38
+ if i < len(ref) and base == ref[i]:
39
+ marked.append(base)
40
+ else:
41
+ marked.append(base.upper())
42
+ return "".join(marked)
43
+
44
+
45
+ def _fmt_float(value: float) -> str:
46
+ return f"{value:.6g}"
47
+
48
+
49
+ def write_aavf(
50
+ sample_name: str,
51
+ reference: FullReference,
52
+ output_path: Path,
53
+ protein_pairs: Sequence[tuple[str, dict[int, dict]]],
54
+ *,
55
+ min_freq: float,
56
+ min_coverage: int,
57
+ sb_alpha: float,
58
+ file_date: str | None = None,
59
+ ) -> None:
60
+ """Write an AAVF v1.0 file for the counted proteins.
61
+
62
+ FILTER codes are the item-1 flags ``lowfreq``, ``sb`` and ``lowcov``.
63
+ CAT and SRVL are never written. Partial / gapped codons map to ``X``.
64
+ True amino-acid indels are not called.
65
+ """
66
+ dated = file_date or date.today().strftime("%Y%m%d")
67
+ fasta_name = Path(reference.path).name
68
+ contig_lines = [
69
+ f"##contig=<ID={record_id.replace(':', '|')},length={len(reference.sequences[record_id])}>"
70
+ for record_id in reference.record_ids
71
+ ]
72
+ header_lines = [
73
+ "##fileformat=AAVFv1.0",
74
+ f"##fileDate={dated}",
75
+ f"##source=codonyat-{PACKAGE_VERSION} (research use only)",
76
+ f"##reference={fasta_name}",
77
+ *contig_lines,
78
+ '##INFO=<ID=RC,Number=1,Type=String,Description="Reference codon">',
79
+ '##INFO=<ID=AC,Number=.,Type=String,Description="Alternate codon(s) encoding ALT">',
80
+ '##INFO=<ID=ACC,Number=.,Type=Integer,Description="Alternate codon counts, same order as AC">',
81
+ '##INFO=<ID=ACF,Number=.,Type=Float,Description="Alternate codon frequencies, same order as AC">',
82
+ '##INFO=<ID=FWC,Number=1,Type=Integer,Description="Forward-strand coverage at the codon">',
83
+ '##INFO=<ID=RVC,Number=1,Type=Integer,Description="Reverse-strand coverage at the codon">',
84
+ '##INFO=<ID=SR,Number=1,Type=Float,Description="Forward/reverse coverage ratio at the codon">',
85
+ '##FILTER=<ID=lowfreq,Description="ALT frequency below --min-freq">',
86
+ '##FILTER=<ID=sb,Description="Strand-bias Fisher p-value below --sb-alpha">',
87
+ '##FILTER=<ID=lowcov,Description="Codon coverage below --min-coverage">',
88
+ "#CHROM\tGENE\tPOS\tREF\tALT\tFILTER\tALT_FREQ\tCOVERAGE\tINFO",
89
+ ]
90
+
91
+ # Stream rows instead of accumulating the whole file in a list; output can
92
+ # be large when many variants are called and there is no need to hold it.
93
+ with open(output_path, "w") as handle:
94
+ handle.write("\n".join(header_lines) + "\n")
95
+
96
+ for protein_name, entries in protein_pairs:
97
+ # aa_pos then genomic position; keeps GENE blocks contiguous.
98
+ ordered = sorted(
99
+ entries.items(),
100
+ key=lambda item: (item[1].get("aa_pos", item[0]), item[0]),
101
+ )
102
+ for pos, entry in ordered:
103
+ depth = int(entry["depth"])
104
+ if depth <= 0:
105
+ continue
106
+ record_id = str(entry.get("record_id", reference.id)).replace(":", "|")
107
+ gene = protein_name.replace(":", "|")
108
+ aa_pos = int(entry.get("aa_pos", 0)) or "."
109
+ ref_codon = reference.get_seq_at(
110
+ pos, 3, entry.get("record_id"), positions=entry.get("codon_coords")
111
+ )
112
+ ref_aa = aavf_amino_acid(ref_codon) if ref_codon else "X"
113
+ fw_depth = int(entry["fw_depth"])
114
+ rv_depth = int(entry["rv_depth"])
115
+ sr = (fw_depth / rv_depth) if rv_depth else 0.0
116
+
117
+ by_aa: dict[str, list[Variant]] = defaultdict(list)
118
+ for variant in entry["variants"].values():
119
+ by_aa[aavf_amino_acid(variant.codon)].append(variant)
120
+
121
+ for alt_aa, variants in by_aa.items():
122
+ alt_count = sum(v.count for v in variants)
123
+ alt_fw = sum(v.fw_reads for v in variants)
124
+ alt_rv = sum(v.rv_reads for v in variants)
125
+ alt_freq = (alt_count / depth) if depth else 0.0
126
+ flags, _p = classify_flags(
127
+ freq=alt_freq,
128
+ coverage=depth,
129
+ fw=alt_fw,
130
+ rv=alt_rv,
131
+ other_fw=fw_depth - alt_fw,
132
+ other_rv=rv_depth - alt_rv,
133
+ min_freq=min_freq,
134
+ min_coverage=min_coverage,
135
+ sb_alpha=sb_alpha,
136
+ )
137
+ filter_field = ";".join(flags) if flags else "PASS"
138
+ ac_parts = []
139
+ acc_parts = []
140
+ acf_parts = []
141
+ for variant in variants:
142
+ codon = variant.codon
143
+ labelled = mark_changed_bases(ref_codon, codon) if is_acgt_codon(codon) else codon.lower()
144
+ ac_parts.append(labelled)
145
+ acc_parts.append(str(int(variant.count)))
146
+ acf_parts.append(_fmt_float(variant.count / depth if depth else 0.0))
147
+ info = (
148
+ f"RC={ref_codon.lower()};"
149
+ f"AC={','.join(ac_parts)};"
150
+ f"ACC={','.join(acc_parts)};"
151
+ f"ACF={','.join(acf_parts)};"
152
+ f"FWC={fw_depth};RVC={rv_depth};SR={_fmt_float(sr)}"
153
+ )
154
+ handle.write(
155
+ "\t".join(
156
+ [
157
+ record_id,
158
+ gene,
159
+ str(aa_pos),
160
+ ref_aa,
161
+ alt_aa,
162
+ filter_field,
163
+ _fmt_float(alt_freq),
164
+ str(depth),
165
+ info,
166
+ ]
167
+ )
168
+ + "\n"
169
+ )
@@ -20,6 +20,7 @@ from .models import Amplicon, Protein, Qual, Variant # noqa: F401
20
20
  from .reference import ( # noqa: F401
21
21
  FullReference,
22
22
  parse_protein_spec,
23
+ proteins_from_fasta_record,
23
24
  resolve_protein_targets,
24
25
  resolve_proteins,
25
26
  )
@@ -106,9 +106,8 @@ def main() -> None:
106
106
  ratio_lower=args.ratio_lower,
107
107
  entropy_threshold=args.entropy_threshold,
108
108
  min_baseq=args.min_baseq,
109
+ collapsed_read_counts=args.collapsed_read_counts,
109
110
  )
110
- container.load()
111
-
112
111
  logger.info("Calculating %s variant frequencies", ",".join(protein_names))
113
112
  container.calculate_for_proteins(proteins, reference)
114
113
 
@@ -9,7 +9,7 @@ from .constants import DEFAULT_MIN_COVERAGE, DEFAULT_MIN_FREQ, DEFAULT_SB_ALPHA
9
9
 
10
10
 
11
11
  def add_research_output_args(parser: ArgumentParser) -> None:
12
- """Attach opt-in 1.3.0 flags. Defaults leave TSV/XML unchanged."""
12
+ """Attach opt-in 1.3.x flags. Defaults leave TSV/XML unchanged."""
13
13
  parser.add_argument(
14
14
  "--min-baseq",
15
15
  type=float,
@@ -50,3 +50,12 @@ def add_research_output_args(parser: ArgumentParser) -> None:
50
50
  default=DEFAULT_SB_ALPHA,
51
51
  help="Fisher exact p-value cutoff for the sb (strand-bias) flag (default: 0.05).",
52
52
  )
53
+ parser.add_argument(
54
+ "--collapsed-read-counts",
55
+ action="store_true",
56
+ default=False,
57
+ help=(
58
+ "If set, treat a read name ending in _<digits> as a collapsed duplicate "
59
+ "count (that many observations). Off by default: each SAM row counts once."
60
+ ),
61
+ )
@@ -6,7 +6,7 @@ DEFAULT_ENTROPY_THRESHOLD = 0.0
6
6
  DEFAULT_PROTEIN = "RT"
7
7
 
8
8
  # Opt-in research statistics (v1.3.0). Unused unless --stats / --aavf / --min-baseq is set.
9
- PACKAGE_VERSION = "1.3.0"
9
+ PACKAGE_VERSION = "1.3.2"
10
10
  WILSON_Z_95 = 1.959963984540054 # Φ^{-1}(0.975)
11
11
  DEFAULT_MIN_FREQ = 0.01
12
12
  DEFAULT_MIN_COVERAGE = 10
@@ -4,7 +4,8 @@ import csv
4
4
  import logging
5
5
  import math
6
6
  import xml.etree.ElementTree as ET
7
- from collections.abc import Sequence
7
+ from bisect import bisect_left, bisect_right
8
+ from collections.abc import Iterator, Sequence
8
9
  from pathlib import Path
9
10
 
10
11
  from .genetic_code import codon_to_aminoacid
@@ -15,6 +16,65 @@ from .sam import SamEntry
15
16
  logger = logging.getLogger(__name__)
16
17
 
17
18
 
19
+ class _ReadCollection(Sequence[SamEntry]):
20
+ """Lazy, re-iterable view over the mapped records in a SAM file.
21
+
22
+ ``SamContainer.load`` used to materialise every mapped read in a list. Each
23
+ ``SamEntry`` cached a reference-position → read-position dict, so a 1 kb
24
+ read cost roughly 100 KB. With a few million reads that easily exhausts
25
+ memory. This view keeps only the path and streams the file on each pass;
26
+ ``read`` aggregation therefore uses O(1) memory in the sample size.
27
+
28
+ It still behaves as a read-only sequence (``len``, iteration, indexing) so
29
+ existing code that inspects ``SamContainer.reads`` keeps working.
30
+ """
31
+
32
+ def __init__(self, path: Path, collapsed_read_counts: bool = False) -> None:
33
+ self.path = Path(path)
34
+ self.collapsed_read_counts = collapsed_read_counts
35
+ self._count: int | None = None
36
+
37
+ def __iter__(self) -> Iterator[SamEntry]:
38
+ count = 0
39
+ with open(self.path) as fh:
40
+ for line in fh:
41
+ if line.startswith("@"):
42
+ continue
43
+ entry = SamEntry(line, collapsed_read_counts=self.collapsed_read_counts)
44
+ if entry.is_mapped():
45
+ count += 1
46
+ yield entry
47
+ # Cache the count after a complete pass so ``len()`` does not rescan.
48
+ self._count = count
49
+
50
+ def __len__(self) -> int:
51
+ if self._count is None:
52
+ # Count without building SamEntry objects (which would decode the
53
+ # CIGAR and sequence for every alignment just to report a length).
54
+ count = 0
55
+ with open(self.path) as fh:
56
+ for line in fh:
57
+ if line.startswith("@"):
58
+ continue
59
+ fields = line.split("\t", 2)
60
+ if not (int(fields[1]) & 4):
61
+ count += 1
62
+ self._count = count
63
+ return self._count
64
+
65
+ def __getitem__(self, index: int | slice):
66
+ if isinstance(index, slice):
67
+ return list(self)[index]
68
+ if index < 0:
69
+ index += len(self)
70
+ if index < 0:
71
+ raise IndexError("read index out of range")
72
+ for current, entry in enumerate(self):
73
+ if current == index:
74
+ return entry
75
+ raise IndexError("read index out of range")
76
+
77
+
18
78
  class SamContainer:
19
79
  """Aggregates SAM entries and computes RT variant statistics for an amplicon set."""
20
80
 
@@ -25,10 +85,14 @@ class SamContainer:
25
85
  ratio_lower: float,
26
86
  entropy_threshold: float,
27
87
  min_baseq: float | None = None,
88
+ collapsed_read_counts: bool = False,
28
89
  ) -> None:
29
90
  """Initialize the container with the SAM path, ratio window, and entropy threshold."""
30
91
  self.sam_path = sam_path
31
- self.reads: list[SamEntry] = []
92
+ # Streaming view: no alignment records are retained between passes.
93
+ self.reads: Sequence[SamEntry] = _ReadCollection(
94
+ sam_path, collapsed_read_counts=collapsed_read_counts
95
+ )
32
96
  self.variants: dict[int, dict] = {}
33
97
  self.protein_variants: dict[str, dict[int, dict]] = {}
34
98
  self._protein_result_list: list[tuple[str, dict[int, dict]]] = []
@@ -36,16 +100,16 @@ class SamContainer:
36
100
  self.ratio_lower = ratio_lower
37
101
  self.entropy_threshold = entropy_threshold
38
102
  self.min_baseq = min_baseq
103
+ self.collapsed_read_counts = collapsed_read_counts
39
104
 
40
105
  def load(self) -> None:
41
- """Load every mapped read from the SAM file so they can be aggregated."""
42
- with open(self.sam_path) as fh:
43
- for line in fh:
44
- if line.startswith("@"):
45
- continue
46
- entry = SamEntry(line)
47
- if entry.is_mapped():
48
- self.reads.append(entry)
106
+ """Prepare the streaming read view (kept for API compatibility).
107
+
108
+ Reads are parsed and aggregated in :meth:`calculate_for_proteins` rather
109
+ than retained. This method still validates the file by counting mapped
110
+ records without storing them, so ``len(container.reads)`` works as before.
111
+ """
112
+ len(self.reads)
49
113
 
50
114
  def calculate_variant_frequencies(
51
115
  self, protein_name: str, reference: FullReference, record_id: str | None = None
@@ -81,19 +145,28 @@ class SamContainer:
81
145
  self.calculate_for_proteins(proteins, reference)
82
146
 
83
147
  def calculate_for_proteins(self, proteins: Sequence[Protein], reference: FullReference) -> None:
84
- """Count every requested protein in one pass over ``self.reads``.
148
+ """Count every requested protein in one streaming pass over the SAM file.
85
149
 
86
150
  Supports multi-record references (reads are matched by RNAME) and
87
- multi-interval proteins.
151
+ multi-interval proteins. Only per-position aggregate state is retained;
152
+ alignment records are released as soon as they have been counted.
153
+
154
+ For each read we look up the codons whose coordinates lie inside the
155
+ read's aligned reference span instead of walking every codon of the
156
+ protein. That keeps both memory and CPU proportional to the covered
157
+ bases rather than to reads × protein length.
88
158
  """
89
159
  if not proteins:
90
160
  raise ValueError("At least one protein is required")
91
161
 
92
- jobs: list[tuple[Protein, dict[int, dict]]] = []
162
+ # Each job: (protein, entries, sorted codon starts, codon spans).
163
+ jobs: list[
164
+ tuple[Protein, dict[int, dict], list[int], list[tuple[int, int]]]
165
+ ] = []
93
166
  for protein in proteins:
94
167
  entries: dict[int, dict] = {}
95
168
  aa_pos = 0
96
- for pos, record_id in self._get_codon_positions(protein, reference):
169
+ for pos, record_id, codon_coords in self._get_codon_positions(protein, reference):
97
170
  aa_pos += 1
98
171
  entries[pos] = {
99
172
  "variants": {},
@@ -103,19 +176,46 @@ class SamContainer:
103
176
  "amplicons": {},
104
177
  "record_id": record_id,
105
178
  "aa_pos": aa_pos,
179
+ "codon_coords": codon_coords,
106
180
  }
107
- jobs.append((protein, entries))
181
+ starts = sorted(entries)
182
+ spans = [
183
+ (
184
+ min(entries[pos]["codon_coords"]),
185
+ max(entries[pos]["codon_coords"]),
186
+ )
187
+ for pos in starts
188
+ ]
189
+ jobs.append((protein, entries, starts, spans))
190
+
191
+ jobs_by_record: dict[str, list[tuple[Protein, dict[int, dict], list[int], list[tuple[int, int]]]]] = {}
192
+ for job in jobs:
193
+ jobs_by_record.setdefault(job[0].record_id, []).append(job)
108
194
 
109
195
  for read in self.reads:
110
- for protein, entries in jobs:
111
- if read.reference != protein.record_id:
112
- continue
113
- for pos, entry in entries.items():
114
- self._accumulate_read_at_position(read, pos, entry)
196
+ if not read.is_mapped():
197
+ continue
198
+ read_jobs = jobs_by_record.get(read.reference)
199
+ if not read_jobs:
200
+ continue
201
+ span_start = read.coordinate
202
+ span_end = read.reference_end
203
+ for _protein, entries, starts, spans in read_jobs:
204
+ # Codons entirely outside the read cannot contribute. Using the
205
+ # first codon coordinate as the sorted key is sufficient: every
206
+ # other coordinate is checked before counting.
207
+ left = bisect_left(starts, span_start)
208
+ right = bisect_right(starts, span_end)
209
+ for index in range(left, right):
210
+ codon_min, codon_max = spans[index]
211
+ if codon_min < span_start or codon_max > span_end:
212
+ continue
213
+ pos = starts[index]
214
+ self._accumulate_read_at_position(read, pos, entries[pos])
115
215
 
116
216
  self.protein_variants = {}
117
217
  self._protein_result_list = []
118
- for protein, entries in jobs:
218
+ for protein, entries, _starts, _spans in jobs:
119
219
  self.variants = entries
120
220
  for pos in self.variants:
121
221
  self.calculate_ratios_nucleotide(pos)
@@ -128,13 +228,16 @@ class SamContainer:
128
228
 
129
229
  def _accumulate_read_at_position(self, read: SamEntry, pos: int, entry: dict) -> None:
130
230
  """Add one mapped read's codon to a protein position, if it covers it."""
131
- if not read.covers(pos):
132
- return
133
- codon = read.codon_at(pos)
231
+ codon_coords = entry.get("codon_coords")
232
+ if codon_coords is None:
233
+ codon_coords = (pos, pos + 1, pos + 2)
234
+ # ``codon_at`` returns None when any of the three coordinates is outside
235
+ # the read's aligned span, so a separate covers() sweep is unnecessary.
236
+ codon = read.codon_at(pos, positions=codon_coords)
134
237
  if not codon or len(codon) < 3:
135
238
  return
136
239
  if self.min_baseq is not None:
137
- quals = read.codon_qualities_at(pos)
240
+ quals = read.codon_qualities_at(pos, positions=codon_coords)
138
241
  if quals is None or any(qual < self.min_baseq for qual in quals):
139
242
  return
140
243
  depth_increment = read.occurences
@@ -166,27 +269,22 @@ class SamContainer:
166
269
  amp_position["depth"] += depth_increment
167
270
 
168
271
  def _get_codon_positions(self, protein, reference: FullReference) -> list[tuple]:
169
- """Generate list of (position, record_id) tuples for all codons in the protein.
272
+ """Generate (start, record_id, codon_coords) for every codon in the protein.
170
273
 
171
- Handles multi-interval proteins including frameshifts.
274
+ Coding bases are concatenated in interval order so the reading frame is
275
+ carried across parts. A codon that sits on a junction keeps the actual
276
+ part coordinates (for example positions 10, 11, 12 of 1-10;11-30).
277
+ Incomplete trailing bases at the end of the protein are dropped.
172
278
  """
173
- positions = []
174
- current_offset = 0
175
-
279
+ positions: list[tuple] = []
280
+ pending: list[int] = []
176
281
  for interval in protein.intervals:
177
- interval_length = interval.end - interval.start + 1
178
- pos = interval.start
179
-
180
- while pos <= interval.end:
181
- remaining_in_interval = interval.end - pos + 1
182
- if remaining_in_interval >= 3:
183
- positions.append((pos, protein.record_id))
184
- pos += 3
185
- else:
186
- break
187
-
188
- current_offset += interval_length
189
-
282
+ for coordinate in range(interval.start, interval.end + 1):
283
+ pending.append(coordinate)
284
+ if len(pending) == 3:
285
+ codon_coords = (pending[0], pending[1], pending[2])
286
+ positions.append((codon_coords[0], protein.record_id, codon_coords))
287
+ pending.clear()
190
288
  return positions
191
289
 
192
290
  def calculate_ratios_nucleotide(self, position: int) -> None:
@@ -361,26 +459,73 @@ class SamContainer:
361
459
  return pairs
362
460
 
363
461
  @staticmethod
364
- def _add_positions(parent: ET.Element, entries: dict[int, dict]) -> None:
462
+ def _position_element(pos: int, entry: dict) -> ET.Element:
463
+ position_elem = ET.Element("Position", index=str(pos))
464
+ ET.SubElement(position_elem, "Depth").text = str(entry["depth"])
465
+ ET.SubElement(position_elem, "FwCover").text = str(entry["fw_depth"])
466
+ ET.SubElement(position_elem, "RvCover").text = str(entry["rv_depth"])
467
+ variants_elem = ET.SubElement(position_elem, "Variants")
468
+ for variant in entry["variants"].values():
469
+ # describe each variant using XML attributes for diagnostics
470
+ var_elem = ET.SubElement(variants_elem, "Variant", codon=variant.codon)
471
+ var_elem.set("count", str(variant.count))
472
+ var_elem.set("fw_reads", str(variant.fw_reads))
473
+ var_elem.set("rv_reads", str(variant.rv_reads))
474
+ return position_elem
475
+
476
+ @classmethod
477
+ def _iter_position_xml(cls, entries: dict[int, dict]) -> Iterator[str]:
478
+ """Serialise one position at a time so the full XML tree is never held."""
365
479
  for pos, entry in sorted(entries.items()):
366
- position_elem = ET.SubElement(parent, "Position", index=str(pos))
367
- ET.SubElement(position_elem, "Depth").text = str(entry["depth"])
368
- ET.SubElement(position_elem, "FwCover").text = str(entry["fw_depth"])
369
- ET.SubElement(position_elem, "RvCover").text = str(entry["rv_depth"])
370
- variants_elem = ET.SubElement(position_elem, "Variants")
371
- for variant in entry["variants"].values():
372
- # describe each variant using XML attributes for diagnostics
373
- var_elem = ET.SubElement(variants_elem, "Variant", codon=variant.codon)
374
- var_elem.set("count", str(variant.count))
375
- var_elem.set("fw_reads", str(variant.fw_reads))
376
- var_elem.set("rv_reads", str(variant.rv_reads))
480
+ yield ET.tostring(
481
+ cls._position_element(pos, entry),
482
+ encoding="unicode",
483
+ short_empty_elements=True,
484
+ )
485
+
486
+ @staticmethod
487
+ def _open_tag(tag: str, **attributes: str) -> str:
488
+ """Return the opening tag for an empty element, in ElementTree style."""
489
+ empty = ET.tostring(
490
+ ET.Element(tag, attributes), encoding="unicode", short_empty_elements=True
491
+ )
492
+ # ElementTree serialises an empty element as ``<tag ... />``.
493
+ return empty[:-3] + ">"
494
+
495
+ def _write_xml_tree(
496
+ self,
497
+ sample_name: str,
498
+ reference: FullReference,
499
+ output_path: Path,
500
+ protein_pairs: Sequence[tuple[str, dict[int, dict]]],
501
+ *,
502
+ wrap_proteins: bool,
503
+ ) -> None:
504
+ with open(output_path, "w", encoding="utf-8", newline="") as handle:
505
+ handle.write("<?xml version='1.0' encoding='utf-8'?>\n")
506
+ handle.write(
507
+ self._open_tag(
508
+ "SamContainer", sample=sample_name, reference=reference.id
509
+ )
510
+ )
511
+ for protein_name, entries in protein_pairs:
512
+ if wrap_proteins:
513
+ handle.write(self._open_tag("Protein", name=protein_name))
514
+ for chunk in self._iter_position_xml(entries):
515
+ handle.write(chunk)
516
+ if wrap_proteins:
517
+ handle.write("</Protein>")
518
+ handle.write("</SamContainer>")
377
519
 
378
520
  def write_xml(self, sample_name: str, reference: FullReference, output_path: Path) -> None:
379
521
  """Dump the per-position state of the (first) protein into XML for diagnostics."""
380
- root = ET.Element("SamContainer", sample=sample_name, reference=reference.id)
381
- self._add_positions(root, self.variants)
382
- tree = ET.ElementTree(root)
383
- tree.write(output_path, encoding="utf-8", xml_declaration=True)
522
+ self._write_xml_tree(
523
+ sample_name,
524
+ reference,
525
+ output_path,
526
+ [("", self.variants)],
527
+ wrap_proteins=False,
528
+ )
384
529
 
385
530
  def write_xml_proteins(
386
531
  self,
@@ -393,12 +538,13 @@ class SamContainer:
393
538
 
394
539
  For a single protein use :meth:`write_xml`, whose layout is unchanged.
395
540
  """
396
- root = ET.Element("SamContainer", sample=sample_name, reference=reference.id)
397
- for protein_name, entries in self._output_pairs(protein_names):
398
- protein_elem = ET.SubElement(root, "Protein", name=protein_name)
399
- self._add_positions(protein_elem, entries)
400
- tree = ET.ElementTree(root)
401
- tree.write(output_path, encoding="utf-8", xml_declaration=True)
541
+ self._write_xml_tree(
542
+ sample_name,
543
+ reference,
544
+ output_path,
545
+ self._output_pairs(protein_names),
546
+ wrap_proteins=True,
547
+ )
402
548
 
403
549
  def write_stats(
404
550
  self,