[Hoeze]

tl;dr: iter_variant_seqs applies the variants of the pairs of scan_variant_pairs to the reference, and yields polars DataFrame batches with the reference and alternative sequences, encoded as uint8 tokens or one-hot arrays. Tokenizer.tokenize_series and Encoder.one_hot_series encode all sequences of a polars Series at once, from its Arrow buffers. scan_vcf_genotypes(path, haplotypes=True) gives one row per haplotype, so the batches can hold the sequence of each haplotype. VariantSeqExtractor.extract raises a ValueError for overlapping variants.

iter_variant_seqs(pairs, reference, *, use_strand, anchor="center", group_by=("interval_idx", "variant_idx"), encoding=Tokenizer(), pad_with_n=False, batch_size=64) is in the new module kipoiseq2.extractors.batches. It takes the pairs as a polars DataFrame or LazyFrame. The pairs with equal values in the group_by columns form a group, and each group gives one row: VariantSeqExtractor.extract applies the variants of the group to its interval. The default group_by gives one row per pair, e.g. for variant effect predictions, and group_by=["interval_idx"] applies all variants of an interval together. A group with several intervals raises a ValueError that suggests adding interval_idx to group_by. An interval without variants has no pairs, so it gives no row.

A batch has the columns of group_by, then chrom, start, end and strand unless group_by contains them, then ref_seq, alt_seq, ref and alt. The rows are sorted by the columns of group_by. encoding is a Tokenizer, an Encoder or None. A Tokenizer gives the tokens of Tokenizer.tokenize_series, e.g. pl.Array(pl.UInt8, length) for the default Tokenizer(). An Encoder gives the one-hot arrays of Encoder.one_hot_series in the dtype of the encoder, e.g. pl.Array(pl.Float32, (length, 4)) for Encoder(ambiguous="fractional", dtype=np.float32). encoding=None gives only the sequences. The encoded columns are Array columns if anchor is not None and all intervals of the pairs have the same width. Then batch["alt"].to_numpy() returns an array of shape (rows, length) or (rows, length, n_columns), without a copy unless the dtype is bool. Otherwise they are List columns, also in a batch whose sequences have the same length by chance. So all batches share one schema.

anchor picks the anchor_pos of extract by the rule of Interval.resize. With use_strand=True, "start" is the 5' end, which is end on "-" intervals. anchor=None lets indels change the length of alt_seq.

iter_variant_seqs collects the pairs and checks the arguments, the columns, the nulls and the groups when it is called. Each pair applies its ALT allele, so a False or null in the column variant_carrier raises a ValueError. Such pairs come from scan_vcf_genotypes(carriers_only=False), and the error suggests the default carriers_only=True. The sequences of a batch are built when the iterator reaches it. During the iteration, the errors of extract and of the encoding get a note that names the group, e.g. for overlapping variants, a REF that differs from the reference, or a letter that encoding lacks.

Within a batch, the reference of each interval is fetched once, as a window that extends beyond each end of the interval by the largest sum of REF lengths among its groups. extract and ref_seq cut their sequences from this window. The batch is built from Series without a polars query. _check_int, which Interval and Variant call, tests for a plain int first, which skips the slow isinstance check against numbers.Integral for most values. On 20,000 pairs with 1 kb windows (80% SNVs, 10% deletions, 10% insertions) and a FastaStringExtractor on a 5 Mb FASTA, iter_variant_seqs gives 57,000 pairs/s with batch_size=64 and 63,000 pairs/s with batch_size=256. A plain loop that calls extract, fetches the reference and tokenizes both sequences with Tokenizer.tokenize gives 70,000 pairs/s, but it builds no DataFrames. A loop that only builds the Interval and Variant objects and calls extract gives 106,000 pairs/s, which is about half of the time per pair of iter_variant_seqs.

Tokenizer.tokenize_series(seqs) and Encoder.one_hot_series(seqs) take a polars Series of strings, and the result keeps its name. If all sequences have the same length, they return an Array column: pl.Array(dtype, length) for tokens, or pl.Array(dtype, (length, n_columns)) for one-hot arrays. Otherwise they return a List column: pl.List(dtype) or pl.List(pl.Array(dtype, n_columns)). An empty Series gives an empty List column, because it has no length for an Array column. For letters of single ASCII characters, the methods read the offsets and the data buffer of the Arrow array of seqs. seqs.to_arrow(compat_level=pl.CompatLevel.oldest()) gives this array as one large_string array, also for a Series of several chunks. One bytes.translate call converts the data buffer, and a List column is built from the offsets and the values without a copy. Other alphabets, e.g. codons, use Tokenizer.tokenize or Encoder.one_hot for each sequence. A Series that is not of strings raises a TypeError, and nulls raise a ValueError. For the first sequence with an unknown letter, the methods raise the ValueError of Tokenizer.tokenize or Encoder.one_hot, with a note that names the row. The ranges extra lists pyarrow, because the methods import it. polars-bio installed it already.

The table compares the time to encode one batch into one Series, in ms. The two "joined" columns join the sequences of a column, encode them in one call and split the result, as an earlier version of this PR did. "old API" encodes with the functions that Tokenizer and Encoder replace, and "new API" with Tokenizer.tokenize and Encoder.one_hot. one-hot uint8 is Encoder(), and one-hot float32 is Encoder(ambiguous="fractional", dtype=np.float32). The sequences are random uppercase ACGTN. In the List rows, their lengths vary by up to 3 %. Each time is the median per call, as the mean of two rounds, and all three give the same values.

batch	column	encoding	joined, old API	joined, new API	Series methods
64 x 196,608 bases	Array	tokens	24.0	7.8	6.0
64 x 196,608 bases	Array	one-hot uint8	30.3	26.5	30.5
64 x 196,608 bases	Array	one-hot float32	43.1	37.5	39.7
64 x 196,608 bases	List	tokens	24.4	5.4	8.5
64 x 196,608 bases	List	one-hot uint8	46.2	38.0	25.8
64 x 196,608 bases	List	one-hot float32	98.1	82.3	37.2
10,000 x 1 kb	Array	tokens	18.7	4.4	4.3
10,000 x 1 kb	Array	one-hot uint8	20.7	19.2	20.7
10,000 x 1 kb	Array	one-hot float32	29.5	27.0	30.0
10,000 x 1 kb	List	tokens	43.8	28.4	4.3
10,000 x 1 kb	List	one-hot uint8	75.6	73.0	21.4
10,000 x 1 kb	List	one-hot float32	112	109	30.4

The Array one-hot rows spend most of their time in the take that writes the output array. The times of 64 x 196,608 bases swing by up to 3 ms between processes. Interleaved in one process, the joined route and the Series methods take the same time for these tokens. A batch of 64 x 1 kb takes 43 instead of 30 us, which is the fixed cost of to_arrow and of the checks. The pairs/s of iter_variant_seqs stay the same within the noise of 10 %.

scan_vcf_genotypes(path, haplotypes=True) gives one row per ALT allele, sample and haplotype. The new column haplotype is 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, and carriers_only keeps only the haplotypes that carry it. A missing GT gives one row with a null haplotype. haplotype ignores the phasing, as bcftools consensus -H does: for an unphased GT such as 0/1, the order of the alleles in GT sets the haplotype. Unphased GTs raise no error: the frame is lazy, so an error would only appear at collect time as a polars error, and a homozygous unphased GT such as 1/1 is unambiguous. GT stays in the frame, so ~pl.col("GT").str.contains("/") keeps the phased and haploid GTs. Passed as the variants of scan_variant_pairs, the genotypes give the pair columns variant_sample and variant_haplotype. So group_by=["interval_idx", "variant_sample", "variant_haplotype"] gives the sequence of each haplotype that carries variants in the interval. The default haplotypes=False keeps the rows as before.

VariantSeqExtractor.extract raises a ValueError if the REF of two variants share a base, e.g. for two ALT alleles of one VCF record. Before, it applied both: two SNVs at one base gave two bases, in the order of the variants. GenomeKit rejects overlapping variants too. Variants that only touch, such as an SNV next to a deletion, still apply together. iter_variant_seqs shares this rule, because it calls extract. The commit is a fix, because the old output was wrong. MIGRATING.md lists the change under the breaking changes of 0.2, so the footer BREAKING CHANGE describes it.

README.md gets a section "Batches of reference and alternative sequences", with a 1 kb window per variant and the haplotypes of genes as examples, and a paragraph on haplotypes=True. A variant within 500 bases of a chromosome end gets a window that runs past the chromosome, so the window example passes pad_with_n=True. The README section "Encoding sequences" shows the Series methods. In MIGRATING.md, iter_variant_seqs replaces SingleVariantVCFSeqExtractor and SingleSeqVCFSeqExtractor. For their sample_id, it groups the rows of scan_vcf_genotypes by sample or by haplotype. The API page "Extractors" lists kipoiseq2.extractors.batches.

BREAKING CHANGE: VariantSeqExtractor.extract raises a ValueError for two variants whose REF share a base, e.g. two ALT alleles of one VCF record. In 0.1, it applied both. Apply such variants in separate calls, e.g. with the default group_by of iter_variant_seqs, which gives one row per pair.

fix(extractors): reject overlapping variants in VariantSeqExtractor
feat(vcf): give one row per haplotype in scan_vcf_genotypes
feat(encoders): encode polars Series from their Arrow buffers