From 4b79e7f1bd111642aff65c99e4ce6c488ac78d8f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florian=20R=2E=20H=C3=B6lzlwimmer?= Date: Sat, 17 Oct 2026 00:46:57 +0200 Subject: [PATCH 1/4] fix(extractors): reject overlapping variants in VariantSeqExtractor VariantSeqExtractor.extract applied overlapping variants one after the other. Two SNVs at one base gave two bases, and the sequence depended on the order of the variants. extract now raises a ValueError if the REF of two variants share a base, e.g. for two ALT alleles of one VCF record. GenomeKit rejects overlapping variants too. Variants that only touch, such as an SNV next to a deletion, still apply together. --- MIGRATING.md | 2 ++ src/kipoiseq2/extractors/vcf_seq.py | 11 ++++++++-- tests/extractors/test_vcf_seq_extractor.py | 25 ++++++++++++++++++++++ 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/MIGRATING.md b/MIGRATING.md index 34c5461..eed829a 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -217,6 +217,8 @@ kipoiseq2 reads VCF files with polars-bio instead of cyvcf2, so the `vcf` extra - `VariantSeqExtractor.extract` checks the variants. It raises a ValueError for a variant on another chromosome than `interval`, for a variant outside the chromosome, and for a variant whose REF differs from the reference, ignoring case. In 0.1, it replaced the reference bases at the position of such a variant with its ALT. It raises a TypeError if `interval` is not an `Interval` or a variant is not a `Variant`. +- `VariantSeqExtractor.extract` raises a ValueError for two variants that overlap, i.e. whose REF share a base, e.g. two ALT alleles of one VCF record. + In 0.1, it applied both: two SNVs at one base gave two bases, in the order of the variants. - `VariantSeqExtractor.extract(interval, variants, *, anchor_pos=None, pad_with_n=False)` takes its options only as keywords. `anchor` and `fixed_len` become `anchor_pos`: `anchor=a` becomes `anchor_pos=a`, and `anchor=a, fixed_len=False` becomes `anchor_pos=None`. `is_padding` is renamed to `pad_with_n`. diff --git a/src/kipoiseq2/extractors/vcf_seq.py b/src/kipoiseq2/extractors/vcf_seq.py index 423b1fc..e9cf4d6 100644 --- a/src/kipoiseq2/extractors/vcf_seq.py +++ b/src/kipoiseq2/extractors/vcf_seq.py @@ -1,3 +1,4 @@ +import itertools from collections.abc import Iterable from typing import Literal @@ -16,7 +17,8 @@ _Piece = tuple[int, int] | str """A part of the result: a reference range ``(start, end)``, 0-based and half-open, or an ALT sequence. -Overlapping variants can give a range whose ``start`` is greater than its ``end``. It stands for no bases. +A variant that spans an end of the sequence can give a range whose ``start`` is greater than its ``end``. +It stands for no bases. """ @@ -76,6 +78,7 @@ def extract( :param interval: the region of the sequence :param variants: the variants to apply, on the chromosome of ``interval``. Variants whose REF does not overlap ``interval`` have no effect. + Two variants must not overlap: their REF must not share a base. :param anchor_pos: None, or a 0-based position from ``interval.start`` to ``interval.end``. If None, the length of the sequence follows the variants. Otherwise the sequence has the width of ``interval``, anchored at ``anchor_pos`` as described above. @@ -90,7 +93,7 @@ def extract( :raises ValueError: if ``anchor_pos`` lies outside ``interval``, if ``interval.chrom`` is not in ``reference.chrom_sizes``, if a variant lies on another chromosome than ``interval`` or outside the chromosome, - if the REF of a variant differs from the reference, ignoring case, + if two variants overlap, if the REF of a variant differs from the reference, ignoring case, or if the sequence needs bases outside the chromosome and ``pad_with_n`` is False """ if not isinstance(interval, Interval): @@ -111,6 +114,10 @@ def extract( raise ValueError(f"variant {v} lies on another chromosome than the interval {interval}") if v.start < 0 or v.end > chrom_size: raise ValueError(f"variant {v} lies outside {v.chrom}, which has {chrom_size} bases") + # sorted by start, two variants overlap only if two neighbors overlap + for v1, v2 in itertools.pairwise(sorted(variants, key=lambda v: v.start)): + if v2.start < v1.end: + raise ValueError(f"variants {v1} and {v2} overlap, but the REF of two variants must not share a base") edits = [(v.start, v.end, v.alt) for v in variants] if anchor_pos is None: diff --git a/tests/extractors/test_vcf_seq_extractor.py b/tests/extractors/test_vcf_seq_extractor.py index d4bc159..1dcad24 100644 --- a/tests/extractors/test_vcf_seq_extractor.py +++ b/tests/extractors/test_vcf_seq_extractor.py @@ -273,6 +273,31 @@ def test_extract_ref_differs(variant_seq_extractor): variant_seq_extractor.extract(Interval("chr1", 20, 30), [Variant("chr1", 13, "TA", "T")], anchor_pos=20) +# the reference from 0 to 10 is ACGTAACGTA +OVERLAPPING_VARIANTS = { + "same_base": [Variant("chr1", 5, "A", "C"), Variant("chr1", 5, "A", "G")], + "snv_in_deletion": [Variant("chr1", 5, "AAC", "A"), Variant("chr1", 6, "A", "T")], + "snv_at_deletion_end": [Variant("chr1", 2, "CGTAAC", "C"), Variant("chr1", 7, "C", "T")], + "insertions_at_one_base": [Variant("chr1", 5, "A", "AT"), Variant("chr1", 5, "A", "AGG")], + "duplicate": [Variant("chr1", 5, "A", "C"), Variant("chr1", 5, "A", "C")], +} + + +@pytest.mark.parametrize("variants", OVERLAPPING_VARIANTS.values(), ids=OVERLAPPING_VARIANTS.keys()) +@pytest.mark.parametrize("anchor_pos", [None, 0, 5, 10]) +def test_extract_overlapping_variants(variant_seq_extractor, variants, anchor_pos): + # in 0.1, both ALTs were inserted, so the sequence could be too long + for order in (variants, variants[::-1]): + with pytest.raises(ValueError, match=r"^variants chr1:\S+ and chr1:\S+ overlap, but the REF of two variants"): + variant_seq_extractor.extract(Interval("chr1", 0, 10), order, anchor_pos=anchor_pos) + + +def test_extract_adjacent_variants(variant_seq_extractor): + # variants that touch but do not share a base do not overlap + variants = [Variant("chr1", 5, "A", "T"), Variant("chr1", 6, "AC", "A"), Variant("chr1", 4, "T", "TT")] + assert variant_seq_extractor.extract(Interval("chr1", 0, 10), variants) == "ACGTTTAGTA" + + @pytest.mark.parametrize("ref", ["C", "c"]) def test_extract_ref_ignores_case(ref): # sample.fasta has the sequence ACGTacgt, so with uppercase=False the reference has a lowercase c at position 6 From 70180fb7a83bbd9d890ce577f255dd35453a6aa5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florian=20R=2E=20H=C3=B6lzlwimmer?= Date: Sat, 17 Oct 2026 00:55:42 +0200 Subject: [PATCH 2/4] feat(extractors): yield batches of encoded ref and alt sequences iter_variant_seqs(pairs, reference, *, use_strand, anchor, group_by, encoding, pad_with_n, batch_size) turns interval-variant pairs into polars DataFrame batches. The pairs with equal values in the group_by columns give one row, with all their variants applied to their interval by VariantSeqExtractor.extract. The default group_by gives one row per pair, and group_by=["interval_idx"] applies all variants of an interval together. A batch holds the group keys, chrom, start, end and strand, ref_seq and alt_seq, and the encoded columns ref and alt. A Tokenizer (default Tokenizer()) gives tokens, and an Encoder gives one-hot arrays, both in their dtype. If all sequences have the same length, the encoded columns are pl.Array columns, and to_numpy() returns a (rows, length) or (rows, length, n_columns) array without a copy. Otherwise they are pl.List columns. anchor picks the anchor_pos of VariantSeqExtractor.extract by the rule of Interval.resize, so with use_strand, "start" is the 5' end of "-" intervals. anchor=None lets indels change the length of alt_seq. iter_variant_seqs collects and checks the pairs when it is called. A group with several intervals raises a ValueError. Errors of extract and of the encoding, e.g. for overlapping variants, get a note that names the group. Each pair applies its ALT allele, so pairs with a False in the column variant_carrier raise a ValueError. Such pairs come from scan_vcf_genotypes(carriers_only=False). Within a batch, the reference of each interval is fetched once, and ref_seq is cut from it. All sequences of a column are encoded in one call. _check_int tests for a plain int before it tests for numbers.Integral, which makes Interval and Variant faster to build. With 1 kb windows, iter_variant_seqs is faster than a loop that calls extract, fetches the reference and tokenizes both sequences. --- MIGRATING.md | 4 +- README.md | 37 ++- docs/api/extractors.md | 4 + src/kipoiseq2/dataclasses.py | 3 +- src/kipoiseq2/extractors/__init__.py | 1 + src/kipoiseq2/extractors/batches.py | 402 ++++++++++++++++++++++ tests/extractors/test_batches.py | 476 +++++++++++++++++++++++++++ 7 files changed, 924 insertions(+), 3 deletions(-) create mode 100644 src/kipoiseq2/extractors/batches.py create mode 100644 tests/extractors/test_batches.py diff --git a/MIGRATING.md b/MIGRATING.md index eed829a..46148e1 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -25,7 +25,8 @@ kipoiseq2 reads VCF files with polars-bio instead of cyvcf2, so the `vcf` extra The carrier check works per ALT allele, not per record: a sample with GT 0/2 carries the second ALT allele only. - `to_vcf` becomes `polars_bio.sink_vcf`, e.g. `pb.sink_vcf(pb.scan_vcf(path).filter(pl.col("qual") > 10), out_path)`. - `VariantIntervalQueryable` and the query classes of `kipoiseq.extractors.vcf_query` are gone: filter the polars frames instead. -- `SingleVariantVCFSeqExtractor` and `SingleSeqVCFSeqExtractor` are gone: match the variants with `scan_variant_pairs`, and apply them with `VariantSeqExtractor`. +- `SingleVariantVCFSeqExtractor` and `SingleSeqVCFSeqExtractor` are gone: match the variants with `scan_variant_pairs`, and apply them with `iter_variant_seqs`, which yields batches of reference and alternative sequences. + `VariantSeqExtractor` applies the variants of one interval. - `Variant.from_cyvcf` and `Variant.from_cyvcf_and_given_alt` are gone. A `Variant` has only `chrom`, `pos`, `ref` and `alt`, so `id`, `qual`, `filter`, `info`, `source` and `copy()` are gone. The VCF columns ID, QUAL, FILTER and INFO stay in the frame of `scan_vcf_variants`. @@ -41,6 +42,7 @@ kipoiseq2 reads VCF files with polars-bio instead of cyvcf2, so the `vcf` extra ``` The sort gives the pairs in VCF order, and per variant in the order of the intervals. `SingleVariantMatcher` grouped them by chromosome and strand. For the variants of each interval, as `MultiVariantsMatcher` gave them, group the pairs by `interval_idx`. An interval without variants has no pairs. + `iter_variant_seqs(pairs, reference, use_strand=True, group_by=["interval_idx"])` applies the variants of each interval together. - `scan_variant_pairs` takes the intervals as `intervals`, a polars DataFrame with the columns `chrom`, `start`, `end`, optionally `strand` and any other columns, instead of `pranges`. For a PyRanges object `pr`, pass `intervals=pl.from_pandas(pr.df).rename({"Chromosome": "chrom", "Start": "start", "End": "end", "Strand": "strand"})`. For a list of Interval objects, pass `pl.DataFrame(intervals)`. diff --git a/README.md b/README.md index 3f31c6f..93f912d 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ Requires Python >= 3.12. ```bash pip install kipoiseq2 # FASTA extractors and transforms (numpy, pyfaidx) -pip install 'kipoiseq2[ranges]' # + scan_vcf_variants, scan_vcf_genotypes and scan_variant_pairs (polars, polars-bio) +pip install 'kipoiseq2[ranges]' # + scan_vcf_variants, scan_vcf_genotypes, scan_variant_pairs and iter_variant_seqs (polars, polars-bio) ``` kipoiseq2 does not depend on pandas or pyranges. @@ -172,6 +172,41 @@ A null strand in the intervals frame means `"."`, as a missing strand column. polars-bio prints a progress bar on stderr for each collect. Set the environment variable `TQDM_DISABLE=1` to turn it off. +### Batches of reference and alternative sequences + +`iter_variant_seqs` applies the variants of the pairs to the reference, and yields the sequences as polars DataFrame batches. +By default, each pair gives one row, and the sequences keep the width of the interval around its center. +A batch holds the keys `interval_idx` and `variant_idx`, the interval, the sequences `ref_seq` and `alt_seq`, and the encoded sequences `ref` and `alt`. +The default encoding `Tokenizer()` gives uint8 tokens. +An `Encoder` as `encoding` gives one-hot arrays in the dtype of the encoder. + +```python +import polars as pl +from kipoiseq2.extractors import FastaStringExtractor, iter_variant_seqs, scan_variant_pairs, scan_vcf_variants + +variants = scan_vcf_variants("variants.vcf.gz").collect() +# a 1 kb window around each variant, in the order of the variants +windows = variants.select("chrom", start=pl.col("start") - 500, end=pl.col("start") + 500) +# keep each variant in its own window only +pairs = scan_variant_pairs(windows, variants).filter(pl.col("interval_idx") == pl.col("variant_idx")) + +with FastaStringExtractor("genome.fa", use_strand=False) as fasta: + for batch in iter_variant_seqs(pairs, fasta, use_strand=True, pad_with_n=True, batch_size=64): + ref = batch["ref"].to_numpy() # uint8 array of shape (64, 1000), without a copy + alt = batch["alt"].to_numpy() +``` + +A variant within 500 bases of a chromosome end gets a window that runs past the chromosome. +`pad_with_n=True` fills the bases outside the chromosome with N. + +`group_by` sets the rows: the pairs with equal values in these columns give one row, with all their variants applied. +`group_by=["interval_idx"]` applies all variants of an interval together. +Variants that overlap within a row raise a ValueError, as in `VariantSeqExtractor`. +Each pair applies its ALT allele. So for the genotypes of `scan_vcf_genotypes`, keep the default `carriers_only=True`: a False or a null in the column `variant_carrier` raises a ValueError. +`encoding=Encoder()` gives uint8 one-hot arrays of shape `(rows, width, 4)`, and `encoding=None` gives only the sequences. +`anchor="start"` or `"end"` keeps the 5' or the 3' end of the interval, and `anchor=None` lets indels change the length of `alt_seq`. +Sequences of different lengths give polars List columns instead of Array columns. + More examples: - The tests in [tests/](tests/) show the usage of every extractor and transform. - API docs: diff --git a/docs/api/extractors.md b/docs/api/extractors.md index 9e6d29f..00705be 100644 --- a/docs/api/extractors.md +++ b/docs/api/extractors.md @@ -20,4 +20,8 @@ .. automodule:: kipoiseq2.extractors.vcf_matching :members: :undoc-members: + +.. automodule:: kipoiseq2.extractors.batches + :members: + :undoc-members: ``` diff --git a/src/kipoiseq2/dataclasses.py b/src/kipoiseq2/dataclasses.py index 288ab5e..ecdfc8e 100644 --- a/src/kipoiseq2/dataclasses.py +++ b/src/kipoiseq2/dataclasses.py @@ -92,7 +92,8 @@ def from_str(cls, s: str) -> "Variant": def _check_int(value: object, name: str) -> None: """Raise a TypeError if ``value`` is not an int. A bool does not count as an int.""" - if not isinstance(value, numbers.Integral) or isinstance(value, bool): + # the type check is much faster than isinstance with numbers.Integral, and covers most values + if type(value) is not int and (not isinstance(value, numbers.Integral) or isinstance(value, bool)): raise TypeError(f"{name} must be an int, not {type(value).__name__}") diff --git a/src/kipoiseq2/extractors/__init__.py b/src/kipoiseq2/extractors/__init__.py index 6015163..9f31ae8 100644 --- a/src/kipoiseq2/extractors/__init__.py +++ b/src/kipoiseq2/extractors/__init__.py @@ -1,4 +1,5 @@ from .base import * +from .batches import * from .fasta import * from .vcf import * from .vcf_matching import * diff --git a/src/kipoiseq2/extractors/batches.py b/src/kipoiseq2/extractors/batches.py new file mode 100644 index 0000000..46cf481 --- /dev/null +++ b/src/kipoiseq2/extractors/batches.py @@ -0,0 +1,402 @@ +"""Build batches of reference and alternative sequences as polars DataFrames. + +:func:`iter_variant_seqs` applies the variants of interval-variant pairs, e.g. +of :func:`~kipoiseq2.extractors.vcf_matching.scan_variant_pairs`, to the +reference, and encodes the sequences for model input. It needs polars, which +the ``ranges`` extra installs. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any + +import numpy as np +from numpy.typing import NDArray + +from kipoiseq2.dataclasses import _ANCHORS, Interval, Variant, _Anchor, _check_int, _is_reverse +from kipoiseq2.encoders import Encoder, Tokenizer +from kipoiseq2.extractors.base import SequenceExtractor +from kipoiseq2.extractors.vcf import _check_names +from kipoiseq2.extractors.vcf_seq import VariantSeqExtractor +from kipoiseq2.transforms import reverse_complement + +if TYPE_CHECKING: + import polars as pl + +__all__ = ["iter_variant_seqs"] + +_INTERVAL_COLUMNS = ("chrom", "start", "end", "strand") +"""The columns of the pairs that hold the interval.""" + +_VARIANT_COLUMNS = ("variant_pos", "variant_ref", "variant_alt") +"""The columns of the pairs that hold the variant.""" + +_SEQ_COLUMNS = ("ref_seq", "alt_seq", "ref", "alt") +"""The columns that :func:`iter_variant_seqs` adds to the group keys and the interval.""" + +_CARRIER_COLUMN = "variant_carrier" +"""The column of the pairs that tells whether the sample carries the ALT allele, +from the carrier column of :func:`~kipoiseq2.extractors.vcf.scan_vcf_genotypes`.""" + + +def iter_variant_seqs( + pairs: pl.DataFrame | pl.LazyFrame, + reference: SequenceExtractor, + *, + use_strand: bool, + anchor: _Anchor | None = "center", + group_by: Sequence[str] = ("interval_idx", "variant_idx"), + encoding: Tokenizer | Encoder | None = Tokenizer(), + pad_with_n: bool = False, + batch_size: int = 64, +) -> Iterator[pl.DataFrame]: + """Yield the reference and alternative sequences of interval-variant pairs as DataFrame batches. + + The rows of ``pairs`` with equal values in the columns ``group_by`` form a group. + Each group gives one output row: its variants are applied to its interval with + :meth:`VariantSeqExtractor.extract `. + All rows of a group must have the same interval, and the variants of a group must not overlap. + The default ``group_by`` gives one row per pair, e.g. for variant effect predictions. + ``group_by=["interval_idx"]`` applies all variants of an interval together. + An interval without variants has no pairs, so it gives no row. + Each pair applies its ALT allele, so ``pairs`` must not contain samples that do not carry it. + For the genotypes of :func:`~kipoiseq2.extractors.vcf.scan_vcf_genotypes`, keep ``carriers_only=True``: + a False or a null in the column variant_carrier raises a ValueError. + + ``pairs`` is collected into memory when ``iter_variant_seqs`` is called, and the groups are checked then. + The sequences of a batch are built when the iterator reaches it. + Within a batch, the reference of each interval is fetched once. + + The encoded columns are polars Array columns if all sequences have the same length, + i.e. if ``anchor`` is not None and all intervals of ``pairs`` have the same width. + Then ``batch["alt"].to_numpy()`` returns a read-only array of shape ``(rows, length)`` for a Tokenizer, + or ``(rows, length, n_columns)`` for an Encoder, without a copy. + Otherwise the encoded columns are polars List columns, with one item per letter. + + :param pairs: DataFrame or LazyFrame with the columns chrom, start, end, strand, + variant_pos, variant_ref, variant_alt and the columns of ``group_by``, + e.g. of :func:`~kipoiseq2.extractors.vcf_matching.scan_variant_pairs` + :param reference: the reference sequence, e.g. a :class:`~kipoiseq2.extractors.fasta.FastaStringExtractor`. + It gets intervals with the strand ``"."``, as from :class:`~kipoiseq2.extractors.vcf_seq.VariantSeqExtractor`. + :param use_strand: if True, reverse complement the sequences of "-" intervals, and let ``anchor`` follow the strand + :param anchor: the part of the interval that keeps its position: ``"start"``, ``"center"`` or ``"end"``. + Then alt_seq has the width of the interval. An indel pulls in or pushes out bases at the end of the interval + on its side of the anchor. + ``anchor`` picks the ``anchor_pos`` of ``VariantSeqExtractor.extract`` by the rule of + :meth:`Interval.resize `: + ``"start"`` is the 5' end, which is ``end`` on "-" intervals if ``use_strand`` is True, and ``start`` otherwise. + ``"end"`` is the other end. ``"center"`` is ``interval.center(use_strand=use_strand)``. + None gives alt_seq the length that the variants give, as ``anchor_pos=None``. + :param group_by: the columns of ``pairs`` that form the groups + :param encoding: a :class:`~kipoiseq2.encoders.Tokenizer` for the tokens of + :meth:`~kipoiseq2.encoders.Tokenizer.tokenize`, an :class:`~kipoiseq2.encoders.Encoder` for the one-hot arrays of + :meth:`~kipoiseq2.encoders.Encoder.one_hot`, or None for no encoded columns. + The dtype of ``encoding`` sets the dtype of the encoded columns, e.g. ``Encoder(dtype=np.float16)``. + :param pad_with_n: if True, use N for the bases outside the chromosome. Otherwise raise a ValueError for them. + :param batch_size: number of rows per batch. The last batch can have fewer rows. + :returns: iterator of DataFrames, sorted by the columns of ``group_by``, with the columns: + + - the columns of ``group_by`` + - chrom, start, end and strand of the interval, unless ``group_by`` contains them + - ref_seq: the reference sequence of the interval + - alt_seq: the sequence of the interval with the variants of the group applied + - ref and alt: ref_seq and alt_seq encoded, unless ``encoding`` is None. + For a Tokenizer, they are ``pl.Array(dtype, length)`` with the polars type of its dtype, + e.g. ``pl.UInt8`` for the default ``Tokenizer()``. + For an Encoder, they are ``pl.Array(dtype, (length, n_columns))`` with the polars type of its dtype. + length is the number of letters. + :raises TypeError: if ``pairs`` is not a polars DataFrame or LazyFrame, + if ``reference`` is not a :class:`~kipoiseq2.extractors.base.SequenceExtractor`, + if ``group_by`` is not a sequence of strings, if ``encoding`` is not a Tokenizer, an Encoder or None, + or if ``batch_size`` is not an int + :raises ValueError: if ``anchor`` is not ``"start"``, ``"center"``, ``"end"`` or None, + if ``group_by`` is empty, repeats a column or names ref_seq, alt_seq, ref or alt, + if ``pairs`` lacks a column, if a required column of ``pairs`` contains nulls, + if the column variant_carrier of ``pairs`` contains False or null, + if a group has several intervals, or if ``batch_size`` is less than 1. + While iterating, it raises the errors of ``VariantSeqExtractor.extract`` and of the encoding, + with a note that names the group, e.g. a ValueError for overlapping variants, + for a REF that differs from the reference, or for a letter that ``encoding`` does not contain. + + Example: + + >>> pairs = scan_variant_pairs(windows, scan_vcf_variants("variants.vcf.gz")) # doctest: +SKIP + >>> with FastaStringExtractor("genome.fa", use_strand=False) as fasta: # doctest: +SKIP + ... for batch in iter_variant_seqs(pairs, fasta, use_strand=True, batch_size=32): + ... alt = batch["alt"].to_numpy() # uint8 tokens of shape (32, window width) + """ + import polars as pl + + if not isinstance(pairs, (pl.DataFrame, pl.LazyFrame)): + raise TypeError(f"pairs must be a polars DataFrame or LazyFrame, not {type(pairs).__name__}") + if not isinstance(reference, SequenceExtractor): + raise TypeError(f"reference must be a SequenceExtractor, not {type(reference).__name__}") + if anchor is not None and anchor not in _ANCHORS: + raise ValueError(f'anchor must be "start", "center", "end" or None, got {anchor!r}') + _check_names(group_by, "group_by") + keys = list(group_by) + if not keys: + raise ValueError("group_by must name at least one column") + if len(set(keys)) < len(keys): + raise ValueError(f"group_by must not repeat a column, got {keys}") + clash = [c for c in keys if c in _SEQ_COLUMNS] + if clash: + raise ValueError(f"group_by must not contain {clash}, because the batches use these names for the sequences") + encoder: Callable[[str], NDArray[Any]] | None = None + letter_length = 1 + if isinstance(encoding, Tokenizer): + encoder = encoding.tokenize + letter_length = len(encoding.alphabet.letters[0]) + elif isinstance(encoding, Encoder): + encoder = encoding.one_hot + letter_length = len(encoding.alphabet.letters[0]) + elif encoding is not None: + raise TypeError(f"encoding must be a Tokenizer, an Encoder or None, not {type(encoding).__name__}") + _check_int(batch_size, "batch_size") + if batch_size < 1: + raise ValueError(f"batch_size must be at least 1, got {batch_size}") + + columns = pairs.collect_schema().names() + missing = [c for c in (*_INTERVAL_COLUMNS, *_VARIANT_COLUMNS, *keys) if c not in columns] + if missing: + raise ValueError(f"pairs lacks the columns {missing}") + + # The keys get private names, so they cannot clash with the other columns. + key_names = [f"key_{i}" for i in range(len(keys))] + carrier = [_CARRIER_COLUMN] if _CARRIER_COLUMN in columns else [] + frame = ( + pairs.lazy() + .select(*(pl.col(c).alias(k) for c, k in zip(keys, key_names)), *_INTERVAL_COLUMNS, *_VARIANT_COLUMNS, *carrier) + .collect() + ) + if carrier and not frame.get_column(_CARRIER_COLUMN).fill_null(False).all(): + raise ValueError( + f"pairs has rows whose {_CARRIER_COLUMN} is False or null, " + "but iter_variant_seqs applies the ALT allele of every row. " + "Scan the genotypes with scan_vcf_genotypes(..., carriers_only=True)." + ) + null_counts = frame.select(pl.col(*_INTERVAL_COLUMNS, *_VARIANT_COLUMNS).null_count()).row(0, named=True) + nulls = [c for c, n in null_counts.items() if n] + if nulls: + raise ValueError(f"pairs has nulls in the columns {nulls}") + groups = ( + frame.group_by(key_names) + .agg( + pl.struct(_INTERVAL_COLUMNS).n_unique().alias("n_intervals"), + pl.col(_INTERVAL_COLUMNS).first(), + pl.col(_VARIANT_COLUMNS), + # extract fetches at most the sum of the REF lengths beyond each end of the interval + pl.col("variant_ref").str.len_chars().sum().alias("margin"), + ) + .sort(key_names) + ) + several = groups.filter(pl.col("n_intervals") > 1) + if len(several): + group = dict(zip(keys, several.select(key_names).row(0))) + raise ValueError( + f"the group {group} of pairs has {several['n_intervals'][0]} intervals, but a group needs one interval. " + "Add interval_idx to group_by." + ) + same_length = anchor is not None and groups.select((pl.col("end") - pl.col("start")).n_unique()).item() <= 1 + return _iter_batches( + groups, + reference, + keys=keys, + key_names=key_names, + use_strand=use_strand, + anchor=anchor, + encoder=encoder, + letter_length=letter_length, + same_length=same_length, + pad_with_n=pad_with_n, + batch_size=batch_size, + ) + + +def _iter_batches( + groups: pl.DataFrame, + reference: SequenceExtractor, + *, + keys: list[str], + key_names: list[str], + use_strand: bool, + anchor: _Anchor | None, + encoder: Callable[[str], NDArray[Any]] | None, + letter_length: int, + same_length: bool, + pad_with_n: bool, + batch_size: int, +) -> Iterator[pl.DataFrame]: + """Yield the batches of :func:`iter_variant_seqs`. + + :param groups: one row per group, sorted, with the keys, n_intervals, the interval columns, + the variant columns as lists, and margin + :param keys: the columns of ``group_by`` + :param key_names: the names of the keys in ``groups`` + :param encoder: the function that encodes a sequence, or None for no encoded columns + :param letter_length: the length of a letter of ``encoder`` + :param same_length: if True, all sequences have the same length, so the encoded columns are Array columns + """ + import polars as pl + + window_reference = _WindowReference(reference) + extractor = VariantSeqExtractor(window_reference, use_strand=use_strand) + for batch in groups.iter_slices(batch_size): + # Series.to_list and the DataFrame constructor run no polars query, unlike select + chroms, starts, ends, strands, positions, ref_alleles, alt_alleles, margins = ( + batch.get_column(c).to_list() for c in (*_INTERVAL_COLUMNS, *_VARIANT_COLUMNS, "margin") + ) + window_margins: dict[tuple[str, int, int], int] = {} + for coords, margin in zip(zip(chroms, starts, ends), margins): + window_margins[coords] = max(margin, window_margins.get(coords, 0)) + windows = {coords: window_reference.fetch(*coords, margin) for coords, margin in window_margins.items()} + + ref_seqs: list[str] = [] + alt_seqs: list[str] = [] + rows = zip(chroms, starts, ends, strands, positions, ref_alleles, alt_alleles) + for i, (chrom, start, end, strand, row_positions, row_refs, row_alts) in enumerate(rows): + try: + window_reference.window = windows[chrom, start, end] + interval = Interval(chrom, start, end, strand=strand) + ref_seq = _ref_seq(window_reference, extractor, interval, pad_with_n=pad_with_n) + variants = [Variant(chrom, *v) for v in zip(row_positions, row_refs, row_alts)] + # a resize to width 0 keeps only the position of the anchor + anchor_pos = None if anchor is None else interval.resize(0, anchor=anchor, use_strand=use_strand).start + alt_seq = extractor.extract(interval, variants, anchor_pos=anchor_pos, pad_with_n=pad_with_n) + except (TypeError, ValueError) as e: + e.add_note(_group_note(batch, i, keys, key_names)) + raise + ref_seqs.append(ref_seq) + alt_seqs.append(alt_seq) + + columns = [ + *(batch.get_column(name).alias(k) for k, name in zip(keys, key_names)), + *(batch.get_column(c) for c in _INTERVAL_COLUMNS if c not in keys), + pl.Series("ref_seq", ref_seqs, dtype=pl.String), + pl.Series("alt_seq", alt_seqs, dtype=pl.String), + ] + if encoder is not None: + try: + columns += [ + _encoded_series("ref", ref_seqs, encoder, letter_length=letter_length, same_length=same_length), + _encoded_series("alt", alt_seqs, encoder, letter_length=letter_length, same_length=same_length), + ] + except ValueError: + # encode row by row to find the group of the error + for i, seqs in enumerate(zip(ref_seqs, alt_seqs)): + try: + for seq in seqs: + encoder(seq) + except ValueError as e: + e.add_note(_group_note(batch, i, keys, key_names)) + raise + raise + yield pl.DataFrame(columns) + + +def _group_note(batch: pl.DataFrame, i: int, keys: list[str], key_names: list[str]) -> str: + """Return the note that names the group of row ``i`` of ``batch`` for an error.""" + group = {k: batch.get_column(name)[i] for k, name in zip(keys, key_names)} + return f"in the group {group} of pairs" + + +def _ref_seq( + window_reference: _WindowReference, extractor: VariantSeqExtractor, interval: Interval, *, pad_with_n: bool +) -> str: + """Return the reference sequence of ``interval``, as ``extractor.extract(interval, [])``. + + It cuts the sequence from the window if the window covers ``interval``. + The window covers every interval within the chromosome, so only an interval past an end of the chromosome + needs ``extract``, which pads it or raises the error. + """ + seq = window_reference.cut(interval) + if seq is None: + return extractor.extract(interval, [], pad_with_n=pad_with_n) + if _is_reverse(interval, use_strand=extractor.use_strand): + return reverse_complement(seq) + return seq + + +def _encoded_series( + name: str, seqs: list[str], encoder: Callable[[str], NDArray[Any]], *, letter_length: int, same_length: bool +) -> pl.Series: + """Return the encoded sequences of a batch as one polars Series. + + ``encoder`` encodes the joined sequences in one call, which is faster than one call per sequence. + + :param seqs: the sequences of the rows + :param encoder: the function that encodes a sequence into an array with one item per letter + :param letter_length: the length of a letter of ``encoder`` + :param same_length: if True, all sequences have the same length, and the Series is an Array column. + Otherwise it is a List column with one item per letter. + :raises ValueError: if ``encoder`` cannot encode a sequence + """ + import polars as pl + + lengths = [len(seq) // letter_length for seq in seqs] + letters = encoder("".join(seqs)) + if len(letters) != sum(lengths): + # the joined sequences can split into letters even if a sequence cannot + raise ValueError(f"the length of a sequence is not a multiple of the letter length {letter_length}") + if same_length: + # polars keeps the array without a copy + return pl.Series(name, letters.reshape(len(seqs), lengths[0], *letters.shape[1:])) + arrays = np.split(letters, np.cumsum(lengths[:-1])) + inner = pl.Series(np.empty(0, dtype=letters.dtype)).dtype + # polars builds an Array column if all arrays have the same length, so cast it + series = pl.Series(name, [a.ravel() for a in arrays]).cast(pl.List(inner)) + if letters.ndim == 2: + # one Array of n_columns values per letter + series = series.list.eval(pl.element().reshape((-1, letters.shape[1]))) + return series + + +class _WindowReference: + """The reference of :func:`iter_variant_seqs`, which serves the sequences of an interval from one fetch. + + :meth:`extract` cuts the sequence from :attr:`window` if the window covers it, + and fetches it from the reference otherwise. + + :param reference: the reference sequence + """ + + def __init__(self, reference: SequenceExtractor) -> None: + self._reference = reference + self.window: tuple[str, int, str] = ("", 0, "") + """The current window: its chromosome, its start and its sequence on the "+" strand.""" + + @property + def chrom_sizes(self) -> Mapping[str, int]: + """Length of each chromosome, by name.""" + return self._reference.chrom_sizes + + def fetch(self, chrom: str, start: int, end: int, margin: int) -> tuple[str, int, str]: + """Fetch the window from ``start - margin`` to ``end + margin``, within the chromosome. + + :returns: the window for :attr:`window`. For a chromosome that is not in the reference, an empty window. + """ + chrom_size = self.chrom_sizes.get(chrom) + if chrom_size is None: + # extract raises the error for the chromosome + return chrom, 0, "" + window_start = min(max(start - margin, 0), chrom_size) + window_end = max(min(end + margin, chrom_size), window_start) + return chrom, window_start, self._reference.extract(Interval(chrom, window_start, window_end)) + + def cut(self, interval: Interval) -> str | None: + """Return the "+" strand of ``interval`` if :attr:`window` covers it, and None otherwise.""" + chrom, window_start, seq = self.window + start = interval.start - window_start + end = interval.end - window_start + if interval.chrom == chrom and start >= 0 and end <= len(seq): + return seq[start:end] + return None + + def extract(self, interval: Interval) -> str: + """Return the "+" strand of ``interval``. + + :param interval: an interval with the strand ``"."``, from VariantSeqExtractor + """ + seq = self.cut(interval) + return self._reference.extract(interval) if seq is None else seq diff --git a/tests/extractors/test_batches.py b/tests/extractors/test_batches.py new file mode 100644 index 0000000..a66feb7 --- /dev/null +++ b/tests/extractors/test_batches.py @@ -0,0 +1,476 @@ +import random + +import numpy as np +import polars as pl +import pytest +from conftest import sample_5kb_fasta_file + +from kipoiseq2.alphabets import Alphabet +from kipoiseq2.dataclasses import Interval, Variant +from kipoiseq2.encoders import Encoder, Tokenizer +from kipoiseq2.extractors import ( + FastaStringExtractor, + VariantSeqExtractor, + iter_variant_seqs, + scan_variant_pairs, + scan_vcf_genotypes, + scan_vcf_variants, +) +from kipoiseq2.transforms import reverse_complement + +FLOAT32_ONE_HOT = Encoder(ambiguous="fractional", dtype=np.float32) +"""float32 one-hot arrays, with N as 0.25 in every column.""" + +# random sequences, so that a shifted sequence differs from the original +_rng = random.Random(0) +CHROMS = {"chr1": "".join(_rng.choices("ACGT", k=200)), "chr2": "".join(_rng.choices("ACGT", k=60))} + + +class SeqReference: + """SequenceExtractor of the sequences in CHROMS, which records the intervals that it extracts.""" + + def __init__(self, seqs=CHROMS): + self.seqs = seqs + self.calls = [] + + @property + def chrom_sizes(self): + return {chrom: len(seq) for chrom, seq in self.seqs.items()} + + def extract(self, interval): + self.calls.append(interval) + return self.seqs[interval.chrom][interval.start : interval.end] + + +def variant(chrom, pos, ref_len, alt): + """Return the variant at ``pos`` with the REF of CHROMS. ``alt`` may use ``{ref}`` and ``{first}``.""" + ref = CHROMS[chrom][pos - 1 : pos - 1 + ref_len] + return Variant(chrom, pos, ref, alt.format(ref=ref, first=ref[0])) + + +# The variants do not overlap each other, so all variants of an interval can apply together. +VARIANTS = [ + variant("chr1", 9, 3, "{first}"), # deletion across the start of the first interval + variant("chr1", 12, 1, "N"), + variant("chr1", 20, 1, "{first}GAT"), + variant("chr1", 25, 3, "{first}"), + variant("chr1", 31, 2, "NN"), + variant("chr1", 38, 5, "{first}"), # deletion across the end of the first interval + variant("chr1", 55, 1, "{first}CC"), + variant("chr1", 101, 1, "N"), + variant("chr1", 190, 1, "N"), # in no interval + variant("chr2", 8, 1, "N"), + variant("chr2", 15, 6, "{first}"), + variant("chr2", 29, 1, "{first}A"), +] + +# widths 30, 21, 30, 30, 1 and 10, the last interval without variants +INTERVALS = pl.DataFrame( + { + "chrom": ["chr1", "chr1", "chr1", "chr2", "chr1", "chr1"], + "start": [10, 15, 40, 6, 100, 150], + "end": [40, 36, 70, 36, 101, 160], + "strand": ["+", "-", ".", "-", "+", "+"], + "name": ["a", "b", "c", "d", "e", "f"], + } +) + +# the intervals of the same width 30 +SAME_WIDTH = INTERVALS.filter(pl.col("end") - pl.col("start") == 30) + + +def pairs_of(intervals, variants=VARIANTS): + return scan_variant_pairs(intervals, pl.DataFrame(variants)) + + +def collect(batches): + return pl.concat(list(batches)) + + +def expected_anchor_pos(interval, anchor, use_strand): + """Return the position that Interval.resize keeps for ``anchor``.""" + return interval.resize(0, anchor=anchor, use_strand=use_strand).start + + +@pytest.mark.parametrize("use_strand", [True, False]) +@pytest.mark.parametrize("anchor", ["start", "center", "end", None]) +@pytest.mark.parametrize("group_by", [("interval_idx", "variant_idx"), ("interval_idx",)]) +def test_iter_variant_seqs_equals_extract(use_strand, anchor, group_by): + pairs = pairs_of(INTERVALS).collect() + reference = SeqReference() + out = collect(iter_variant_seqs(pairs, reference, use_strand=use_strand, anchor=anchor, group_by=group_by)) + # one batch, so one fetch per interval + assert len(reference.calls) == out.select("chrom", "start", "end").n_unique() + + assert len(out) == pairs.select(group_by).n_unique() + assert out.select(group_by).rows() == sorted(out.select(group_by).rows()) + extractor = VariantSeqExtractor(reference, use_strand=use_strand) + for row in out.iter_rows(named=True): + group = pairs.filter(*(pl.col(k) == row[k] for k in group_by)) + variants = [ + Variant(r["chrom"], r["variant_pos"], r["variant_ref"], r["variant_alt"]) + for r in group.iter_rows(named=True) + ] + interval = Interval(row["chrom"], row["start"], row["end"], strand=row["strand"]) + anchor_pos = None if anchor is None else expected_anchor_pos(interval, anchor, use_strand) + assert row["alt_seq"] == extractor.extract(interval, variants, anchor_pos=anchor_pos) + ref_seq = reference.extract(interval) + assert row["ref_seq"] == (reverse_complement(ref_seq) if use_strand and interval.strand == "-" else ref_seq) + assert row["ref"] == Tokenizer().tokenize(row["ref_seq"]).tolist() + assert row["alt"] == Tokenizer().tokenize(row["alt_seq"]).tolist() + + +@pytest.mark.parametrize("use_strand", [True, False]) +@pytest.mark.parametrize("anchor", ["start", "center", "end", None]) +def test_iter_variant_seqs_equals_fasta_extract(alt_kinds_vcf, use_strand, anchor): + # the reference is a FastaStringExtractor with use_strand=True, which gets intervals with the strand "." + intervals = pl.DataFrame({"chrom": ["chr1"] * 3, "start": [0, 3, 10], "end": [12, 13, 22], "strand": list("+-.")}) + pairs = scan_variant_pairs(intervals, scan_vcf_variants(alt_kinds_vcf)) + with FastaStringExtractor(sample_5kb_fasta_file, use_strand=True) as reference: + out = collect(iter_variant_seqs(pairs, reference, use_strand=use_strand, anchor=anchor)) + pairs = pairs.collect() + assert len(out) == len(pairs) + with FastaStringExtractor(sample_5kb_fasta_file, use_strand=use_strand) as fasta: + extractor = VariantSeqExtractor(fasta, use_strand=use_strand) + for row in out.iter_rows(named=True): + pair = pairs.row( + by_predicate=(pl.col("interval_idx") == row["interval_idx"]) + & (pl.col("variant_idx") == row["variant_idx"]), + named=True, + ) + interval = Interval(row["chrom"], row["start"], row["end"], strand=row["strand"]) + variants = [Variant(row["chrom"], pair["variant_pos"], pair["variant_ref"], pair["variant_alt"])] + anchor_pos = None if anchor is None else expected_anchor_pos(interval, anchor, use_strand) + assert row["alt_seq"] == extractor.extract(interval, variants, anchor_pos=anchor_pos) + assert row["ref_seq"] == fasta.extract(interval) + + +@pytest.mark.parametrize( + "anchor, use_strand, alt_seq", + [ + # anchor_pos 12, the 5' end of the "-" interval: the deletion pulls in the base at 3 + ("start", True, reverse_complement("ACCCGGGG")), + # anchor_pos 4: the deletion pulls in the base at 12 + ("end", True, reverse_complement("CCCGGGGT")), + ("start", False, "CCCGGGGT"), + ("end", False, "ACCCGGGG"), + # the center 8 lies right of the deletion + ("center", True, reverse_complement("ACCCGGGG")), + (None, True, reverse_complement("CCCGGGG")), + ], +) +def test_iter_variant_seqs_anchor_follows_strand(anchor, use_strand, alt_seq): + reference = SeqReference({"t": "AAAACCCCGGGGTTTT"}) + pairs = pl.DataFrame( + { + "chrom": ["t"], + "start": [4], + "end": [12], + "strand": ["-"], + "interval_idx": [0], + "variant_pos": [7], + "variant_ref": ["CC"], + "variant_alt": ["C"], + "variant_idx": [0], + } + ) + (batch,) = iter_variant_seqs(pairs, reference, use_strand=use_strand, anchor=anchor) + assert batch["ref_seq"].to_list() == [reverse_complement("CCCCGGGG") if use_strand else "CCCCGGGG"] + assert batch["alt_seq"].to_list() == [alt_seq] + + +def test_iter_variant_seqs_columns(): + pairs = pairs_of(SAME_WIDTH) + (batch,) = iter_variant_seqs(pairs, SeqReference(), use_strand=True) + assert batch.schema == pl.Schema( + { + "interval_idx": pl.UInt32, + "variant_idx": pl.UInt32, + "chrom": pl.String, + "start": pl.Int64, + "end": pl.Int64, + "strand": pl.String, + "ref_seq": pl.String, + "alt_seq": pl.String, + "ref": pl.Array(pl.UInt8, 30), + "alt": pl.Array(pl.UInt8, 30), + } + ) + + +def test_iter_variant_seqs_group_by_interval_columns(): + # the keys come first, and the interval columns that are keys appear once + pairs = pairs_of(SAME_WIDTH) + (batch,) = iter_variant_seqs(pairs, SeqReference(), use_strand=True, group_by=["name", "chrom"], encoding=None) + assert batch.columns == ["name", "chrom", "start", "end", "strand", "ref_seq", "alt_seq"] + assert batch["name"].to_list() == ["a", "c", "d"] + + +def test_iter_variant_seqs_several_variants_per_interval(): + pairs = pairs_of(INTERVALS.filter(pl.col("name") == "a")) + (batch,) = iter_variant_seqs(pairs, SeqReference(), use_strand=True, anchor=None, group_by=["interval_idx"]) + # the deletions across the ends of the interval keep only their bases within the interval + seq = CHROMS["chr1"] + parts = ["N", seq[12:19], seq[19], "GAT", seq[20:24], seq[24], seq[27:30], "NN", seq[32:37], seq[37]] + assert batch["alt_seq"].to_list() == ["".join(parts)] + + +@pytest.mark.parametrize("batch_size", [1, 2, 3, 7, 8, 100]) +def test_iter_variant_seqs_batch_size(batch_size): + pairs = pairs_of(INTERVALS) + n_rows = len(pairs.collect()) + batches = list(iter_variant_seqs(pairs, SeqReference(), use_strand=True, batch_size=batch_size)) + assert [len(b) for b in batches] == [min(batch_size, n_rows - i) for i in range(0, n_rows, batch_size)] + assert collect(batches).equals( + collect(iter_variant_seqs(pairs, SeqReference(), use_strand=True, batch_size=n_rows)) + ) + + +def test_iter_variant_seqs_fetches_each_interval_once_per_batch(): + # the first interval has 6 pairs + pairs = pairs_of(INTERVALS.filter(pl.col("name") == "a")) + for batch_size, n_batches in [(6, 1), (4, 2), (1, 6)]: + reference = SeqReference() + batches = list(iter_variant_seqs(pairs, reference, use_strand=True, batch_size=batch_size)) + assert len(batches) == n_batches + assert len(reference.calls) == n_batches + + +def test_iter_variant_seqs_empty_pairs(): + pairs = pairs_of(INTERVALS.filter(pl.col("name") == "f")) + assert pairs.collect().is_empty() + assert list(iter_variant_seqs(pairs, SeqReference(), use_strand=True)) == [] + + +def test_iter_variant_seqs_tokens_round_trip(): + (batch,) = iter_variant_seqs(pairs_of(SAME_WIDTH), SeqReference(), use_strand=True, batch_size=100) + for column, seq_column in [("ref", "ref_seq"), ("alt", "alt_seq")]: + tokens = batch[column].to_numpy() + assert tokens.shape == (len(batch), 30) + assert tokens.dtype == np.uint8 + # np.eye(15, 4) has the row of Encoder() for each token of DNA, and decode gives N for a row of zeros + assert [Encoder().decode(np.eye(15, 4, dtype=np.uint8)[t]) for t in tokens] == batch[seq_column].to_list() + + +def test_iter_variant_seqs_one_hot_round_trip(): + (batch,) = iter_variant_seqs( + pairs_of(SAME_WIDTH), SeqReference(), use_strand=True, encoding=FLOAT32_ONE_HOT, batch_size=100 + ) + assert batch.schema["alt"] == pl.Array(pl.Float32, (30, 4)) + alt = batch["alt"].to_numpy() + assert alt.shape == (len(batch), 30, 4) + assert alt.dtype == np.float32 + assert [FLOAT32_ONE_HOT.decode(a) for a in alt] == batch["alt_seq"].to_list() + np.testing.assert_array_equal(alt, np.stack([FLOAT32_ONE_HOT.one_hot(s) for s in batch["alt_seq"]])) + + +@pytest.mark.parametrize("encoding", [Tokenizer(), FLOAT32_ONE_HOT]) +def test_iter_variant_seqs_to_numpy_without_copy(encoding): + (batch,) = iter_variant_seqs(pairs_of(SAME_WIDTH), SeqReference(), use_strand=True, encoding=encoding) + # each call returns a view of the same polars buffer + assert np.shares_memory(batch["alt"].to_numpy(), batch["alt"].to_numpy()) + assert not batch["alt"].to_numpy().flags.writeable + + +# a Tokenizer gives tokens and an Encoder gives one-hot arrays, both in their dtype +@pytest.mark.parametrize( + "encoding, polars_dtype", + [ + (Tokenizer(), pl.Array(pl.UInt8, 30)), + (Tokenizer(dtype=np.int64), pl.Array(pl.Int64, 30)), + (Encoder(), pl.Array(pl.UInt8, (30, 4))), + (FLOAT32_ONE_HOT, pl.Array(pl.Float32, (30, 4))), + (Encoder(ambiguous="fractional", dtype="float64"), pl.Array(pl.Float64, (30, 4))), + (Encoder(dtype=np.float16), pl.Array(pl.Float16, (30, 4))), + ], +) +def test_iter_variant_seqs_dtype(encoding, polars_dtype): + batches = iter_variant_seqs(pairs_of(SAME_WIDTH), SeqReference(), use_strand=True, encoding=encoding) + assert collect(batches).schema["alt"] == polars_dtype + + +@pytest.mark.parametrize("anchor", ["center", None]) +def test_iter_variant_seqs_letters_of_two_bases(anchor): + base_pairs = Tokenizer(Alphabet([a + b for a in "ACGTN" for b in "ACGTN"])) + batches = iter_variant_seqs( + pairs_of(SAME_WIDTH), SeqReference(), use_strand=True, anchor=anchor, encoding=base_pairs + ) + if anchor is None: + # the first pair has the deletion of 2 bases across the start of the interval, which leaves 29 bases + with pytest.raises(ValueError, match="multiple of the letter length 2") as excinfo: + list(batches) + assert excinfo.value.__notes__ == ["in the group {'interval_idx': 0, 'variant_idx': 0} of pairs"] + else: + (batch,) = batches + assert batch.schema["alt"] == pl.Array(pl.UInt8, 15) + assert batch["alt"].to_list() == [base_pairs.tokenize(s).tolist() for s in batch["alt_seq"]] + + +@pytest.mark.parametrize("encoding", [Tokenizer(), FLOAT32_ONE_HOT]) +def test_iter_variant_seqs_unknown_letter(encoding): + # X at position 50 lies in the second interval of SAME_WIDTH, but in no REF + seq = CHROMS["chr1"] + reference = SeqReference({**CHROMS, "chr1": seq[:50] + "X" + seq[51:]}) + with pytest.raises(ValueError, match=r"not in the alphabet: \['X'\]") as excinfo: + list(iter_variant_seqs(pairs_of(SAME_WIDTH), reference, use_strand=True, encoding=encoding)) + assert excinfo.value.__notes__ == ["in the group {'interval_idx': 1, 'variant_idx': 5} of pairs"] + + +@pytest.mark.parametrize("encoding", [Tokenizer(), FLOAT32_ONE_HOT]) +def test_iter_variant_seqs_variable_length(encoding): + # with anchor=None, the indels change the length + out = collect( + iter_variant_seqs(pairs_of(SAME_WIDTH), SeqReference(), use_strand=True, anchor=None, encoding=encoding) + ) + inner = pl.UInt8 if isinstance(encoding, Tokenizer) else pl.Array(pl.Float32, 4) + assert out.schema["alt"] == pl.List(inner) + assert out["alt"].list.len().to_list() == out["alt_seq"].str.len_chars().to_list() + assert len(set(out["alt_seq"].str.len_chars())) > 1 + for row in out.iter_rows(named=True): + expected = ( + encoding.tokenize(row["alt_seq"]) if isinstance(encoding, Tokenizer) else encoding.one_hot(row["alt_seq"]) + ) + np.testing.assert_array_equal(np.array(row["alt"]), expected) + + +def test_iter_variant_seqs_list_for_different_widths(): + # with an anchor, alt_seq has the width of the interval, but the widths differ + out = collect(iter_variant_seqs(pairs_of(INTERVALS), SeqReference(), use_strand=True)) + assert out.schema["alt"] == pl.List(pl.UInt8) + assert out["alt"].list.len().to_list() == (out["end"] - out["start"]).to_list() + + +def test_iter_variant_seqs_list_with_equal_lengths_in_batch(): + # a batch whose sequences have equal lengths still gives a List column + batches = list(iter_variant_seqs(pairs_of(INTERVALS), SeqReference(), use_strand=True, batch_size=1)) + assert {b.schema["alt"] for b in batches} == {pl.List(pl.UInt8)} + + +def test_iter_variant_seqs_no_encoding(): + (batch,) = iter_variant_seqs(pairs_of(SAME_WIDTH), SeqReference(), use_strand=True, encoding=None) + assert batch.columns[-2:] == ["ref_seq", "alt_seq"] + + +@pytest.mark.parametrize("pad_with_n", [False, True]) +def test_iter_variant_seqs_pad_with_n(pad_with_n): + # the deletion pulls in bases from beyond the end of chr2 + intervals = pl.DataFrame({"chrom": ["chr2"], "start": [50], "end": [60], "strand": ["+"]}) + pairs = pairs_of(intervals, [variant("chr2", 55, 3, "{first}")]) + batches = iter_variant_seqs(pairs, SeqReference(), use_strand=True, anchor="start", pad_with_n=pad_with_n) + if pad_with_n: + assert collect(batches)["alt_seq"].to_list() == [CHROMS["chr2"][50:55] + CHROMS["chr2"][57:] + "NN"] + else: + with pytest.raises(ValueError, match="needs bases outside chr2") as excinfo: + list(batches) + assert excinfo.value.__notes__ == ["in the group {'interval_idx': 0, 'variant_idx': 0} of pairs"] + + +def test_iter_variant_seqs_ref_seq_past_chromosome_end(): + # the window misses an interval past the end of chr2, so each pair gets its ref_seq from extract + intervals = pl.DataFrame({"chrom": ["chr2"], "start": [50], "end": [70], "strand": ["-"]}) + pairs = pairs_of(intervals, [variant("chr2", 52, 1, "N"), variant("chr2", 56, 1, "N")]) + reference = SeqReference() + out = collect(iter_variant_seqs(pairs, reference, use_strand=True, pad_with_n=True)) + interval = Interval("chr2", 50, 70, strand="-") + expected = VariantSeqExtractor(reference, use_strand=True).extract(interval, [], pad_with_n=True) + assert expected == reverse_complement(CHROMS["chr2"][50:] + "N" * 10) + assert out["ref_seq"].to_list() == [expected, expected] + + +def test_iter_variant_seqs_rejects_non_carriers(alt_kinds_vcf): + # with carriers_only=False, the genotypes contain samples that do not carry the ALT allele + intervals = pl.DataFrame({"chrom": ["chr1"], "start": [0], "end": [30]}) + genotypes = scan_vcf_genotypes(alt_kinds_vcf, carriers_only=False) + pairs = scan_variant_pairs(intervals, genotypes) + message = "^pairs has rows whose variant_carrier is False or null, .* carriers_only=True" + with pytest.raises(ValueError, match=message): + iter_variant_seqs(pairs, SeqReference(), use_strand=True) + carriers = pairs.filter("variant_carrier").collect() + assert len(carriers) == len(scan_variant_pairs(intervals, scan_vcf_genotypes(alt_kinds_vcf)).collect()) + # a null does not tell that the sample carries the ALT allele + unknown = carriers.with_columns(pl.lit(None, dtype=pl.Boolean).alias("variant_carrier")) + with pytest.raises(ValueError, match=message): + iter_variant_seqs(unknown, SeqReference(), use_strand=True) + with FastaStringExtractor(sample_5kb_fasta_file, use_strand=False) as fasta: + out = collect(iter_variant_seqs(carriers, fasta, use_strand=True, encoding=None)) + assert len(out) == len(carriers) + + +def test_iter_variant_seqs_overlapping_variants_in_group(): + # the variants of interval_idx 0 do not overlap, but these two do + variants = [variant("chr1", 12, 1, "A"), variant("chr1", 12, 1, "C"), variant("chr1", 30, 1, "N")] + pairs = pairs_of(INTERVALS.filter(pl.col("name") == "a"), variants) + # each pair alone is fine + assert len(collect(iter_variant_seqs(pairs, SeqReference(), use_strand=True))) == 3 + with pytest.raises(ValueError, match="overlap") as excinfo: + list(iter_variant_seqs(pairs, SeqReference(), use_strand=True, group_by=["interval_idx"])) + assert excinfo.value.__notes__ == ["in the group {'interval_idx': 0} of pairs"] + + +def test_iter_variant_seqs_ref_differs(): + pairs = pairs_of(SAME_WIDTH).with_columns(variant_ref=pl.lit("NNN")) + with pytest.raises(ValueError, match="differs from the reference"): + list(iter_variant_seqs(pairs, SeqReference(), use_strand=True)) + + +def test_iter_variant_seqs_unknown_chrom(): + pairs = pairs_of(SAME_WIDTH) + reference = SeqReference({"chr1": CHROMS["chr1"]}) + with pytest.raises(ValueError, match="^chrom 'chr2' is not in the reference") as excinfo: + list(iter_variant_seqs(pairs, reference, use_strand=True)) + assert excinfo.value.__notes__ == ["in the group {'interval_idx': 2, 'variant_idx': 9} of pairs"] + + +def test_iter_variant_seqs_group_with_several_intervals(): + pairs = pairs_of(INTERVALS) + with pytest.raises( + ValueError, match=r"^the group \{'chrom': 'chr1'\} of pairs has 4 intervals, but a group needs one interval" + ): + iter_variant_seqs(pairs, SeqReference(), use_strand=True, group_by=["chrom"]) + + +@pytest.mark.parametrize( + "kwargs, error, match", + [ + ({"pairs": [1]}, TypeError, "^pairs must be a polars DataFrame or LazyFrame, not list$"), + ({"reference": sample_5kb_fasta_file}, TypeError, "^reference must be a SequenceExtractor, not str$"), + ({"group_by": "interval_idx"}, TypeError, "^group_by must be a sequence of strings, not a str"), + ({"group_by": [1]}, TypeError, "^group_by must be a sequence of strings, but contains 1$"), + ({"encoding": {"A": [1]}}, TypeError, "^encoding must be a Tokenizer, an Encoder or None, not dict$"), + ({"encoding": Alphabet(["A"])}, TypeError, "^encoding must be a Tokenizer, an Encoder or None, not Alphabet$"), + ({"dtype": np.uint8}, TypeError, "unexpected keyword argument 'dtype'"), + ({"batch_size": 1.0}, TypeError, "^batch_size must be an int, not float$"), + ({"anchor": "middle"}, ValueError, '^anchor must be "start", "center", "end" or None, got \'middle\'$'), + ({"group_by": []}, ValueError, "^group_by must name at least one column$"), + ({"group_by": ["interval_idx"] * 2}, ValueError, "^group_by must not repeat a column"), + ({"group_by": ["interval_idx", "ref"]}, ValueError, r"^group_by must not contain \['ref'\]"), + ({"group_by": ["gene_id"]}, ValueError, r"^pairs lacks the columns \['gene_id'\]$"), + ({"batch_size": 0}, ValueError, "^batch_size must be at least 1, got 0$"), + ], +) +def test_iter_variant_seqs_invalid_arguments(kwargs, error, match): + arguments = {"pairs": pairs_of(SAME_WIDTH), "reference": SeqReference(), "use_strand": True, **kwargs} + pairs = arguments.pop("pairs") + reference = arguments.pop("reference") + with pytest.raises(error, match=match): + iter_variant_seqs(pairs, reference, **arguments) + + +@pytest.mark.parametrize("column", ["strand", "variant_pos", "variant_alt"]) +def test_iter_variant_seqs_missing_column(column): + pairs = pairs_of(SAME_WIDTH).drop(column) + with pytest.raises(ValueError, match=rf"^pairs lacks the columns \['{column}'\]$"): + iter_variant_seqs(pairs, SeqReference(), use_strand=True) + + +def test_iter_variant_seqs_nulls(): + pairs = pairs_of(SAME_WIDTH).with_columns(variant_alt=pl.lit(None, dtype=pl.String)) + with pytest.raises(ValueError, match=r"^pairs has nulls in the columns \['variant_alt'\]$"): + iter_variant_seqs(pairs, SeqReference(), use_strand=True) + + +def test_iter_variant_seqs_options_are_keyword_only(): + with pytest.raises(TypeError, match="positional"): + iter_variant_seqs(pairs_of(SAME_WIDTH), SeqReference(), True) + with pytest.raises(TypeError, match="use_strand"): + iter_variant_seqs(pairs_of(SAME_WIDTH), SeqReference()) From 305e3c2abdb517bb2b919e7e41b84aa48b870d35 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florian=20R=2E=20H=C3=B6lzlwimmer?= Date: Sat, 17 Oct 2026 01:08:42 +0200 Subject: [PATCH 3/4] feat(vcf): give one row per haplotype in scan_vcf_genotypes scan_vcf_genotypes(path, haplotypes=True) gives one row per ALT allele, sample and haplotype, with the column haplotype: the position of the allele in GT, so 0 or 1 for a diploid GT. carrier then tells whether this haplotype carries the ALT allele, and carriers_only keeps only the haplotypes that carry it. A missing GT gives one row with a null haplotype. iter_variant_seqs gives the sequence of each haplotype if the genotypes are the variants of scan_variant_pairs and group_by is interval_idx, variant_sample and variant_haplotype. haplotype ignores the phasing, as bcftools consensus -H does: for an unphased GT such as 0/1, the order of the alleles in GT sets the haplotype. GT stays in the frame, so unphased GTs can be filtered. --- MIGRATING.md | 2 + README.md | 21 ++++++++++ src/kipoiseq2/extractors/batches.py | 3 ++ src/kipoiseq2/extractors/vcf.py | 38 +++++++++++++++--- tests/extractors/test_batches.py | 30 +++++++++++++++ tests/extractors/test_vcf.py | 60 +++++++++++++++++++++++++---- 6 files changed, 142 insertions(+), 12 deletions(-) diff --git a/MIGRATING.md b/MIGRATING.md index 46148e1..608d89b 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -27,6 +27,8 @@ kipoiseq2 reads VCF files with polars-bio instead of cyvcf2, so the `vcf` extra - `VariantIntervalQueryable` and the query classes of `kipoiseq.extractors.vcf_query` are gone: filter the polars frames instead. - `SingleVariantVCFSeqExtractor` and `SingleSeqVCFSeqExtractor` are gone: match the variants with `scan_variant_pairs`, and apply them with `iter_variant_seqs`, which yields batches of reference and alternative sequences. `VariantSeqExtractor` applies the variants of one interval. + For the variants that a sample carries, as with `sample_id`, take the variants from `scan_vcf_genotypes(path)` and group by `["interval_idx", "variant_sample"]`. + For one sequence per haplotype, take them from `scan_vcf_genotypes(path, haplotypes=True)` and group by `["interval_idx", "variant_sample", "variant_haplotype"]`. - `Variant.from_cyvcf` and `Variant.from_cyvcf_and_given_alt` are gone. A `Variant` has only `chrom`, `pos`, `ref` and `alt`, so `id`, `qual`, `filter`, `info`, `source` and `copy()` are gone. The VCF columns ID, QUAL, FILTER and INFO stay in the frame of `scan_vcf_variants`. diff --git a/README.md b/README.md index 93f912d..d3c1d34 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,11 @@ carriers = scan_vcf_genotypes("variants.vcf.gz", format_fields=["GT", "GQ"]) carriers.sink_csv("carriers.csv") ``` +`haplotypes=True` gives one row per haplotype of each sample, with the column `haplotype`: the position of the allele in GT, e.g. 0 or 1 for a diploid GT. +`carrier` then tells whether this haplotype carries the ALT allele. +`haplotype` ignores the phasing, as `bcftools consensus -H` does: for an unphased GT such as 0/1, the order of the alleles in GT sets the haplotype. +`~pl.col("GT").str.contains("/")` keeps the phased and haploid GTs. + ### Matching variants with intervals `scan_variant_pairs` finds every interval-variant pair with one [polars-bio](https://biodatageeks.org/polars-bio/) overlap join. @@ -203,6 +208,22 @@ A variant within 500 bases of a chromosome end gets a window that runs past the `group_by=["interval_idx"]` applies all variants of an interval together. Variants that overlap within a row raise a ValueError, as in `VariantSeqExtractor`. Each pair applies its ALT allele. So for the genotypes of `scan_vcf_genotypes`, keep the default `carriers_only=True`: a False or a null in the column `variant_carrier` raises a ValueError. +For the sequences of the haplotypes, take the variants from `scan_vcf_genotypes(path, haplotypes=True)`, and group by sample and haplotype: + +```python +from kipoiseq2.extractors import scan_vcf_genotypes + +genes = pl.DataFrame({"chrom": ["chr1"], "start": [1000], "end": [3000], "strand": ["-"], "gene_id": ["g1"]}) +# the pairs get the columns variant_sample and variant_haplotype +pairs = scan_variant_pairs(genes, scan_vcf_genotypes("variants.vcf.gz", haplotypes=True)) + +with FastaStringExtractor("genome.fa", use_strand=False) as fasta: + group_by = ["interval_idx", "variant_sample", "variant_haplotype"] + for batch in iter_variant_seqs(pairs, fasta, use_strand=True, group_by=group_by): + alt = batch["alt"].to_numpy() # one row per gene and haplotype that carries variants in it +``` + +A haplotype without variants in an interval gives no row. Its sequence is the `ref_seq` of the interval. `encoding=Encoder()` gives uint8 one-hot arrays of shape `(rows, width, 4)`, and `encoding=None` gives only the sequences. `anchor="start"` or `"end"` keeps the 5' or the 3' end of the interval, and `anchor=None` lets indels change the length of `alt_seq`. Sequences of different lengths give polars List columns instead of Array columns. diff --git a/src/kipoiseq2/extractors/batches.py b/src/kipoiseq2/extractors/batches.py index 46cf481..b787733 100644 --- a/src/kipoiseq2/extractors/batches.py +++ b/src/kipoiseq2/extractors/batches.py @@ -59,6 +59,9 @@ def iter_variant_seqs( All rows of a group must have the same interval, and the variants of a group must not overlap. The default ``group_by`` gives one row per pair, e.g. for variant effect predictions. ``group_by=["interval_idx"]`` applies all variants of an interval together. + For the haplotypes of the samples, pass ``scan_vcf_genotypes(path, haplotypes=True)`` as the variants of + :func:`~kipoiseq2.extractors.vcf_matching.scan_variant_pairs`, + and group by ``["interval_idx", "variant_sample", "variant_haplotype"]``. An interval without variants has no pairs, so it gives no row. Each pair applies its ALT allele, so ``pairs`` must not contain samples that do not carry it. For the genotypes of :func:`~kipoiseq2.extractors.vcf.scan_vcf_genotypes`, keep ``carriers_only=True``: diff --git a/src/kipoiseq2/extractors/vcf.py b/src/kipoiseq2/extractors/vcf.py index 65d5a54..14723bd 100644 --- a/src/kipoiseq2/extractors/vcf.py +++ b/src/kipoiseq2/extractors/vcf.py @@ -2,7 +2,8 @@ :func:`scan_vcf_variants` reads a VCF file with polars-bio as a polars LazyFrame, with one row per ALT allele. :func:`scan_vcf_genotypes` adds one row per sample, -with the FORMAT fields of the sample and whether it carries the ALT allele. +or per haplotype of a sample, with the FORMAT fields of the sample and whether +it carries the ALT allele. Both need the ``ranges`` extra (polars and polars-bio). """ @@ -143,6 +144,7 @@ def scan_vcf_genotypes( info_fields: Sequence[str] = (), sequence_alts_only: bool = True, carriers_only: bool = True, + haplotypes: bool = False, ) -> pl.LazyFrame: """Read the genotypes of a VCF file as a polars LazyFrame, with one row per ALT allele and sample. @@ -154,6 +156,16 @@ def scan_vcf_genotypes( the ALT alleles ``A,,C``, a sample with GT 0/3 carries C, and a sample with GT 0/2 carries none of the kept ALT alleles. + With ``haplotypes=True``, each allele of GT gives a row, with its + position in GT as haplotype: 0 or 1 for a diploid GT. carrier then + tells whether this haplotype carries the ALT allele. So GT 0|1 gives a + row for haplotype 0 that does not carry the ALT allele, and a row for + haplotype 1 that carries it. haplotype ignores the phasing, as + ``bcftools consensus -H`` does: for an unphased GT such as 0/1, the + order of the alleles in GT sets the haplotype. + ``~pl.col("GT").str.contains("/")`` keeps the phased and haploid GTs. + A missing GT gives one row with a null haplotype. + :param path: path of the VCF file. :param samples: names of the samples to read. The default reads all samples. :param format_fields: FORMAT fields to read. GT is always read. @@ -162,7 +174,10 @@ def scan_vcf_genotypes( bases, as in :func:`scan_vcf_variants`. With False, a sample with GT 0/2 for the ALT alleles ``A,,C`` carries ````. :param carriers_only: keep only the rows where the sample carries the ALT allele. - :returns: LazyFrame with the columns of :func:`scan_vcf_variants`, then sample, GT, + With ``haplotypes=True``, keep only the haplotypes that carry it. + :param haplotypes: give one row per ALT allele, sample and haplotype, as described above. + :returns: LazyFrame with the columns of :func:`scan_vcf_variants`, then sample, + haplotype (only with ``haplotypes=True``), GT, the other FORMAT fields with the value of the sample, and carrier. FORMAT fields with several values per sample, such as AD, are lists. :raises TypeError: if ``samples``, ``format_fields`` or ``info_fields`` is not @@ -194,7 +209,20 @@ def scan_vcf_genotypes( # for a single-sample VCF, polars-bio returns the FORMAT fields as columns columns = [c for c in columns if c not in fields] genotypes = variants.with_columns(sample=pl.lit(sample_names[0])) - alleles = pl.col("GT").str.extract_all(r"\d+") - carrier = alleles.list.contains(pl.col("allele_idx").cast(pl.String)).fill_null(False) - genotypes = genotypes.select(*columns, "sample", *fields, carrier.alias("carrier")) + allele_idx = pl.col("allele_idx").cast(pl.String) + if haplotypes: + # the missing alleles (.) give a row too, so that haplotype is the position of the allele in GT + alleles = pl.col("GT").str.extract_all(r"\d+|\.") + # carrier holds the allele of the haplotype until the select compares it with allele_idx + genotypes = genotypes.with_columns(carrier=alleles, haplotype=pl.int_ranges(alleles.list.len())).explode( + "carrier", "haplotype", empty_as_null=False + ) + carrier = pl.col("carrier") == allele_idx + haplotype_columns = ["haplotype"] + else: + carrier = pl.col("GT").str.extract_all(r"\d+").list.contains(allele_idx) + haplotype_columns = [] + genotypes = genotypes.select( + *columns, "sample", *haplotype_columns, *fields, carrier.fill_null(False).alias("carrier") + ) return genotypes.filter("carrier") if carriers_only else genotypes diff --git a/tests/extractors/test_batches.py b/tests/extractors/test_batches.py index a66feb7..79af1a5 100644 --- a/tests/extractors/test_batches.py +++ b/tests/extractors/test_batches.py @@ -145,6 +145,36 @@ def test_iter_variant_seqs_equals_fasta_extract(alt_kinds_vcf, use_strand, ancho assert row["ref_seq"] == fasta.extract(interval) +def test_iter_variant_seqs_haplotypes(tmp_path): + # S1 carries the first variant on haplotype 1 and the other two on haplotype 0. S3 carries none. + records = [ + (variant("chr1", 12, 1, "N"), "0|1", "1|1", "0|0"), + (variant("chr1", 20, 1, "{first}GAT"), "1|0", "0|0", "0|0"), + (variant("chr1", 25, 3, "{first}"), "1|0", "0|0", "0|0"), + ] + header = [ + "##fileformat=VCFv4.2", + "##contig=", + '##FORMAT=', + "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS1\tS2\tS3", + ] + lines = ["\t".join(["chr1", str(v.pos), ".", v.ref, v.alt, ".", ".", ".", "GT", *gts]) for v, *gts in records] + vcf = tmp_path / "haplotypes.vcf" + vcf.write_text("\n".join(header + lines) + "\n") + intervals = pl.DataFrame({"chrom": ["chr1"], "start": [10], "end": [40], "strand": ["+"]}) + pairs = scan_variant_pairs(intervals, scan_vcf_genotypes(str(vcf), haplotypes=True)) + reference = SeqReference() + group_by = ["interval_idx", "variant_sample", "variant_haplotype"] + out = collect(iter_variant_seqs(pairs, reference, use_strand=True, group_by=group_by)) + # the haplotypes without variants in the interval, e.g. those of S3, give no row + assert out.select("variant_sample", "variant_haplotype").rows() == [("S1", 0), ("S1", 1), ("S2", 0), ("S2", 1)] + v1, v2, v3 = (v for v, *_ in records) + interval = Interval("chr1", 10, 40, strand="+") + extractor = VariantSeqExtractor(reference, use_strand=True) + expected = [extractor.extract(interval, vs, anchor_pos=interval.center()) for vs in [[v2, v3], [v1], [v1], [v1]]] + assert out["alt_seq"].to_list() == expected + + @pytest.mark.parametrize( "anchor, use_strand, alt_seq", [ diff --git a/tests/extractors/test_vcf.py b/tests/extractors/test_vcf.py index b5510e6..7ef82d6 100644 --- a/tests/extractors/test_vcf.py +++ b/tests/extractors/test_vcf.py @@ -219,9 +219,53 @@ def test_scan_vcf_genotypes_all_rows(edge_case_vcf): assert df.filter((pl.col("pos") == 40) & (pl.col("sample") == "S3")).select("GT", "carrier").row(0) == (None, False) -def test_scan_vcf_genotypes_alt_kinds(alt_kinds_vcf): - df = scan_vcf_genotypes(alt_kinds_vcf).collect() - # for the ALT alleles A,,C, S1 with GT 2/3 carries C, and S2 with GT 0/1 carries A +def test_scan_vcf_genotypes_haplotypes(edge_case_vcf): + df = scan_vcf_genotypes(edge_case_vcf, haplotypes=True).collect() + assert df.select("pos", "alt", "sample", "haplotype", "GT").rows() == [ + (10, "C", "S2", 0, "1|1"), + (10, "C", "S2", 1, "1|1"), + # haplotype 1 of 0/2 carries the second ALT allele + (10, "G", "S1", 1, "0/2"), + (40, "N", "S1", 1, "0/1"), + # haploid + (5, "T", "S1", 0, "1"), + # the missing allele keeps its position in GT + (5, "T", "S3", 1, ".|1"), + # unphased: the order in GT sets the haplotype + (5, "TAA", "S2", 1, "3/2"), + (5, "N", "S2", 0, "3/2"), + ] + assert df["carrier"].all() + # the carriers per sample are the carriers per haplotype without haplotype + per_sample = scan_vcf_genotypes(edge_case_vcf).collect() + assert df.drop("haplotype").unique(maintain_order=True).equals(per_sample) + + +def test_scan_vcf_genotypes_haplotypes_all_rows(edge_case_vcf): + df = scan_vcf_genotypes(edge_case_vcf, format_fields=["GQ", "AD"], carriers_only=False, haplotypes=True).collect() + assert df.columns[10:] == ["sample", "haplotype", "GT", "GQ", "AD", "carrier"] + assert df.schema["haplotype"] == pl.Int64 + # the ALT alleles C and G of 3 diploid samples, N of 2 diploid samples and 1 missing GT, + # and T, TAA and N of 1 haploid and 2 diploid samples + assert df.height == 2 * 3 * 2 + (2 * 2 + 1) + 3 * (1 + 2 * 2) + assert df.filter(pl.col("alt") == "C").select("sample", "haplotype", "GT", "GQ", "carrier").rows() == [ + ("S1", 0, "0/2", 30, False), + ("S1", 1, "0/2", 30, False), + ("S2", 0, "1|1", 20, True), + ("S2", 1, "1|1", 20, True), + ("S3", 0, "./.", None, False), + ("S3", 1, "./.", None, False), + ] + # a missing GT gives one row with a null haplotype + missing = df.filter((pl.col("pos") == 40) & (pl.col("sample") == "S3")) + assert missing.select("haplotype", "GT", "carrier").rows() == [(None, None, False)] + + +@pytest.mark.parametrize("haplotypes", [False, True]) +def test_scan_vcf_genotypes_alt_kinds(alt_kinds_vcf, haplotypes): + df = scan_vcf_genotypes(alt_kinds_vcf, haplotypes=haplotypes).collect() + # for the ALT alleles A,,C, S1 with GT 2/3 carries C, and S2 with GT 0/1 carries A. + # Each carrier has one haplotype with the ALT allele, so haplotypes=True gives the same rows. assert df.select("pos", "alt", "allele_idx", "sample", "GT").rows() == [ (2, "A", 1, "S1", "0/1"), (4, "ACG", 1, "S1", "0/1"), @@ -262,16 +306,18 @@ def test_scan_vcf_genotypes_keyword_only(edge_case_vcf): scan_vcf_genotypes(edge_case_vcf, use_zero_based=False) -def test_scan_vcf_genotypes_samples(edge_case_vcf, tmp_path): - s2 = scan_vcf_genotypes(edge_case_vcf, samples=["S2"], format_fields=["GQ", "AD"], carriers_only=False).collect() - all_samples = scan_vcf_genotypes(edge_case_vcf, format_fields=["GQ", "AD"], carriers_only=False).collect() +@pytest.mark.parametrize("haplotypes", [False, True]) +def test_scan_vcf_genotypes_samples(edge_case_vcf, tmp_path, haplotypes): + options = {"format_fields": ["GQ", "AD"], "carriers_only": False, "haplotypes": haplotypes} + s2 = scan_vcf_genotypes(edge_case_vcf, samples=["S2"], **options).collect() + all_samples = scan_vcf_genotypes(edge_case_vcf, **options).collect() assert s2.equals(all_samples.filter(pl.col("sample") == "S2")) # polars-bio returns the FORMAT fields of a single-sample VCF as columns instead of a struct lines = [line.split("\t") for line in EDGE_CASE_VCF.splitlines()] single_sample_vcf = tmp_path / "single_sample.vcf" single_sample_vcf.write_text("".join("\t".join(line[:9] + line[10:11]) + "\n" for line in lines)) - single = scan_vcf_genotypes(str(single_sample_vcf), format_fields=["GQ", "AD"], carriers_only=False).collect() + single = scan_vcf_genotypes(str(single_sample_vcf), **options).collect() assert single.equals(s2) From be8d7207e4a2abd82334cf867f6983164356f311 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florian=20R=2E=20H=C3=B6lzlwimmer?= Date: Sat, 17 Oct 2026 14:26:54 +0200 Subject: [PATCH 4/4] feat(encoders): encode polars Series from their Arrow buffers Tokenizer.tokenize_series(seqs) and Encoder.one_hot_series(seqs) encode all sequences of a polars String Series at once. iter_variant_seqs uses them for the columns ref and alt. They build List columns without a copy, so the List columns of 10,000 sequences of 1 kb encode 3.4 to 6.6 times faster than with _encoded_series. The result keeps the name of seqs. If all sequences have the same length, it is an Array column: pl.Array(dtype, length) for tokens, or pl.Array(dtype, (length, n_columns)) for one-hot arrays. Otherwise it is a List column: pl.List(dtype) or pl.List(pl.Array(dtype, n_columns)). An empty Series gives an empty List column, because it has no length for an Array column. For letters of single ASCII characters, the methods read the offsets and the data buffer of seqs.to_arrow(), also of chunked and sliced arrays. One bytes.translate call converts the data buffer. A List column is built from the offsets and the values without a copy. Other alphabets, e.g. codons, use tokenize or one_hot for each sequence. A Series that is not of strings raises a TypeError, and nulls raise a ValueError. For the first sequence with an unknown letter, the methods raise the ValueError of tokenize or one_hot, with a note that names the row. In one_hot_series, a dtype without a polars type raises a TypeError, e.g. np.longdouble on x86-64. iter_variant_seqs drops _encoded_series. If the intervals have different widths, but the sequences of a batch have the same length by chance, it casts the Array column to a List column. So all batches share one schema. The ranges extra lists pyarrow, because the methods import it. polars-bio installed it already. Time to encode one batch into one Series, in ms. Each time is the median per call, as the mean of two rounds. old is _encoded_series at b872c8b, with tokenize, one_hot_table(DNA, neutral_value=0, dtype=np.uint8) and DNA_ONE_HOT. ref is _encoded_series at d7254f8, with Tokenizer().tokenize, Encoder().one_hot and Encoder(ambiguous="fractional", dtype=np.float32).one_hot. new is the Series methods of the same encoders. The sequences are random uppercase ACGTN. In the List cases, their lengths vary by up to 3 %. All three give the same values. 64 x 196,608 bases 10,000 x 1 kb old ref new old ref new Array tokens 24.0 7.8 6.0 18.7 4.4 4.3 Array one-hot uint8 30.3 26.5 30.5 20.7 19.2 20.7 Array one-hot float32 43.1 37.5 39.7 29.5 27.0 30.0 List tokens 24.4 5.4 8.5 43.8 28.4 4.3 List one-hot uint8 46.2 38.0 25.8 75.6 73.0 21.4 List one-hot float32 98.1 82.3 37.2 112 109 30.4 The Array one-hot cases spend most of their time in the take that writes the output array. The times of 64 x 196,608 bases swing by up to 3 ms between processes. Interleaved in one process, ref and new take the same time for these tokens: 9.7 and 9.2 ms as Array columns, 6.1 and 6.2 ms as List columns. A batch of 64 x 1 kb takes 43 instead of 30 us, which is the fixed cost of to_arrow and of the checks. iter_variant_seqs gives 57,000 pairs/s with batch_size=64 and 63,000 pairs/s with batch_size=256 for 20,000 pairs with 1 kb windows. This is the same as at b872c8b and d7254f8 within the noise of 10 %. --- README.md | 18 ++- pyproject.toml | 2 + src/kipoiseq2/encoders.py | 167 +++++++++++++++++++- src/kipoiseq2/extractors/batches.py | 134 ++++++++-------- tests/extractors/test_batches.py | 15 +- tests/test_encoders_series.py | 229 ++++++++++++++++++++++++++++ uv.lock | 2 + 7 files changed, 492 insertions(+), 75 deletions(-) create mode 100644 tests/test_encoders_series.py diff --git a/README.md b/README.md index d3c1d34..2d0e627 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,8 @@ Requires Python >= 3.12. ```bash pip install kipoiseq2 # FASTA extractors and transforms (numpy, pyfaidx) -pip install 'kipoiseq2[ranges]' # + scan_vcf_variants, scan_vcf_genotypes, scan_variant_pairs and iter_variant_seqs (polars, polars-bio) +pip install 'kipoiseq2[ranges]' # + scan_vcf_variants, scan_vcf_genotypes, scan_variant_pairs, iter_variant_seqs + # and the Series methods of the encoders (polars, polars-bio, pyarrow) ``` kipoiseq2 does not depend on pandas or pyranges. @@ -99,6 +100,20 @@ rows[4:, 4] = 1 rows[Tokenizer().tokenize("ACGTNR")] # float32 array of shape (6, 5) ``` +`Tokenizer.tokenize_series` and `Encoder.one_hot_series` encode all sequences of a polars Series at once. +They need the `ranges` extra. +If all sequences have the same length, they return a polars Array column, and its `to_numpy()` needs no copy. +Otherwise they return a List column. + +```python +import polars as pl + +seqs = pl.Series("seq", ["ACGTN", "NNACG"]) +Tokenizer().tokenize_series(seqs).to_numpy() # uint8 array of shape (2, 5) +Encoder().one_hot_series(seqs).dtype # Array(UInt8, shape=(5, 4)) +Tokenizer().tokenize_series(pl.Series(["ACGT", "A"])).dtype # List(UInt8) +``` + ### Reading VCF files as tables `scan_vcf_variants` reads a VCF file with polars-bio as a polars LazyFrame, with one row per ALT allele. @@ -184,6 +199,7 @@ By default, each pair gives one row, and the sequences keep the width of the int A batch holds the keys `interval_idx` and `variant_idx`, the interval, the sequences `ref_seq` and `alt_seq`, and the encoded sequences `ref` and `alt`. The default encoding `Tokenizer()` gives uint8 tokens. An `Encoder` as `encoding` gives one-hot arrays in the dtype of the encoder. +Each batch encodes its sequences with `tokenize_series` or `one_hot_series`. ```python import polars as pl diff --git a/pyproject.toml b/pyproject.toml index ac4f2ec..43e123a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,8 @@ dependencies = [ ranges = [ "polars>=1.37.1", "polars-bio>=0.31", + # tokenize_series and one_hot_series read the Arrow buffers of polars Series + "pyarrow", ] [project.urls] diff --git a/src/kipoiseq2/encoders.py b/src/kipoiseq2/encoders.py index e8f4fcd..f41d6fa 100644 --- a/src/kipoiseq2/encoders.py +++ b/src/kipoiseq2/encoders.py @@ -12,6 +12,9 @@ or read them with ``FastaStringExtractor``, which uppercases by default. An unknown letter raises a ValueError that names it. +:meth:`Tokenizer.tokenize_series` and :meth:`Encoder.one_hot_series` encode all sequences of a polars Series at once. +They need polars, which the ``ranges`` extra installs. + Example: >>> from kipoiseq2 import Encoder, Tokenizer @@ -25,10 +28,12 @@ array([0, 1, 2, 3, 4, 5], dtype=uint8) """ +from __future__ import annotations + import functools from collections.abc import Container, Mapping from types import MappingProxyType -from typing import Any, Literal, get_args +from typing import TYPE_CHECKING, Any, Literal, get_args import numpy as np from numpy.typing import DTypeLike, NDArray @@ -36,6 +41,10 @@ from kipoiseq2.alphabets import DNA, Alphabet from kipoiseq2.dataclasses import _check_int +if TYPE_CHECKING: + import polars as pl + import pyarrow as pa + __all__ = ["Encoder", "Tokenizer"] _SENTINEL = 255 @@ -263,6 +272,73 @@ def tokenize(self, seq: str) -> NDArray[np.integer[Any]]: except KeyError: raise self._unknown_letters_error(seq) from None + def tokenize_series(self, seqs: pl.Series) -> pl.Series: + """Convert the sequences of a polars Series to tokens. + + Each sequence gets the tokens of :meth:`tokenize`, but this method converts all sequences at once. + For letters of single ASCII characters, it translates the Arrow data buffer of ``seqs`` with one call. + The result keeps the name of ``seqs``. + If all sequences have the same length, the result is an Array column ``pl.Array(dtype, length)``. + Then ``.to_numpy()`` returns an array of shape ``(len(seqs), length)`` without a copy. + Otherwise the result is a List column ``pl.List(dtype)``. + dtype is the polars type of :attr:`dtype`, e.g. ``pl.UInt8``, and length is the number of letters. + An empty ``seqs`` gives an empty List column, because it has no length for an Array column. + + :param seqs: polars Series of strings without nulls + :returns: polars Series with the tokens of each sequence + :raises TypeError: if ``seqs`` is not a polars Series of strings + :raises ValueError: if ``seqs`` contains nulls, or if :meth:`tokenize` raises a ValueError for a sequence. + For the first such sequence, it raises the ValueError of :meth:`tokenize`, with a note that names the row. + + Example: + + >>> import polars as pl + >>> tokens = Tokenizer().tokenize_series(pl.Series("seq", ["ACG", "TNA"])) + >>> tokens.dtype, tokens.to_list() + (Array(UInt8, shape=(3,)), [[0, 1, 2], [3, 4, 0]]) + """ + tokens, offsets = self._series_tokens(seqs) + return _series(seqs.name, tokens, offsets) + + def _series_tokens(self, seqs: pl.Series) -> tuple[NDArray[np.integer[Any]], NDArray[np.int64]]: + """Return the tokens of all sequences of ``seqs``, and the offsets of the sequences in the tokens. + + :returns: the tokens, and ``len(seqs) + 1`` offsets that start at 0. + The tokens of ``seqs[i]`` are ``tokens[offsets[i]:offsets[i + 1]]``. + :raises TypeError: as :meth:`tokenize_series` + :raises ValueError: as :meth:`tokenize_series` + """ + import polars as pl + + _check_series(seqs) + if self._table is None: + # the slow path tokenizes each sequence on its own + arrays = [] + for i, seq in enumerate(seqs.to_list()): + try: + arrays.append(self.tokenize(seq)) + except ValueError as e: + e.add_note(f"in the row {i} of seqs") + raise + offsets = np.zeros(len(arrays) + 1, dtype=np.int64) + np.cumsum([len(tokens) for tokens in arrays], out=offsets[1:]) + return np.concatenate([np.empty(0, dtype=self._dtype), *arrays]), offsets + # the oldest compat_level gives large_string, while a newer one can give string_view + offsets, data = _arrow_strings(seqs.to_arrow(compat_level=pl.CompatLevel.oldest())) + # translate needs bytes, so bytes copies the data first. + # The table has only ASCII letters, so a non-ASCII byte gives the sentinel too. + rows = bytes(data).translate(self._table) + bad = rows.find(_SENTINEL) + if bad >= 0: + # the byte offsets give the row of the first unknown letter + i = int(np.searchsorted(offsets, bad, side="right")) - 1 + error = self._unknown_letters_error(seqs[i]) + error.add_note(f"in the row {i} of seqs") + raise error + # all bytes are ASCII letters, so the byte offsets are the offsets of the letters + tokens = np.frombuffer(rows, dtype=np.uint8) + return (tokens if self._ids is None else self._ids.take(tokens)), offsets + def _unknown_letters_error(self, seq: str) -> ValueError: """Return the ValueError that names the letters of ``seq`` that are not in :attr:`alphabet`.""" n = self._letter_length @@ -271,6 +347,64 @@ def _unknown_letters_error(self, seq: str) -> ValueError: return _unknown_letters_error(unknown, self._tokens, "letters", "the alphabet") +def _check_series(seqs: object) -> None: + """Raise a TypeError if ``seqs`` is not a polars Series of strings, and a ValueError if it contains nulls.""" + import polars as pl + + if not isinstance(seqs, pl.Series): + raise TypeError(f"seqs must be a polars Series, not {type(seqs).__name__}") + if seqs.dtype != pl.String: + raise TypeError(f"seqs must be a polars Series of strings, got the dtype {seqs.dtype}") + if seqs.null_count(): + raise ValueError(f"seqs must not contain nulls, but the row {seqs.is_null().arg_true()[0]} is null") + + +def _arrow_strings(arr: pa.Array) -> tuple[NDArray[np.int64], memoryview]: + """Return the byte offsets and the bytes of the strings of an Arrow array. + + :param arr: Arrow array of large_string without nulls, e.g. of ``pl.Series.to_arrow`` + with ``compat_level=pl.CompatLevel.oldest()`` + :returns: ``len(arr) + 1`` offsets that start at 0, and the bytes of all strings. + The string ``i`` is ``data[offsets[i]:offsets[i + 1]]``. + """ + _, offsets_buffer, data_buffer = arr.buffers() + # a sliced array starts at arr.offset in its buffers + offsets = np.frombuffer(offsets_buffer, dtype=np.int64, count=len(arr) + 1, offset=8 * arr.offset) + start, end = int(offsets[0]), int(offsets[-1]) + return offsets - start, memoryview(data_buffer)[start:end] + + +def _series(name: str, values: NDArray[Any], offsets: NDArray[np.int64]) -> pl.Series: + """Return the encoded sequences as one polars Series, without a copy of ``values``. + + :param name: the name of the Series + :param values: the tokens or the one-hot rows of the letters of all sequences, one sequence after the other + :param offsets: ``n + 1`` offsets that start at 0. The sequence ``i`` is ``values[offsets[i]:offsets[i + 1]]``. + :returns: an Array column if the ``n`` sequences have the same length, and a List column otherwise + :raises TypeError: if polars has no type for the dtype of ``values`` + """ + import polars as pl + import pyarrow as pa + + if not values.dtype.isnative: + # Arrow holds numbers in the native byte order + values = values.astype(values.dtype.newbyteorder("=")) + try: + pa.from_numpy_dtype(values.dtype) + except NotImplementedError: + raise TypeError(f"polars has no type for the dtype {values.dtype}") from None + lengths = np.diff(offsets) + if lengths.size and (lengths == lengths[0]).all(): + # polars keeps the array without a copy, except for bool, which Arrow packs into bits + return pl.Series(name, values.reshape(len(lengths), int(lengths[0]), *values.shape[1:])) + # pa.array keeps the values without a copy, except for bool, which Arrow packs into bits + items = pa.array(values.reshape(-1)) + if values.ndim == 2: + # one array of n_columns values per letter + items = pa.FixedSizeListArray.from_arrays(items, values.shape[1]) + return pl.Series(name, pa.LargeListArray.from_arrays(pa.array(offsets), items)) + + class Encoder: """One-hot encode sequences. @@ -403,6 +537,37 @@ def one_hot(self, seq: str) -> NDArray[Any]: """ return self._rows.take(self._tokenizer.tokenize(seq), axis=0) + def one_hot_series(self, seqs: pl.Series) -> pl.Series: + """One-hot encode the sequences of a polars Series. + + Each sequence gets the array of :meth:`one_hot`, but this method encodes all sequences at once, + as :meth:`Tokenizer.tokenize_series` does. + The result keeps the name of ``seqs``. + If all sequences have the same length, the result is an Array column ``pl.Array(dtype, (length, n_columns))``. + Then ``.to_numpy()`` returns an array of shape ``(len(seqs), length, n_columns)``, + without a copy unless :attr:`dtype` is bool. + Otherwise the result is a List column ``pl.List(pl.Array(dtype, n_columns))``. + dtype is the polars type of :attr:`dtype`, e.g. ``pl.UInt8``, length is the number of letters, + and n_columns is ``len(alphabet.letters)``. + An empty ``seqs`` gives an empty List column, because it has no length for an Array column. + + :param seqs: polars Series of strings without nulls + :returns: polars Series with the one-hot array of each sequence + :raises TypeError: if ``seqs`` is not a polars Series of strings, + or if polars has no type for :attr:`dtype`, e.g. for ``np.longdouble`` on x86-64 + :raises ValueError: if ``seqs`` contains nulls, or if :meth:`one_hot` raises a ValueError for a sequence. + For the first such sequence, it raises the ValueError of :meth:`one_hot`, with a note that names the row. + + Example: + + >>> import polars as pl + >>> one_hot = Encoder().one_hot_series(pl.Series("seq", ["A", "CN"])) + >>> one_hot.dtype, one_hot.to_list() + (List(Array(UInt8, shape=(4,))), [[[1, 0, 0, 0]], [[0, 1, 0, 0], [0, 0, 0, 0]]]) + """ + tokens, offsets = self._tokenizer._series_tokens(seqs) + return _series(seqs.name, self._rows.take(tokens, axis=0), offsets) + def decode(self, x: NDArray[Any]) -> str: """Convert a one-hot encoded array back to a sequence. diff --git a/src/kipoiseq2/extractors/batches.py b/src/kipoiseq2/extractors/batches.py index b787733..5dd164d 100644 --- a/src/kipoiseq2/extractors/batches.py +++ b/src/kipoiseq2/extractors/batches.py @@ -11,9 +11,6 @@ from collections.abc import Callable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any -import numpy as np -from numpy.typing import NDArray - from kipoiseq2.dataclasses import _ANCHORS, Interval, Variant, _Anchor, _check_int, _is_reverse from kipoiseq2.encoders import Encoder, Tokenizer from kipoiseq2.extractors.base import SequenceExtractor @@ -23,6 +20,7 @@ if TYPE_CHECKING: import polars as pl + from numpy.typing import NDArray __all__ = ["iter_variant_seqs"] @@ -74,8 +72,8 @@ def iter_variant_seqs( The encoded columns are polars Array columns if all sequences have the same length, i.e. if ``anchor`` is not None and all intervals of ``pairs`` have the same width. Then ``batch["alt"].to_numpy()`` returns a read-only array of shape ``(rows, length)`` for a Tokenizer, - or ``(rows, length, n_columns)`` for an Encoder, without a copy. - Otherwise the encoded columns are polars List columns, with one item per letter. + or ``(rows, length, n_columns)`` for an Encoder, without a copy unless the dtype is bool. + Otherwise the encoded columns of every batch are polars List columns, with one item per letter. :param pairs: DataFrame or LazyFrame with the columns chrom, start, end, strand, variant_pos, variant_ref, variant_alt and the columns of ``group_by``, @@ -93,8 +91,8 @@ def iter_variant_seqs( None gives alt_seq the length that the variants give, as ``anchor_pos=None``. :param group_by: the columns of ``pairs`` that form the groups :param encoding: a :class:`~kipoiseq2.encoders.Tokenizer` for the tokens of - :meth:`~kipoiseq2.encoders.Tokenizer.tokenize`, an :class:`~kipoiseq2.encoders.Encoder` for the one-hot arrays of - :meth:`~kipoiseq2.encoders.Encoder.one_hot`, or None for no encoded columns. + :meth:`~kipoiseq2.encoders.Tokenizer.tokenize_series`, an :class:`~kipoiseq2.encoders.Encoder` for the one-hot + arrays of :meth:`~kipoiseq2.encoders.Encoder.one_hot_series`, or None for no encoded columns. The dtype of ``encoding`` sets the dtype of the encoded columns, e.g. ``Encoder(dtype=np.float16)``. :param pad_with_n: if True, use N for the bases outside the chromosome. Otherwise raise a ValueError for them. :param batch_size: number of rows per batch. The last batch can have fewer rows. @@ -105,9 +103,9 @@ def iter_variant_seqs( - ref_seq: the reference sequence of the interval - alt_seq: the sequence of the interval with the variants of the group applied - ref and alt: ref_seq and alt_seq encoded, unless ``encoding`` is None. - For a Tokenizer, they are ``pl.Array(dtype, length)`` with the polars type of its dtype, - e.g. ``pl.UInt8`` for the default ``Tokenizer()``. - For an Encoder, they are ``pl.Array(dtype, (length, n_columns))`` with the polars type of its dtype. + For a Tokenizer, they are ``pl.Array(dtype, length)`` or ``pl.List(dtype)``, with the polars type of its + dtype, e.g. ``pl.UInt8`` for the default ``Tokenizer()``. + For an Encoder, they are ``pl.Array(dtype, (length, n_columns))`` or ``pl.List(pl.Array(dtype, n_columns))``. length is the number of letters. :raises TypeError: if ``pairs`` is not a polars DataFrame or LazyFrame, if ``reference`` is not a :class:`~kipoiseq2.extractors.base.SequenceExtractor`, @@ -146,14 +144,13 @@ def iter_variant_seqs( clash = [c for c in keys if c in _SEQ_COLUMNS] if clash: raise ValueError(f"group_by must not contain {clash}, because the batches use these names for the sequences") - encoder: Callable[[str], NDArray[Any]] | None = None - letter_length = 1 + # the method for a Series encodes the batches, and the method for one sequence finds the group of an error + encode_series: Callable[[pl.Series], pl.Series] | None = None + encode: Callable[[str], NDArray[Any]] | None = None if isinstance(encoding, Tokenizer): - encoder = encoding.tokenize - letter_length = len(encoding.alphabet.letters[0]) + encode_series, encode = encoding.tokenize_series, encoding.tokenize elif isinstance(encoding, Encoder): - encoder = encoding.one_hot - letter_length = len(encoding.alphabet.letters[0]) + encode_series, encode = encoding.one_hot_series, encoding.one_hot elif encoding is not None: raise TypeError(f"encoding must be a Tokenizer, an Encoder or None, not {type(encoding).__name__}") _check_int(batch_size, "batch_size") @@ -209,8 +206,8 @@ def iter_variant_seqs( key_names=key_names, use_strand=use_strand, anchor=anchor, - encoder=encoder, - letter_length=letter_length, + encode_series=encode_series, + encode=encode, same_length=same_length, pad_with_n=pad_with_n, batch_size=batch_size, @@ -225,8 +222,8 @@ def _iter_batches( key_names: list[str], use_strand: bool, anchor: _Anchor | None, - encoder: Callable[[str], NDArray[Any]] | None, - letter_length: int, + encode_series: Callable[[pl.Series], pl.Series] | None, + encode: Callable[[str], NDArray[Any]] | None, same_length: bool, pad_with_n: bool, batch_size: int, @@ -237,9 +234,10 @@ def _iter_batches( the variant columns as lists, and margin :param keys: the columns of ``group_by`` :param key_names: the names of the keys in ``groups`` - :param encoder: the function that encodes a sequence, or None for no encoded columns - :param letter_length: the length of a letter of ``encoder`` - :param same_length: if True, all sequences have the same length, so the encoded columns are Array columns + :param encode_series: the method that encodes a Series of sequences, or None for no encoded columns + :param encode: the method that encodes one sequence, or None for no encoded columns + :param same_length: if True, all sequences have the same length, so the encoded columns are Array columns. + Otherwise they are List columns. """ import polars as pl @@ -273,31 +271,61 @@ def _iter_batches( ref_seqs.append(ref_seq) alt_seqs.append(alt_seq) + ref_series = pl.Series("ref_seq", ref_seqs, dtype=pl.String) + alt_series = pl.Series("alt_seq", alt_seqs, dtype=pl.String) columns = [ *(batch.get_column(name).alias(k) for k, name in zip(keys, key_names)), *(batch.get_column(c) for c in _INTERVAL_COLUMNS if c not in keys), - pl.Series("ref_seq", ref_seqs, dtype=pl.String), - pl.Series("alt_seq", alt_seqs, dtype=pl.String), + ref_series, + alt_series, ] - if encoder is not None: + if encode_series is not None and encode is not None: try: columns += [ - _encoded_series("ref", ref_seqs, encoder, letter_length=letter_length, same_length=same_length), - _encoded_series("alt", alt_seqs, encoder, letter_length=letter_length, same_length=same_length), + _encoded(encode_series, ref_series, same_length=same_length).alias("ref"), + _encoded(encode_series, alt_series, same_length=same_length).alias("alt"), ] except ValueError: - # encode row by row to find the group of the error - for i, seqs in enumerate(zip(ref_seqs, alt_seqs)): - try: - for seq in seqs: - encoder(seq) - except ValueError as e: - e.add_note(_group_note(batch, i, keys, key_names)) - raise + _raise_with_group(encode, ref_seqs, alt_seqs, batch, keys, key_names) raise yield pl.DataFrame(columns) +def _encoded(encode_series: Callable[[pl.Series], pl.Series], seqs: pl.Series, *, same_length: bool) -> pl.Series: + """Return the encoded ``seqs``, as an Array column if ``same_length`` is True, and as a List column otherwise. + + :raises ValueError: if ``encode_series`` cannot encode a sequence + """ + import polars as pl + + series = encode_series(seqs) + if not same_length and isinstance(series.dtype, pl.Array): + # the sequences of this batch have the same length by chance, but all batches share one schema + series = series.cast(pl.List(series.dtype.inner)) + return series + + +def _raise_with_group( + encode: Callable[[str], NDArray[Any]], + ref_seqs: list[str], + alt_seqs: list[str], + batch: pl.DataFrame, + keys: list[str], + key_names: list[str], +) -> None: + """Encode the sequences row by row, and raise the first ValueError with a note that names the group of its row. + + If no sequence raises a ValueError, it returns, and the caller raises its own error. + """ + for i, (ref_seq, alt_seq) in enumerate(zip(ref_seqs, alt_seqs)): + try: + encode(ref_seq) + encode(alt_seq) + except ValueError as e: + e.add_note(_group_note(batch, i, keys, key_names)) + raise + + def _group_note(batch: pl.DataFrame, i: int, keys: list[str], key_names: list[str]) -> str: """Return the note that names the group of row ``i`` of ``batch`` for an error.""" group = {k: batch.get_column(name)[i] for k, name in zip(keys, key_names)} @@ -321,40 +349,6 @@ def _ref_seq( return seq -def _encoded_series( - name: str, seqs: list[str], encoder: Callable[[str], NDArray[Any]], *, letter_length: int, same_length: bool -) -> pl.Series: - """Return the encoded sequences of a batch as one polars Series. - - ``encoder`` encodes the joined sequences in one call, which is faster than one call per sequence. - - :param seqs: the sequences of the rows - :param encoder: the function that encodes a sequence into an array with one item per letter - :param letter_length: the length of a letter of ``encoder`` - :param same_length: if True, all sequences have the same length, and the Series is an Array column. - Otherwise it is a List column with one item per letter. - :raises ValueError: if ``encoder`` cannot encode a sequence - """ - import polars as pl - - lengths = [len(seq) // letter_length for seq in seqs] - letters = encoder("".join(seqs)) - if len(letters) != sum(lengths): - # the joined sequences can split into letters even if a sequence cannot - raise ValueError(f"the length of a sequence is not a multiple of the letter length {letter_length}") - if same_length: - # polars keeps the array without a copy - return pl.Series(name, letters.reshape(len(seqs), lengths[0], *letters.shape[1:])) - arrays = np.split(letters, np.cumsum(lengths[:-1])) - inner = pl.Series(np.empty(0, dtype=letters.dtype)).dtype - # polars builds an Array column if all arrays have the same length, so cast it - series = pl.Series(name, [a.ravel() for a in arrays]).cast(pl.List(inner)) - if letters.ndim == 2: - # one Array of n_columns values per letter - series = series.list.eval(pl.element().reshape((-1, letters.shape[1]))) - return series - - class _WindowReference: """The reference of :func:`iter_variant_seqs`, which serves the sequences of an interval from one fetch. diff --git a/tests/extractors/test_batches.py b/tests/extractors/test_batches.py index 79af1a5..0e58e7e 100644 --- a/tests/extractors/test_batches.py +++ b/tests/extractors/test_batches.py @@ -370,10 +370,19 @@ def test_iter_variant_seqs_list_for_different_widths(): assert out["alt"].list.len().to_list() == (out["end"] - out["start"]).to_list() -def test_iter_variant_seqs_list_with_equal_lengths_in_batch(): +@pytest.mark.parametrize( + "encoding, polars_dtype", [(Tokenizer(), pl.List(pl.UInt8)), (FLOAT32_ONE_HOT, pl.List(pl.Array(pl.Float32, 4)))] +) +def test_iter_variant_seqs_list_with_equal_lengths_in_batch(encoding, polars_dtype): # a batch whose sequences have equal lengths still gives a List column - batches = list(iter_variant_seqs(pairs_of(INTERVALS), SeqReference(), use_strand=True, batch_size=1)) - assert {b.schema["alt"] for b in batches} == {pl.List(pl.UInt8)} + batches = list( + iter_variant_seqs(pairs_of(INTERVALS), SeqReference(), use_strand=True, encoding=encoding, batch_size=1) + ) + assert {b.schema[c] for b in batches for c in ["ref", "alt"]} == {polars_dtype} + out = pl.concat(batches) + encode = encoding.tokenize if isinstance(encoding, Tokenizer) else encoding.one_hot + for column, seq_column in [("ref", "ref_seq"), ("alt", "alt_seq")]: + assert out[column].to_list() == [encode(seq).tolist() for seq in out[seq_column]] def test_iter_variant_seqs_no_encoding(): diff --git a/tests/test_encoders_series.py b/tests/test_encoders_series.py new file mode 100644 index 0000000..de2fea4 --- /dev/null +++ b/tests/test_encoders_series.py @@ -0,0 +1,229 @@ +import pickle +import random +import re + +import numpy as np +import polars as pl +import pyarrow as pa +import pytest + +from kipoiseq2 import DNA, Alphabet, Encoder, Tokenizer +from kipoiseq2.encoders import _series + +CODONS = Alphabet(["ACG", "TGA"], ambiguous={"NNN": ["ACG", "TGA"]}) +"""An alphabet of letters with three characters.""" + +NON_ASCII = Alphabet(DNA.letters, ambiguous={**DNA.ambiguous, "é": DNA.letters}) +""":data:`DNA` with the extra code ``é`` for all letters, which turns off the translate table.""" + +ENCODERS = [ + Tokenizer(), + Tokenizer(dtype=np.int64), + # the token 255 is the sentinel of the translate table, so tokenize takes the tokens from an array + Tokenizer(token_map={"A": 7, "C": 8, "G": 255, "T": 10, **dict.fromkeys(DNA.ambiguous, 11)}), + Tokenizer(dtype=">i4"), + Tokenizer(NON_ASCII), + Tokenizer(CODONS, dtype=np.int16), + Encoder(), + Encoder(ambiguous="fractional", dtype=np.float32), + Encoder(dtype=bool), + Encoder(ambiguous="fractional", dtype=np.float16), + Encoder(ambiguous="fractional", dtype=">f4"), + Encoder(ambiguous="ones"), + Encoder(ambiguous="ones", dtype=bool), + Encoder(NON_ASCII), + Encoder(NON_ASCII, ambiguous="ones"), + Encoder(CODONS, ambiguous="fractional", dtype=np.float64), +] +"""Encoders with the translate table and with the slow path, in several dtypes.""" + +LENGTHS = { + "same lengths": [5, 5, 5], + "different lengths": [5, 0, 3, 1, 8], + "one sequence": [7], + "empty sequences": [0, 0], +} +"""The numbers of letters of the sequences of the test Series.""" + + +def _encode(encoder, seq): + return encoder.tokenize(seq) if isinstance(encoder, Tokenizer) else encoder.one_hot(seq) + + +def _encode_series(encoder, seqs): + return encoder.tokenize_series(seqs) if isinstance(encoder, Tokenizer) else encoder.one_hot_series(seqs) + + +def _seqs(encoder, lengths): + """Return random sequences of the letters of ``encoder`` with ``lengths`` letters each.""" + rng = random.Random(0) + letters = [*encoder.alphabet.letters, *encoder.alphabet.ambiguous] + return ["".join(rng.choices(letters, k=length)) for length in lengths] + + +def _polars_dtype(encoder, lengths): + """Return the polars dtype of the encoded sequences with ``lengths`` letters.""" + inner = pl.Series(np.empty(0, dtype=encoder.dtype.newbyteorder("="))).dtype + if isinstance(encoder, Encoder): + n_columns = len(encoder.alphabet.letters) + if lengths and len(set(lengths)) == 1: + return pl.Array(inner, (lengths[0], n_columns)) + return pl.List(pl.Array(inner, n_columns)) + if lengths and len(set(lengths)) == 1: + return pl.Array(inner, lengths[0]) + return pl.List(inner) + + +def _check(encoder, series, out): + """Check that ``out`` holds the arrays of the per-sequence method for the sequences of ``series``.""" + expected = [_encode(encoder, seq) for seq in series.to_list()] + assert out.name == series.name + assert out.dtype == _polars_dtype(encoder, [len(x) for x in expected]) + assert out.to_list() == [x.tolist() for x in expected] + + +@pytest.mark.parametrize("lengths", LENGTHS.values(), ids=LENGTHS.keys()) +@pytest.mark.parametrize("encoder", ENCODERS, ids=repr) +def test_series_equals_the_method_per_sequence(encoder, lengths): + series = pl.Series("seq", _seqs(encoder, lengths)) + _check(encoder, series, _encode_series(encoder, series)) + + +@pytest.mark.parametrize("encoder", [Tokenizer(), Encoder(ambiguous="fractional", dtype=np.float32)], ids=repr) +def test_series_array_column_to_numpy(encoder): + series = pl.Series("seq", ["ACGTN", "NNACG", "TTTTT"]) + x = _encode_series(encoder, series).to_numpy() + np.testing.assert_array_equal(x, np.stack([_encode(encoder, seq) for seq in series]), strict=True) + + +@pytest.mark.parametrize("encoder", [Tokenizer(), Tokenizer(dtype=np.int64), Encoder(), Encoder(dtype=np.float32)]) +def test_series_array_column_to_numpy_without_copy(encoder): + out = _encode_series(encoder, pl.Series(["ACGTN", "NNACG"])) + # each call returns a view of the same polars buffer + assert np.shares_memory(out.to_numpy(), out.to_numpy()) + + +@pytest.mark.parametrize( + "values, offsets", + [ + (np.arange(12, dtype=np.uint8), [0, 6, 12]), + (np.arange(12, dtype=np.uint8), [0, 4, 12]), + (np.arange(24, dtype=np.float32).reshape(6, 4), [0, 3, 6]), + (np.arange(24, dtype=np.float32).reshape(6, 4), [0, 2, 6]), + ], + ids=["Array of tokens", "List of tokens", "Array of one-hot rows", "List of one-hot rows"], +) +def test_series_keeps_the_values_without_a_copy(values, offsets): + arrow = _series("seq", values, np.array(offsets, dtype=np.int64)).to_arrow() + while pa.types.is_nested(arrow.type): + arrow = arrow.values + assert arrow.buffers()[1].address == values.ctypes.data + + +@pytest.mark.parametrize("encoder", [Tokenizer(), Tokenizer(CODONS), Encoder(), Encoder(CODONS)], ids=repr) +def test_series_of_several_chunks_or_a_slice(encoder): + seqs = _seqs(encoder, [5, 0, 3, 1, 8, 2]) + chunks = pl.concat([pl.Series("seq", seqs[:2]), pl.Series("seq", seqs[2:])], rechunk=False) + assert chunks.n_chunks() == 2 + _check(encoder, chunks, _encode_series(encoder, chunks)) + sliced = pl.Series("seq", seqs).slice(1, 4) + _check(encoder, sliced, _encode_series(encoder, sliced)) + + +def _sliced_arrow(seqs): + """Return the Arrow array of ``seqs`` as a slice of a longer array, whose other strings have unknown letters.""" + return pa.array(["XXX", *seqs, "UU"], type=pa.large_string()).slice(1, len(seqs)) + + +@pytest.mark.parametrize("encoder", [Tokenizer(), Encoder()], ids=repr) +def test_series_from_a_sliced_arrow_array(monkeypatch, encoder): + # polars gives offset 0 for a sliced Series, but an Arrow array can start at an offset in its buffers + series = pl.Series("seq", ["ACGTN", "", "GGA", "NNNNNNN", "T", "AC"]) + monkeypatch.setattr(pl.Series, "to_arrow", lambda self, *args, **kwargs: _sliced_arrow(self.to_list())) + assert series.to_arrow().offset == 1 + # the unknown letters outside of the slice must not raise an error + _check(encoder, series, _encode_series(encoder, series)) + + +@pytest.mark.parametrize( + "encoder, dtype", + [ + (Tokenizer(), pl.List(pl.UInt8)), + (Tokenizer(dtype=np.int64), pl.List(pl.Int64)), + (Tokenizer(CODONS, dtype=np.int64), pl.List(pl.Int64)), + (Encoder(), pl.List(pl.Array(pl.UInt8, 4))), + (Encoder(ambiguous="fractional", dtype=np.float32), pl.List(pl.Array(pl.Float32, 4))), + (Encoder(CODONS), pl.List(pl.Array(pl.UInt8, 2))), + ], + ids=repr, +) +def test_series_empty(encoder, dtype): + # an empty Series has no length for an Array column, so it gives an empty List column + out = _encode_series(encoder, pl.Series("seq", [], dtype=pl.String)) + assert out.name == "seq" + assert out.dtype == dtype + assert out.is_empty() + + +@pytest.mark.parametrize("encoder", [Tokenizer(), Tokenizer(NON_ASCII), Encoder(), Encoder(NON_ASCII)], ids=repr) +@pytest.mark.parametrize( + "seqs, row", + [ + (["ACGT", "ACUT", "AXGT"], 1), + (["ACGT", "", "acgt"], 2), + (["ñA", "A"], 0), + (["ACGT", "AC\xffT"], 1), + ], +) +def test_series_unknown_letters(encoder, seqs, row): + # the error is the one of the method per sequence for the first sequence with an unknown letter + with pytest.raises(ValueError) as expected: + _encode(encoder, seqs[row]) + with pytest.raises(ValueError) as excinfo: + _encode_series(encoder, pl.Series(seqs)) + assert excinfo.value.args == expected.value.args + assert excinfo.value.__notes__ == [f"in the row {row} of seqs"] + + +@pytest.mark.parametrize("encoder", [Tokenizer(CODONS), Encoder(CODONS)], ids=repr) +def test_series_length_not_a_multiple_of_the_letter_length(encoder): + with pytest.raises(ValueError, match=r"multiple of the letter length 3, got len\(seq\)=4") as excinfo: + _encode_series(encoder, pl.Series(["ACGTGA", "ACGT"])) + assert excinfo.value.__notes__ == ["in the row 1 of seqs"] + + +@pytest.mark.parametrize("encoder", [Tokenizer(), Encoder(), Tokenizer(CODONS)], ids=repr) +@pytest.mark.parametrize( + "seqs, match", + [ + (["ACGT"], "seqs must be a polars Series, not list"), + ("ACGT", "seqs must be a polars Series, not str"), + (pl.Series(["ACGT"]).to_frame(), "seqs must be a polars Series, not DataFrame"), + (pl.Series([1, 2]), "seqs must be a polars Series of strings, got the dtype Int64"), + (pl.Series([b"ACGT"]), "seqs must be a polars Series of strings, got the dtype Binary"), + (pl.Series(["ACGT"], dtype=pl.Categorical), "seqs must be a polars Series of strings, got the dtype Categ"), + ], +) +def test_series_not_a_series_of_strings(encoder, seqs, match): + with pytest.raises(TypeError, match=f"^{re.escape(match)}"): + _encode_series(encoder, seqs) + + +@pytest.mark.parametrize("encoder", [Tokenizer(), Encoder(), Tokenizer(CODONS)], ids=repr) +def test_series_nulls(encoder): + with pytest.raises(ValueError, match=r"^seqs must not contain nulls, but the row 1 is null$"): + _encode_series(encoder, pl.Series(["ACG", None, None])) + + +@pytest.mark.skipif(np.dtype(np.longdouble).itemsize <= 8, reason="np.longdouble is float64 on this platform") +def test_one_hot_series_dtype_without_polars_type(): + with pytest.raises(TypeError, match=f"^polars has no type for the dtype {np.dtype(np.longdouble)}$"): + Encoder(dtype=np.longdouble).one_hot_series(pl.Series(["ACGT"])) + + +def test_series_after_pickle(): + series = pl.Series("seq", ["ACGTN", "AC"]) + token_map = {"A": 7, "C": 8, "G": 9, "T": 10, **dict.fromkeys(DNA.ambiguous, 11)} + for encoder in [Tokenizer(token_map=token_map), Encoder(dtype=np.float32)]: + copied = pickle.loads(pickle.dumps(encoder)) + assert _encode_series(copied, series).to_list() == _encode_series(encoder, series).to_list() diff --git a/uv.lock b/uv.lock index b8bfd7d..1ed41e7 100644 --- a/uv.lock +++ b/uv.lock @@ -492,6 +492,7 @@ dependencies = [ ranges = [ { name = "polars" }, { name = "polars-bio" }, + { name = "pyarrow" }, ] [package.dev-dependencies] @@ -521,6 +522,7 @@ requires-dist = [ { name = "numpy" }, { name = "polars", marker = "extra == 'ranges'", specifier = ">=1.37.1" }, { name = "polars-bio", marker = "extra == 'ranges'", specifier = ">=0.31" }, + { name = "pyarrow", marker = "extra == 'ranges'" }, { name = "pyfaidx" }, ] provides-extras = ["ranges"]