feat(extractors): yield batches of encoded ref and alt sequences

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`.

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: <https://gagneurlab.github.io/kipoiseq2/>

docs/api/extractors.md

@@ -20,4 +20,8 @@
 .. automodule:: kipoiseq2.extractors.vcf_matching
    :members:
    :undoc-members:
+
+.. automodule:: kipoiseq2.extractors.batches
+   :members:
+   :undoc-members:
 ```

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]

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__}")