cstag.to_vcf

  1from __future__ import annotations
  2
  3import re
  4from collections import Counter, defaultdict, deque
  5from collections.abc import Iterable
  6from dataclasses import dataclass
  7from itertools import chain
  8from typing import cast, overload
  9
 10from .consensus import normalize_read_lengths
 11from .split import split
 12from .utils.validator import validate_cs_tag, validate_long_format, validate_pos
 13
 14
 15@dataclass(frozen=True)
 16class CsInfo:
 17    cs_tag: str
 18    pos_start: int
 19    pos_end: int
 20    chrom: str | None = None
 21
 22
 23@dataclass(frozen=True)
 24class VcfInfo:
 25    dp: int | None = None
 26    rd: int | None = None
 27    ad: int | None = None
 28    vaf: float | None = None
 29
 30
 31@dataclass(frozen=True)
 32class Vcf:
 33    chrom: str | None = None
 34    pos: int | None = None
 35    ref: str | None = None
 36    alt: str | None = None
 37    info: VcfInfo = VcfInfo()
 38
 39
 40def remove_spaces_around_newlines(text: str) -> str:
 41    return re.sub(r"\s*\n\s*", "\n", text)
 42
 43
 44###########################################################
 45# Get variant annotations
 46###########################################################
 47
 48
 49def find_ref_for_insertion(cs_tag_split: list[str], idx: int) -> str | None:
 50    idx_ref = idx - 1
 51    while idx_ref >= 0:
 52        cs = cs_tag_split[idx_ref]
 53        if cs[0] in ["=", "-"]:
 54            return cs[-1].upper()
 55        if cs.startswith("*"):
 56            return cs[1].upper()
 57        idx_ref -= 1
 58    return None
 59
 60
 61def find_ref_for_deletion(cs_tag_split: list[str], idx: int) -> str:
 62    ref = deque([cs_tag_split[idx][1:].upper()])
 63    idx_ref = idx - 1
 64    while idx_ref >= 0:
 65        cs = cs_tag_split[idx_ref]
 66        if cs.startswith("="):
 67            ref.appendleft(cs[-1].upper())
 68            break
 69        if cs.startswith("*"):
 70            ref.appendleft(cs[1].upper())
 71            break
 72        idx_ref -= 1
 73    return "".join(ref)
 74
 75
 76def get_variant_annotations(cs_tag_split: list[str], position: int) -> list[Vcf]:
 77    variant_annotations: list[Vcf] = []
 78    pos = position
 79    for idx, cs in enumerate(cs_tag_split):
 80        if cs.startswith("="):
 81            pos += len(cs) - 1
 82        elif cs.startswith("*"):
 83            ref, alt = cs[1].upper(), cs[2].upper()
 84            variant_annotations.append(Vcf(pos=pos, ref=ref, alt=alt))
 85            pos += 1
 86        elif cs.startswith("+"):
 87            ref = cast(str, find_ref_for_insertion(cs_tag_split, idx))
 88            alt = ref + cs[1:].upper()
 89            variant_annotations.append(Vcf(pos=pos - 1, ref=ref, alt=alt))
 90        elif cs.startswith("-"):
 91            ref = find_ref_for_deletion(cs_tag_split, idx)
 92            variant_annotations.append(Vcf(pos=pos - 1, ref=ref, alt=ref[0]))
 93        elif cs.startswith("~"):
 94            continue
 95
 96    return variant_annotations
 97
 98
 99###########################################################
100# Format the cs tags
101###########################################################
102
103
104def get_pos_end(cs_tag: str, pos: int) -> int:
105    """Get 1-index end positions"""
106    pos_end = pos - 1
107    for cs in split(cs_tag):
108        if cs[0] in ["=", "-"]:
109            pos_end += len(cs) - 1
110        if cs[0] == "*":
111            pos_end += 1
112        else:
113            continue
114    return pos_end
115
116
117def format_cs_tags(
118    cs_tags: list[str],
119    chroms: list[str] | list[int],
120    positions: list[int],
121) -> list[CsInfo]:
122    """Format and filter cs_tags, and create a list of CsInfo objects.
123
124    This function takes lists of cs_tags, chromosomes, and positions. It filters
125    out any cs_tags containing a tilde ("~") and creates a list of CsInfo objects.
126
127    Args:
128        cs_tags (list[str]): List of cs_tags as strings.
129        chroms (list[str] | list[int]): List of chromosomes as strings or integers.
130        positions (list[int]): List of starting positions as integers.
131
132    Returns:
133        list[CsInfo]: A list of CsInfo objects, each containing information about
134        a cs_tag, its chromosome, and its start and end positions.
135    """
136
137    # Convert all chromosomes to string type
138    chrom_strings = [str(chrom) for chrom in chroms]
139    # Create a list of CsInfo objects, filtering out any with a splicing ("~") in the cs_tag
140    cs_info_list = [
141        CsInfo(cs_tag=cs, chrom=chrom, pos_start=pos, pos_end=get_pos_end(cs, pos))
142        for cs, chrom, pos in zip(cs_tags, chrom_strings, positions, strict=False)
143        if "~" not in cs
144    ]
145    return cs_info_list
146
147
148###########################################################
149# Group by chrom and overlapping intervals
150###########################################################
151
152
153def group_by_chrom(
154    cs_tags_formatted: list[CsInfo],
155) -> dict[str | None, list[CsInfo]]:
156    """Group cs tags by chromosomes"""
157    cs_tags_grouped: defaultdict[str | None, list[CsInfo]] = defaultdict(list)
158    for cs in cs_tags_formatted:
159        cs_tags_grouped[cs.chrom].append(
160            CsInfo(
161                cs_tag=cs.cs_tag,
162                pos_start=cs.pos_start,
163                pos_end=cs.pos_end,
164                chrom=cs.chrom,
165            )
166        )
167    return dict(cs_tags_grouped)
168
169
170def group_by_overlapping_intervals(
171    cs_tags_grouped: list[CsInfo],
172) -> list[list[CsInfo]]:
173    # Sort the list by the starting point
174    sorted_data = sorted(cs_tags_grouped, key=lambda x: x.pos_start)
175    # Initialize the list of grouped intervals
176    grouped_intervals: list[list[CsInfo]] = []
177    # Initialize the first group with the first element
178    current_group = [sorted_data[0]]
179    # Loop through the sorted list starting from the second element
180    for i in range(1, len(sorted_data)):
181        overlaps = False
182        for j in current_group:
183            # Check if the intervals overlap
184            if (
185                sorted_data[i].pos_start <= j.pos_end
186                and sorted_data[i].pos_end >= j.pos_start
187            ):
188                overlaps = True
189                break
190        if overlaps:
191            # Add the interval to the current group
192            current_group.append(sorted_data[i])
193        else:
194            # Add the current group to the list of grouped intervals
195            grouped_intervals.append(current_group)
196            # Start a new group
197            current_group = [sorted_data[i]]
198    # Add the last group to the list of grouped intervals
199    grouped_intervals.append(current_group)
200
201    return grouped_intervals
202
203
204###########################################################
205# Add VCF info
206###########################################################
207
208
209def replace_mutation_to_atmark(cs_tags: Iterable[str | None]) -> str:
210    """Replaces mutations with '@'."""
211    return "".join(cs if cs in {"A", "C", "G", "T"} else "@" for cs in cs_tags)
212
213
214def call_reference_depth(
215    variant_annotations: list[Vcf],
216    cs_tags_list: list[str],
217    positions_list: list[int],
218) -> dict[tuple[str, int], int]:
219    cs_tags_normalized_length = normalize_read_lengths(cs_tags_list, positions_list)
220    cs_replaced = [
221        replace_mutation_to_atmark(cs_tags) for cs_tags in cs_tags_normalized_length
222    ]
223
224    reference_depth: defaultdict[tuple[str, int], int] = defaultdict(int)
225    unique_variants = set(variant_annotations)
226    for v in unique_variants:
227        variant_pos = cast(int, v.pos)
228        variant_ref = cast(str, v.ref)
229        v_idx = variant_pos - min(positions_list)
230        for cs in cs_replaced:
231            if variant_ref == cs[v_idx : v_idx + len(variant_ref)]:
232                reference_depth[(variant_ref, variant_pos)] += 1
233
234    return dict(reference_depth)
235
236
237def add_vcf_fields(
238    variant_annotations: list[Vcf],
239    chrom: str,
240    reference_depth: dict[tuple[str, int], int],
241) -> list[Vcf]:
242    """Add Chrom and VCF info (AD, RD, DP, and VAF) to immutable Vcf dataclass"""
243    variant_counter = Counter((v.pos, v.ref, v.alt) for v in variant_annotations)
244
245    updated_annotations: list[Vcf] = []
246    for v in set(variant_annotations):
247        ad = variant_counter[(v.pos, v.ref, v.alt)]
248        rd = reference_depth.get((cast(str, v.ref), cast(int, v.pos)), 0)
249        dp = rd + ad
250        vaf = round(ad / dp, 3) if dp else 0
251
252        # Creating a new VcfInfo object
253        updated_info = VcfInfo(dp=dp, rd=rd, ad=ad, vaf=vaf)
254
255        # Creating a new Vcf object
256        updated_variant = Vcf(
257            chrom=chrom, pos=v.pos, ref=v.ref, alt=v.alt, info=updated_info
258        )
259
260        updated_annotations.append(updated_variant)
261
262    return updated_annotations
263
264
265###########################################################
266# Process cs tag (One)
267###########################################################
268
269
270def process_cs_tag(cs_tag: str, chrom: str | int, pos: int) -> str:
271    validate_cs_tag(cs_tag)
272    validate_long_format(cs_tag)
273    validate_pos(pos)
274    chrom = str(chrom)
275
276    cs_tag_split = split(cs_tag)
277
278    # Call POS, REF, ALT
279    variants = get_variant_annotations(cs_tag_split, pos)
280
281    # Write VCF
282    HEADER = "##fileformat=VCFv4.2\n#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n"
283    vcf = remove_spaces_around_newlines(HEADER).strip().split("\n")
284    for v in variants:
285        vcf.append(f"{chrom}\t{v.pos}\t.\t{v.ref}\t{v.alt}\t.\t.\t.")
286
287    return "\n".join(vcf)
288
289
290###########################################################
291# Process cs tags (Many)
292###########################################################
293
294
295def chrom_sort_key(chrom: str) -> int:
296    """Convert a chromosome string to an integer for sorting."""
297    return int(chrom.replace("chr", ""))
298
299
300def process_cs_tags(
301    cs_tags: list[str],
302    chroms: list[str] | list[int],
303    positions: list[int],
304) -> str:
305    # validate inputs
306    for cs_tag in cs_tags:
307        validate_cs_tag(cs_tag)
308        validate_long_format(cs_tag)
309    for pos in positions:
310        validate_pos(pos)
311
312    cs_tags_formatted = format_cs_tags(cs_tags, chroms, positions)
313    cs_tags_grouped_by_chrom = group_by_chrom(cs_tags_formatted)
314
315    vcf_info: list[Vcf] = []
316    for maybe_chrom, cs_tags_grouped in cs_tags_grouped_by_chrom.items():
317        chrom = cast(str, maybe_chrom)
318        for csinfo in group_by_overlapping_intervals(cs_tags_grouped):
319            cs_tags_list = [cs.cs_tag for cs in csinfo]
320            positions_list = [cs.pos_start for cs in csinfo]
321            annotations_by_tag = [
322                get_variant_annotations(split(cs), pos)
323                for cs, pos in zip(cs_tags_list, positions_list, strict=True)
324            ]
325            variant_annotations = list(chain.from_iterable(annotations_by_tag))
326            if not variant_annotations:
327                continue
328            reference_depth = call_reference_depth(
329                variant_annotations, cs_tags_list, positions_list
330            )
331            vcf_info += add_vcf_fields(variant_annotations, chrom, reference_depth)
332
333    # Sort by chrom and pos
334    variants = sorted(
335        vcf_info,
336        key=lambda variant: (
337            chrom_sort_key(cast(str, variant.chrom)),
338            cast(int, variant.pos),
339        ),
340    )
341
342    # Write VCF
343    HEADER = """##fileformat=VCFv4.2
344    ##INFO=<ID=DP,Number=1,Type=Integer,Description="Total Depth">
345    ##INFO=<ID=RD,Number=1,Type=Integer,Description="Depth of Ref allele">
346    ##INFO=<ID=AD,Number=1,Type=Integer,Description="Depth of Alt allele">
347    ##INFO=<ID=VAF,Number=1,Type=Float,Description="Variant allele frequency (AD/DP)">
348    #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n
349    """
350
351    vcf = remove_spaces_around_newlines(HEADER).strip().split("\n")
352    for v in variants:
353        vcf.append(
354            f"{v.chrom}\t{v.pos}\t.\t{v.ref}\t{v.alt}\t.\t.\tDP={v.info.dp};RD={v.info.rd};AD={v.info.ad};VAF={v.info.vaf}"
355        )
356
357    return "\n".join(vcf)
358
359
360###########################################################
361# main
362###########################################################
363
364
365@overload
366def to_vcf(cs_tags: str, chroms: str | int, positions: int) -> str: ...
367
368
369@overload
370def to_vcf(
371    cs_tags: list[str],
372    chroms: list[str] | list[int],
373    positions: list[int],
374) -> str: ...
375
376
377def to_vcf(
378    cs_tags: str | list[str],
379    chroms: str | int | list[str] | list[int],
380    positions: int | list[int],
381) -> str:
382    """
383    Convert cs tag(s) to VCF (Variant Call Format) string.
384
385    Args:
386        cs_tag (str | list[str]): The cs tag representing the sequence alignment.
387        chrom (str | list[str]): The chromosome name.
388        pos (int | list[int]): The starting position for the sequence.
389
390    Returns:
391        str: The VCF-formatted string.
392    Example:
393        >>> import cstag
394        >>> cs_tag = "=AC*gt=T-gg=C+tt=A"
395        >>> chrom = "chr1"
396        >>> pos = 1
397        >>> print(cstag.to_vcf(cs_tag, chrom, pos))
398        ##fileformat=VCFv4.2
399        #CHROM	POS	ID	REF	ALT	QUAL	FILTER	INFO
400        chr1	3	.	G	T	.	.	.
401        chr1	4	.	TGG	T	.	.	.
402        chr1	5	.	C	CTT	.	.	.
403    """
404    if isinstance(cs_tags, str):
405        return process_cs_tag(
406            cs_tags,
407            cast(str | int, chroms),
408            cast(int, positions),
409        )
410    elif isinstance(cs_tags, list):
411        return process_cs_tags(
412            cs_tags,
413            cast(list[str] | list[int], chroms),
414            cast(list[int], positions),
415        )
416    else:
417        raise TypeError(f"cs_tags must be str or list, not {type(cs_tags)}")