diff --git a/MIGRATING.md b/MIGRATING.md index 34c5461..608d89b 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -25,7 +25,10 @@ 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. + 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`. @@ -41,6 +44,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)`. @@ -217,6 +221,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/README.md b/README.md index 3f31c6f..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 and scan_variant_pairs (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. @@ -129,6 +144,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. @@ -172,6 +192,58 @@ 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. +Each batch encodes its sequences with `tokenize_series` or `one_hot_series`. + +```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. +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. + 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/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/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/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/__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..5dd164d --- /dev/null +++ b/src/kipoiseq2/extractors/batches.py @@ -0,0 +1,399 @@ +"""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 + +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 + from numpy.typing import NDArray + +__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. + 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``: + 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 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``, + 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_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. + :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)`` 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`, + 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") + # 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): + encode_series, encode = encoding.tokenize_series, encoding.tokenize + elif isinstance(encoding, Encoder): + 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") + 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, + encode_series=encode_series, + encode=encode, + 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, + encode_series: Callable[[pl.Series], pl.Series] | None, + encode: Callable[[str], NDArray[Any]] | None, + 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 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 + + 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) + + 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), + ref_series, + alt_series, + ] + if encode_series is not None and encode is not None: + try: + columns += [ + _encoded(encode_series, ref_series, same_length=same_length).alias("ref"), + _encoded(encode_series, alt_series, same_length=same_length).alias("alt"), + ] + except ValueError: + _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)} + 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 + + +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/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/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_batches.py b/tests/extractors/test_batches.py new file mode 100644 index 0000000..0e58e7e --- /dev/null +++ b/tests/extractors/test_batches.py @@ -0,0 +1,515 @@ +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) + + +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", + [ + # 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() + + +@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, 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(): + (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()) 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) 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 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"]