Skills › Scientific Agent Skills › genomic-coordinates › _common.py
_common.py
355 lines · skills/genomic-coordinates/scripts/_common.py at 5f2bc03
1"""Shared helpers for the genomic-coordinates scripts.
3Everything here is standard library only. The single organising idea is that all
4intervals are converted to one canonical form on the way in and back out again on
5the way out, so no script ever has to reason about two conventions at once.
7Canonical form: ``(start0, end_exclusive)`` -- 0-based, half-open, the same
8convention BED and Python slices use. ``end0 - start0`` is always the length, and
9a zero-length interval (an insertion point between two bases) is representable.
10"""
12from __future__ import annotations
14import json
15import re
16import sys
17from dataclasses import dataclass
18from pathlib import Path
20# --------------------------------------------------------------------------
21# coordinate conventions
22# --------------------------------------------------------------------------
25@dataclass(frozen=True)
26class Convention:
27 """How one file format or API writes an interval down."""
29 name: str
30 base: int # 0 or 1: what the first base of a contig is called
31 half_open: bool # True: end is exclusive. False: end is inclusive.
32 note: str = ""
34 @property
35 def label(self) -> str:
36 closure = "half-open" if self.half_open else "inclusive"
37 return f"{self.base}-based {closure}"
40# Every entry is sourced from the format's own specification; see
41# references/format-conventions.md for the citation behind each one.
42CONVENTIONS: dict[str, Convention] = {
43 "bed": Convention("bed", 0, True, "BED3/6/12, narrowPeak, broadPeak"),
44 "bedgraph": Convention("bedgraph", 0, True, "bedGraph, and bigWig internals"),
45 "gff": Convention("gff", 1, False, "GFF3"),
46 "gff3": Convention("gff3", 1, False, "GFF3"),
47 "gtf": Convention("gtf", 1, False, "GTF/GFF2, GENCODE, Ensembl"),
48 "vcf": Convention("vcf", 1, False, "POS..POS+len(REF)-1"),
49 "sam": Convention("sam", 1, False, "SAM text POS (BAM stores it 0-based)"),
50 "wig": Convention("wig", 1, False, "fixedStep/variableStep"),
51 "psl": Convention("psl", 0, True, "BLAT"),
52 "genepred": Convention("genepred", 0, True, "genePred, refFlat, UCSC tables"),
53 "interval-list": Convention("interval-list", 1, False, "Picard/GATK"),
54 "maf-tcga": Convention("maf-tcga", 1, False, "Mutation Annotation Format"),
55 "maf-ucsc": Convention("maf-ucsc", 0, True, "Multiple Alignment Format"),
56 "ucsc": Convention("ucsc", 1, False, "browser position box, region string"),
57 "ensembl": Convention("ensembl", 1, False, "REST region string"),
58 "samtools": Convention("samtools", 1, False, "samtools/tabix region string"),
59 "igv": Convention("igv", 1, False, "IGV locus box"),
60 "granges": Convention("granges", 1, False, "Bioconductor GRanges/IRanges"),
61 "pyranges": Convention("pyranges", 0, True, "PyRanges, pybedtools"),
62 "python": Convention("python", 0, True, "list/str slice semantics"),
63}
65# Formats whose canonical text form is "contig:start-end" rather than columns.
66REGION_STRING_FORMATS = {"ucsc", "ensembl", "samtools", "igv"}
69def get_convention(name: str) -> Convention:
70 key = name.strip().lower()
71 if key not in CONVENTIONS:
72 known = ", ".join(sorted(CONVENTIONS))
73 raise ValueError(f"unknown coordinate format {name!r}; known formats: {known}")
74 return CONVENTIONS[key]
77def to_canonical(conv: Convention, start: int, end: int) -> tuple[int, int]:
78 """Convert an interval written in ``conv`` to 0-based half-open."""
79 start0 = start - conv.base
80 end0 = end if conv.half_open else end - conv.base + 1
81 return start0, end0
84def from_canonical(conv: Convention, start0: int, end0: int) -> tuple[int, int]:
85 """Convert a 0-based half-open interval into ``conv``.
87 A zero-length interval yields ``end < start`` in every inclusive convention.
88 That is arithmetically correct and usually means the interval should not have
89 been converted at all; callers are expected to surface it rather than hide it.
90 """
91 start = start0 + conv.base
92 end = end0 if conv.half_open else end0 + conv.base - 1
93 return start, end
96# --------------------------------------------------------------------------
97# region strings
98# --------------------------------------------------------------------------
100# A contig name is either brace-quoted -- htslib's escape for GRCh38 names that
101# contain a colon, such as HLA-DRB1*12:17 -- or runs up to the first colon.
102_REGION_RE = re.compile(
103 r"^\s*(?:\{(?P<braced>[^}]+)\}|(?P<contig>[^\s:]+))"
104 r"(?::(?P<start>[\d,_]+)"
105 r"(?:\s*(?:-|\.\.)\s*(?P<end>[\d,_]+))?)?"
106 r"(?::(?P<strand>-1|1|[-+.]))?\s*$"
107)
110def parse_region(text: str, conv: Convention) -> tuple[str, int, int, str | None]:
111 """Parse ``chr1:1,000-2,000`` into ``(contig, start, end, strand)``.
113 Numbers come back in the convention's own coordinates, unconverted.
115 Neither a bare contig nor a bare ``contig:start`` is accepted. Both mean
116 "to the end of the contig" in samtools, which cannot be converted without a
117 genome file, and inventing an end is exactly the class of bug this skill
118 exists to stop.
119 """
120 stripped = text.strip()
121 ambiguous = ValueError(
122 f"contig name in {text!r} is ambiguous: GRCh38 ALT contigs such as "
123 "HLA-DRB1*12:17 contain colons, so there is no way to tell the name from "
124 "the coordinates. Brace-quote it, as htslib does: {HLA-DRB1*12:17}:100-200"
125 )
126 # Checked before parsing: a '*' in an unbraced name means the split point
127 # cannot be inferred, whether or not the rest happens to parse.
128 if not stripped.startswith("{") and "*" in stripped:
129 raise ambiguous
131 m = _REGION_RE.match(text)
132 if not m:
133 if stripped.count(":") > 1:
134 raise ambiguous
135 raise ValueError(f"cannot parse region string {text!r}")
136 contig = m.group("braced") or m.group("contig")
137 if not m.group("braced"):
138 allowed = 2 if m.group("strand") else 1
139 if stripped.count(":") > allowed:
140 raise ambiguous
141 if m.group("start") is None:
142 raise ValueError(
143 f"region {text!r} names a whole contig; give contig:start-end"
144 )
145 start = int(m.group("start").replace(",", "").replace("_", ""))
146 if m.group("end") is None:
147 raise ValueError(
148 f"region {text!r} has a start but no end. In samtools and tabix that "
149 "means 'from here to the end of the contig', not a single base -- "
150 "write the end explicitly"
151 )
152 end = int(m.group("end").replace(",", "").replace("_", ""))
153 strand = m.group("strand")
154 if conv.name == "ensembl" and strand not in {None, "1", "-1", "+", "-"}:
155 raise ValueError("Ensembl REST strand must be 1 or -1")
156 return contig, start, end, strand
159def format_region(contig: str, start: int, end: int, strand: str | None = None,
160 convention: str = "samtools") -> str:
161 """Render a region string, brace-quoting names that would otherwise be
162 ambiguous (GRCh38 HLA contigs contain colons)."""
163 if convention == "ensembl":
164 direction = {"+": "1", "-": "-1", "1": "1", "-1": "-1"}.get(strand)
165 text = f"{contig}:{start}..{end}"
166 return f"{text}:{direction}" if direction else text
167 name = f"{{{contig}}}" if ":" in contig else contig
168 text = f"{name}:{start}-{end}"
169 return f"{text}:{strand}" if strand else text
172# --------------------------------------------------------------------------
173# contig naming
174# --------------------------------------------------------------------------
176_ROMAN = re.compile(r"^(chr)?([IVXL]+)$", re.IGNORECASE)
179def naming_style(names: list[str]) -> str:
180 """Classify a set of contig names: how a join against another file will fail."""
181 if not names:
182 return "empty"
183 if any(re.match(r"^(NC|NT|NW|GL|KI|CM|GCA|GCF)_?\d", n) for n in names):
184 return "accession"
185 prefixed = sum(1 for n in names if n.startswith("chr"))
186 if prefixed == len(names):
187 return "chr-prefixed"
188 if prefixed == 0:
189 return "plain"
190 return "mixed"
193def strip_chr(name: str) -> str:
194 return name[3:] if name.startswith("chr") else name
197def canonical_contig(name: str) -> str:
198 """Fold ``chr1``/``1`` and ``chrM``/``MT``/``chrMT`` onto one key.
200 For comparing contig *sets* across files only. Never write the result into a
201 data file -- the naming style of the file you are writing has to be preserved.
202 """
203 bare = strip_chr(name).upper()
204 if bare in {"M", "MT"}:
205 return "M"
206 return bare
209# --------------------------------------------------------------------------
210# reference sequence
211# --------------------------------------------------------------------------
214class ReferenceError(Exception):
215 """The reference FASTA could not answer the question that was asked."""
218class Reference:
219 """Random access to a FASTA, through its ``.fai`` index when one exists.
221 Supports uncompressed FASTA only. With an index, only the requested bases
222 are read. Without one the whole FASTA is loaded into memory at construction.
223 Contig names must match exactly; a naming alias is not sequence identity.
224 """
226 def __init__(self, path: str | Path) -> None:
227 self.path = Path(path)
228 if not self.path.exists():
229 raise ReferenceError(f"reference not found: {self.path}")
230 self._index: dict[str, tuple[int, int, int, int]] = {}
231 self._cache: dict[str, str] = {}
232 with self.path.open("rb") as handle:
233 if handle.read(2) == b"\x1f\x8b":
234 raise ReferenceError("compressed FASTA is unsupported; use an uncompressed FASTA")
235 fai = Path(str(self.path) + ".fai")
236 if fai.exists():
237 self._load_fai(fai)
238 else:
239 self._load_all()
241 def _load_fai(self, fai: Path) -> None:
242 for line in fai.read_text().splitlines():
243 if not line.strip():
244 continue
245 parts = line.split("\t")
246 if len(parts) < 5:
247 raise ReferenceError(f"malformed .fai line: {line!r}")
248 name, length, offset, linebases, linewidth = parts[:5]
249 try:
250 meta = tuple(map(int, (length, offset, linebases, linewidth)))
251 except ValueError as exc:
252 raise ReferenceError(f"non-integer .fai fields for {name}") from exc
253 if name in self._index or meta[0] <= 0 or meta[1] < 0 or meta[2] <= 0 or meta[3] < meta[2]:
254 raise ReferenceError(f"duplicate name or invalid .fai geometry for {name}")
255 self._index[name] = meta
257 def _load_all(self) -> None:
258 name = None
259 chunks: list[str] = []
260 with self.path.open() as handle:
261 for line in handle:
262 if line.startswith(">"):
263 if name is not None:
264 self._cache[name] = "".join(chunks)
265 names = line[1:].split()
266 if not names or names[0] in self._cache:
267 raise ReferenceError("empty or duplicate FASTA sequence name")
268 name = names[0]
269 chunks = []
270 elif name is not None:
271 chunks.append(line.strip())
272 if name is not None:
273 self._cache[name] = "".join(chunks)
274 if not self._cache:
275 raise ReferenceError("FASTA contains no sequences")
277 def _resolve(self, contig: str) -> str:
278 if contig in self._index or contig in self._cache:
279 return contig
280 available = sorted(set(self._index) | set(self._cache))
281 shown = ", ".join(available[:8]) + ("..." if len(available) > 8 else "")
282 raise ReferenceError(
283 f"contig {contig!r} is not in {self.path.name} (has: {shown}). "
284 "Exact names are required. Check aliases and assembly with check_contigs.py "
285 "before explicitly renaming inputs."
286 )
288 @property
289 def contigs(self) -> dict[str, int]:
290 if self._index:
291 return {name: meta[0] for name, meta in self._index.items()}
292 return {name: len(seq) for name, seq in self._cache.items()}
294 def length(self, contig: str) -> int:
295 return self.contigs[self._resolve(contig)]
297 def fetch(self, contig: str, start0: int, end0: int) -> str:
298 """Return ``[start0, end0)`` in upper case. Out-of-range is an error."""
299 name = self._resolve(contig)
300 if start0 < 0:
301 raise ReferenceError(f"negative start {start0} on {contig}")
302 size = self.length(name)
303 if end0 > size:
304 raise ReferenceError(
305 f"{contig}:{start0}-{end0} runs past the end of {name} (length {size}); "
306 "the coordinates and this reference are not the same assembly"
307 )
308 if end0 <= start0:
309 return ""
310 if name in self._cache:
311 return self._cache[name][start0:end0].upper()
312 _, offset, linebases, linewidth = self._index[name]
313 with self.path.open("rb") as handle:
314 begin = offset + (start0 // linebases) * linewidth + (start0 % linebases)
315 stop = offset + (end0 // linebases) * linewidth + (end0 % linebases)
316 handle.seek(begin)
317 raw = handle.read(stop - begin)
318 sequence = re.sub(rb"\s", b"", raw).decode("ascii").upper()
319 if len(sequence) != end0 - start0 or set(sequence) - set("ACGTRYSWKMBDHVN"):
320 raise ReferenceError("FASTA/index disagree; rebuild the .fai from this FASTA")
321 return sequence
324# --------------------------------------------------------------------------
325# output
326# --------------------------------------------------------------------------
329def emit(rows: list[dict], columns: list[str], fmt: str, out: str | None) -> None:
330 """Write ``rows`` as TSV or JSON to a path or stdout."""
331 if fmt == "json":
332 text = json.dumps(rows, indent=2)
333 else:
334 lines = ["\t".join(columns)]
335 for row in rows:
336 lines.append("\t".join(str(row.get(col, "")) for col in columns))
337 text = "\n".join(lines)
338 if out:
339 Path(out).write_text(text + "\n")
340 else:
341 sys.stdout.write(text + "\n")
344def iter_data_lines(path: str | Path):
345 """Yield ``(lineno, line)`` for non-blank, non-comment lines of a text file."""
346 with Path(path).open() as handle:
347 for lineno, line in enumerate(handle, start=1):
348 stripped = line.rstrip("\n\r")
349 if stripped == "##FASTA":
350 break
351 if not stripped.strip():
352 continue
353 if stripped.startswith(("#", "track ", "browser ", "@")):
354 continue
355 yield lineno, stripped