A fast dataloader for bigwig files made for machine learning
"""
This script benchmarks both pyBigWig and bigwig-loader
by processing the same batches of intervals over and over.
There is also not anything else going on except loading
epigenetic profiles: no sampling, no loading genetic
sequence.

In this script pyBigWig is processing data in the main
process. A benchmark script with worker processes for
pyBigWig can be found in:

02_benchmark_minimal_pybigwig_multiple_workers.py

For a benchmark script that involves the entire loader
including loading randomly sampling intervals and loading
the genetic sequence see

03_benchmark_entire_loader.py
04_benchmark_entire_loader_pybigwig.py
05_benchmark_entire_loader_pybigwig_multiple_workers.py

Using bigwig-loader:
1 0.057052651047706605 0.00031709791153868065
2 0.057504883874207734 0.00017417943769643463
4 0.05825563352555037 0.00019529730686256182
8 0.06154786795377731 0.0003019231155274113
16 0.06676486572250724 0.00022556050291397823
32 0.07972694663330912 0.0076502898637687425
64 0.09993992112576962 0.00027788778186448623
128 0.1499797949567437 0.00045689017021743373
256 0.2352875482290983 0.0023178531225548945
384 0.27881404412910343 0.0018045323118200274
512 0.3368280011229217 0.0006013540903082275
640 0.4124560124240816 0.002849675913360234
768 0.4759079207666218 0.01028077048736988
896 0.5377236870117486 0.0015763286466726138
1024 0.5991766246035695 0.0009644582857066636
1152 0.6788586993701756 0.0019231941816597235
1280 0.7360170400701463 0.0013234154183372115
1408 0.7953401518054306 0.0021116091192682047
1536 0.8512347645126284 0.0014503305936706486
1664 0.9066219345666469 0.0011453341982749716
1792 0.9620565695688128 0.0014472743645979804
1920 1.0174510131590069 0.001081436407676258
2048 1.074055849108845 0.0019284532899694303
2176 1.129074206855148 0.0026376567781361444
2304 1.1839054396376014 0.001356973514175882
2432 1.228180939052254 0.001749555462204235
2560 1.2815975906327366 0.001366976399215406
2688 1.333120484650135 0.002277772508031197
2816 1.3867404093034565 0.002437589407029183
2944 1.4343125496990978 0.002801285683084347

Using pybigwig:
1 0.007712579797953367 1.3239532677819585e-05
2 0.015228871442377567 8.263952207858359e-05
4 0.030153942853212358 0.00019700686679982924
8 0.060477635357528925 0.0002358331421623131
16 0.12558735040947794 0.00035793705119275106
32 0.25060572251677515 0.0009981795971422773
64 0.4963857516646385 0.0015567211998296923
128 0.9931910111568868 0.006337946291953
256 1.9887389127165078 0.003353472583919708
384 2.978631450049579 0.004801139675230364
512 3.9880202571861445 0.009407197750658897
640 4.9852978646755215 0.00611967283515103
768 5.983420168235898 0.009951637939808588
896 7.002280553430319 0.03586572903432073
1024 7.970334761962294 0.007562513363487385
1152 8.959011241234839 0.007723535605142787
1280 9.959012242592872 0.010164891096016896
1408 10.94503095811233 0.009721375058046907
1536 11.942723208200187 0.013671800221479928
1664 12.922414517775177 0.01037941394057453
1792 13.911210288107394 0.011497199485000838
1920 14.922799918893725 0.019355535832356417
2048 15.89705225303769 0.016396645425789023
2176 16.869348569959403 0.01744892580180062
2304 17.84538498856127 0.017183208089130094
2432 18.836720233224334 0.01547382450705791
2560 19.837676095962525 0.016545546754101328
2688 20.812969361618162 0.017382408372612606
2816 21.805524553731082 0.019967367162665957
2944 22.810655508097263 0.026718328948077626

"""

import time

import cupy as cp
import numpy as np
import pandas as pd
import pyBigWig

from bigwig_loader import config
from bigwig_loader.collection import BigWigCollection
from bigwig_loader.collection import interpret_path


class PyBigWigCollection:
    def __init__(self, bigwig_path, first_n_files=None):
        bigwig_paths = sorted(interpret_path(bigwig_path))
        self.bigwigs = [pyBigWig.open(str(bw)) for bw in bigwig_paths][:first_n_files]

    def get_batch(self, chromosomes, starts, ends):
        out = np.stack(
            [
                bw.values(chrom, start, end, numpy=True)
                for chrom, start, end in zip(chromosomes, starts, ends)
                for bw in self.bigwigs
            ]
        )

        out = out.reshape(
            (out.shape[0] // len(self.bigwigs), len(self.bigwigs), out.shape[1])
        )

        return cp.asarray(out)


def some_intervals():
    df = pd.read_csv("example_data/benchmark_positions.tsv", sep="\t")
    start = df["center"].values - 500
    return list(df["chr"]), start, start + 1000


def run_benchmark(collection, batch_sizes, chromosomes, starts, ends):
    for batch_size in batch_sizes:
        # print("batch_size", batch_size)

        chrom = chromosomes[:batch_size]
        start = starts[:batch_size]
        end = ends[:batch_size]

        # burn in
        for _ in range(20):
            collection.get_batch(chrom, start, end)

        elapsed = []
        start_time = time.perf_counter()
        for _ in range(20):
            collection.get_batch(chrom, start, end)
            end_time = time.perf_counter()
            elapsed.append(end_time - start_time)
            start_time = end_time
        # print("Seconds per batch:", np.mean(elapsed))
        # print("standard deviation:", np.std(elapsed))
        print(batch_size, np.mean(elapsed), np.std(elapsed))


def run_all_benchmarks(bigwig_path=config.bigwig_dir):
    print("Loading from:", config.bigwig_dir)
    pybigwig_collection = PyBigWigCollection(bigwig_path, first_n_files=None)
    bigwig_loader_collection = BigWigCollection(bigwig_path, first_n_files=None)
    chromosomes, starts, ends = some_intervals()
    batch_sizes = [1, 2, 4, 8, 16, 32, 64] + [128 * i for i in range(1, 48)]

    print("Using bigwig-loader:")
    run_benchmark(bigwig_loader_collection, batch_sizes, chromosomes, starts, ends)

    print("Using pybigwig:")
    run_benchmark(pybigwig_collection, batch_sizes, chromosomes, starts, ends)


if __name__ == "__main__":
    run_all_benchmarks()

benchmark/01_benchmark_minimal.py
pfizer-opensource/bigwig-loader
"""
This benchmarks the scenario of using pyBigWig
in a pytorch dataloader with differen number
of workers.

0 32 0.24952834727099307 0.009142287326556482
1 32 0.2645324201394732 0.15919867478319427
2 32 0.1419987794584953 0.1940197516454249
4 32 0.08257061717315362 0.21923183762619575
8 32 0.052774008721686326 0.2243132436523036
16 32 0.038562032580375674 0.25572177004054747
"""

from torch.multiprocessing import set_start_method

try:
    set_start_method("spawn", force=True)
    print("spawned")
except RuntimeError:
    pass

import os
import time
from pathlib import Path
from typing import Iterable

import numpy as np
import pandas as pd
import pyBigWig
from torch.utils.data import DataLoader
from torch.utils.data import IterableDataset

from bigwig_loader import config


def get_bigwig_files_from_path(
    path: Path, file_extensions: Iterable[str] = (".bigWig", ".bw"), crawl: bool = True
) -> list[Path]:
    """
    For a path object, get all bigwig files. If the path is directly pointing
    to a bigwig file, a list with just that file is returned. In case of a
    directory all files are gathered with file extensions that are part of
    file_extensions.
    Args:
        path: or upath.Path object. Either directory or BigWig file
        file_extensions: used to find all BigWig files in a directory.
            Default: (".bigWig", ".bw")
        crawl: whether to find BigWig files in subdirectories. Default: True.

    Returns: list of paths to BigWig files.

    """
    if not path.exists():
        raise FileNotFoundError(f"No such file or directory: {path}")
    elif path.is_dir():
        if crawl:
            pattern = "**/*"
        else:
            pattern = "*"
        return [
            file
            for extension in file_extensions
            for file in path.glob(f"{pattern}{extension}")
        ]
    return [path]


def interpret_path(
    bigwig_path,
    file_extensions: Iterable[str] = (".bigWig", ".bw"),
    crawl: bool = True,
) -> list[Path]:
    """
    Get all bigwig files for a path. Also excepts strings.
    If the path is directly pointing to a bigwig file, a list with just that
    file is returned. In case of a directory all files are gathered with file
    extensions that are part of file_extensions.
    Args:
        bigwig_path: str, Path or upath.Path object. Either directory or BigWig file
        file_extensions: used to find all BigWig files in a directory.
            Default: (".bigWig", ".bw")
        crawl: whether to find BigWig files in subdirectories. Default: True.

    Returns: list of paths to BigWig files.

    """
    if isinstance(bigwig_path, str) or isinstance(bigwig_path, os.PathLike):
        return get_bigwig_files_from_path(
            Path(bigwig_path), file_extensions=file_extensions, crawl=crawl
        )

    elif isinstance(bigwig_path, Iterable):
        return [
            path
            for element in bigwig_path
            for path in interpret_path(
                element, file_extensions=file_extensions, crawl=crawl
            )
        ]
    raise ValueError(
        f"Can not interpret {bigwig_path} as path or a collection of paths."
    )


class PyBigWigCollection:
    def __init__(self, bigwig_path):
        self.bigwig_paths = sorted(interpret_path(bigwig_path))
        self.bigwigs = None

    def init_bigwig_filehandles(self):
        self.bigwigs = [pyBigWig.open(str(bw)) for bw in self.bigwig_paths]

    def get_batch(self, chromosomes, starts, ends):
        if self.bigwigs is None:
            self.init_bigwig_filehandles()
        out = np.stack(
            [
                bw.values(chrom, start, end, numpy=True)
                for chrom, start, end in zip(chromosomes, starts, ends)
                for bw in self.bigwigs
            ]
        )

        out = out.reshape(
            (out.shape[0] // len(self.bigwigs), len(self.bigwigs), out.shape[1])
        )

        return out


class Dataset(IterableDataset):

    """Just a wrapper."""

    def __init__(self, collection, chromosomes, starts, ends):
        self.collection = collection
        self.chromosomes = chromosomes
        self.starts = starts
        self.ends = ends

    def init_collection(self, _):
        self.collection.init_bigwig_filehandles()

    def __len__(self):
        return 128

    def __iter__(self):
        return self

    def __next__(self):
        return self.collection.get_batch(self.chromosomes, self.starts, self.ends)


def some_intervals():
    df = pd.read_csv("example_data/benchmark_positions.tsv", sep="\t")
    start = df["center"].values - 500
    return list(df["chr"]), start, start + 1000


def run_benchmark(batch_sizes, chromosomes, starts, ends, workers):
    for n_workers in workers:
        for batch_size in batch_sizes:
            # print("batch_size", batch_size)

            chrom = chromosomes[:batch_size]
            start = starts[:batch_size]
            end = ends[:batch_size]

            collection = PyBigWigCollection(config.bigwig_dir)

            dataset = Dataset(collection, chrom, start, end)

            # batch_size is set to None because the Dataset returns
            # whole batches at once.
            dataloader = DataLoader(
                dataset,
                num_workers=n_workers,
                batch_size=None,
                worker_init_fn=dataset.init_collection,
            )

            elapsed = []
            start_time = time.perf_counter()
            for i, _ in enumerate(dataloader):
                end_time = time.perf_counter()
                elapsed.append(end_time - start_time)
                start_time = end_time
                if i > 512:
                    break

            print(n_workers, batch_size, np.mean(elapsed), np.std(elapsed))


def run_all_benchmarks(bigwig_path=config.bigwig_dir):
    chromosomes, starts, ends = some_intervals()
    batch_sizes = [32]

    print("Using pybigwig and PyTorch dataloader with multiple number of workers:")
    run_benchmark(
        batch_sizes, chromosomes, starts, ends, workers=[0, 1, 2, 4, 8, 16, 32]
    )


if __name__ == "__main__":
    run_all_benchmarks()

benchmark/02_benchmark_minimal_pybigwig_multiple_workers.py
pfizer-opensource/bigwig-loader
"""
This runs the entire loader producing the input and target
values a neural network can be trained on. In this file
pyBigWig is used as the library processing the bigwig
data.

A more plain benchmark can be found in benchmark_minimal.py
In that file, the same set of intervals is queried over and
over. This eliminates some of the influence of IO on the
numbers and also does not include things that have nothing
to do with bigwig, like the sampling and looking up the
sequence.
"""

import time

import numpy as np
import pandas as pd

from bigwig_loader import config
from bigwig_loader.dataset import BigWigDataset


def run(batch_size):
    train_regions = pd.read_csv("train_regions.tsv", sep="\t")
    print("Loading from:", config.bigwig_dir)

    dataset = BigWigDataset(
        regions_of_interest=train_regions,
        collection=config.bigwig_dir,
        reference_genome_path=config.reference_genome,
        sequence_length=1000,
        center_bin_to_predict=1000,
        batch_size=batch_size,
        batches_per_epoch=10,
        maximum_unknown_bases_fraction=0.1,
        sequence_encoder="onehot",
    )

    elapsed = []
    for i, (sequence, target) in enumerate(dataset):
        if i == 0:
            start_time = time.perf_counter()
            continue
        end_time = time.perf_counter()
        elapsed.append(end_time - start_time)
        start_time = end_time
    print(batch_size, np.mean(elapsed), np.std(elapsed))


if __name__ == "__main__":
    run(batch_size=2048)

benchmark/03_benchmark_entire_loader.py
pfizer-opensource/bigwig-loader
"""
This runs the entire loader producing the input and target
values a neural network can be trained on. In this file
pyBigWig is used as the library processing the bigwig
data.

A more plain benchmark can be found in benchmark_minimal.py
In that file, the same set of intervals is queried over and
over. This eliminates some of the influence of IO on the
numbers and also does not include things that have nothing
to do with bigwig, like the sampling and looking up the
sequence.
"""

import time
from pathlib import Path
from typing import Callable
from typing import Optional
from typing import Union

import cupy as cp
import numpy as np
import pandas as pd
import pyBigWig

from bigwig_loader import config
from bigwig_loader.genome import Genome
from bigwig_loader.position_sampler import PositionSampler


class PyBigWigDataset:
    """
    Dataset over FASTA files and BigWig profiles.
    Args:
        regions_of_interest: pandas dataframe with intervals within the chromosomes to
            sample from. These intervals are not intervals of sequence_length
            but regions of the chromosome that were selected up front. For
            instance regions in which called peaks are located or a threshold
            value is exceeded.
        bigwig_path: path to BigWig Directory or list of BigWig files.
        reference_genome_path: path to fasta file containing the reference genome.
        sequence_length: number of base pairs in input sequence
        center_bin_to_predict: if given, only do prediction on a central window. Should be
            smaller than or equal to sequence_length. If not given will be the same as
            sequence_length.
        window_size: used to down sample the resolution of the target from sequence_length
        batch_size: batch size
        batches_per_epoch: because the length of an epoch is slightly arbitrary here,
            the number of batches can be set by hand. If not the number of batches per
            epoch will be (totol number of bases in combined intervals) // sequence_length // batch_size
        maximum_unknown_bases_fraction: maximum number of bases in an input sequence that
            is unknown.
        sequence_encoder: encoder to apply to the sequence. Default: bigwig_loader.util.onehot_sequences
    """

    def __init__(
        self,
        regions_of_interest: pd.DataFrame,
        collection: Path,
        reference_genome_path: Path,
        sequence_length: int = 1000,
        center_bin_to_predict: Optional[int] = 200,
        batch_size: int = 256,
        batches_per_epoch: Optional[int] = None,
        maximum_unknown_bases_fraction: float = 0.1,
        sequence_encoder: Optional[Union[Callable, str]] = "onehot",
    ):
        super().__init__()
        self.regions_of_interest = regions_of_interest
        self.bigwig_path = collection
        self.reference_genome_path = reference_genome_path
        self.sequence_length = sequence_length
        self.center_bin_to_predict = center_bin_to_predict
        self.batch_size = batch_size
        self.batches_per_epoch = (
            batches_per_epoch
            or regions_of_interest.index[-1] // sequence_length // batch_size
        )
        self._n = 0
        self.maximum_unknown_bases_fraction = maximum_unknown_bases_fraction
        self.sequence_encoder = sequence_encoder

    def _setup_genome(self):
        position_sampler = PositionSampler(regions_of_interest=self.regions_of_interest)
        self.genome = Genome(
            self.reference_genome_path,
            sequence_length=self.sequence_length,
            position_sampler=position_sampler,
            batch_size=self.batch_size,
            maximum_unknown_bases_fraction=self.maximum_unknown_bases_fraction,
            encoder=self.sequence_encoder,
        )

    def _setup_bigwig_collection(self):
        # self.bigwig_collection = BigWigCollection(self.bigwig_path)
        bigwig_paths = list(self.bigwig_path.glob("*.bigWig")) + list(
            self.bigwig_path.glob("*.bw")
        )
        self.bigwigs = [pyBigWig.open(str(bw)) for bw in bigwig_paths]

    def __iter__(self):
        self._setup_genome()
        self._setup_bigwig_collection()
        self._n = 0
        return self

    def __next__(self):
        if self._n < self.batches_per_epoch:
            self._n += 1
            positions, sequences = self.genome.get_batch()
            chromosomes, center = zip(*positions)
            start = np.array(center) - (self.center_bin_to_predict // 2)
            end = start + self.center_bin_to_predict
            target = self.get_target_batch(chromosomes, start, end)
            return sequences, target
        raise StopIteration

    def get_target_batch(self, chromosomes, starts, ends):
        return cp.asarray(
            np.stack(
                [
                    bw.values(chrom, start, end, numpy=True)
                    for bw in self.bigwigs
                    for chrom, start, end in zip(chromosomes, starts, ends)
                ]
            )
        )


def run(batch_size):
    train_regions = pd.read_csv("train_regions.tsv", sep="\t")
    print("Loading from:", config.bigwig_dir)

    dataset = PyBigWigDataset(
        regions_of_interest=train_regions,
        collection=config.bigwig_dir,
        reference_genome_path=config.reference_genome,
        sequence_length=1000,
        center_bin_to_predict=1000,
        batch_size=batch_size,
        batches_per_epoch=10,
        maximum_unknown_bases_fraction=0.1,
        sequence_encoder="onehot",
    )

    elapsed = []
    for i, (sequence, target) in enumerate(dataset):
        if i == 0:
            start_time = time.perf_counter()
            continue
        end_time = time.perf_counter()
        elapsed.append(end_time - start_time)
        start_time = end_time
    print(batch_size, np.mean(elapsed), np.std(elapsed))


if __name__ == "__main__":
    run(batch_size=2048)

benchmark/04_benchmark_entire_loader_pybigwig.py
pfizer-opensource/bigwig-loader
import time
from pathlib import Path
from typing import Callable
from typing import Optional
from typing import Union

import numpy as np
import pandas as pd
import pyBigWig
import torch
from torch.utils.data import DataLoader
from torch.utils.data import IterableDataset

from bigwig_loader import config
from bigwig_loader.genome import Genome
from bigwig_loader.position_sampler import PositionSampler


class PyBigWigDataset:
    """
    Dataset over FASTA files and BigWig profiles.
    Args:
        regions_of_interest: pandas dataframe with intervals within the chromosomes to
            sample from. These intervals are not intervals of sequence_length
            but regions of the chromosome that were selected up front. For
            instance regions in which called peaks are located or a threshold
            value is exceeded.
        bigwig_path: path to BigWig Directory or list of BigWig files.
        reference_genome_path: path to fasta file containing the reference genome.
        sequence_length: number of base pairs in input sequence
        center_bin_to_predict: if given, only do prediction on a central window. Should be
            smaller than or equal to sequence_length. If not given will be the same as
            sequence_length.
        window_size: used to down sample the resolution of the target from sequence_length
        batch_size: batch size
        batches_per_epoch: because the length of an epoch is slightly arbitrary here,
            the number of batches can be set by hand. If not the number of batches per
            epoch will be (totol number of bases in combined intervals) // sequence_length // batch_size
        maximum_unknown_bases_fraction: maximum number of bases in an input sequence that
            is unknown.
        sequence_encoder: encoder to apply to the sequence. Default: bigwig_loader.util.onehot_sequences
    """

    def __init__(
        self,
        regions_of_interest: pd.DataFrame,
        bigwig_path: Path,
        reference_genome_path: Path,
        sequence_length: int = 1000,
        center_bin_to_predict: Optional[int] = 200,
        batch_size: int = 256,
        batches_per_epoch: Optional[int] = None,
        maximum_unknown_bases_fraction: float = 0.1,
        sequence_encoder: Optional[Union[Callable, str]] = "onehot",
    ):
        super().__init__()
        self.regions_of_interest = regions_of_interest
        self.bigwig_path = bigwig_path
        self.reference_genome_path = reference_genome_path
        self.sequence_length = sequence_length
        self.center_bin_to_predict = center_bin_to_predict
        self.batch_size = batch_size
        self.batches_per_epoch = (
            batches_per_epoch
            or regions_of_interest.index[-1] // sequence_length // batch_size
        )
        self._n = 0
        self.maximum_unknown_bases_fraction = maximum_unknown_bases_fraction
        self.sequence_encoder = sequence_encoder

    def _setup_genome(self):
        position_sampler = PositionSampler(regions_of_interest=self.regions_of_interest)
        self.genome = Genome(
            self.reference_genome_path,
            sequence_length=self.sequence_length,
            position_sampler=position_sampler,
            batch_size=self.batch_size,
            maximum_unknown_bases_fraction=self.maximum_unknown_bases_fraction,
            encoder=self.sequence_encoder,
        )

    def _setup_bigwig_collection(self):
        # self.bigwig_collection = BigWigCollection(self.bigwig_path)
        bigwig_paths = list(self.bigwig_path.glob("*.bigWig")) + list(
            self.bigwig_path.glob("*.bw")
        )
        self.bigwigs = [pyBigWig.open(str(bw)) for bw in bigwig_paths]

    def __iter__(self):
        self._setup_genome()
        self._setup_bigwig_collection()
        self._n = 0
        return self

    def __next__(self):
        if self._n < self.batches_per_epoch:
            self._n += 1
            positions, sequences = self.genome.get_batch()
            chromosomes, center = zip(*positions)
            start = np.array(center) - (self.center_bin_to_predict // 2)
            end = start + self.center_bin_to_predict
            target = self.get_target_batch(chromosomes, start, end)
            return sequences, target
        raise StopIteration

    def get_target_batch(self, chromosomes, starts, ends):
        return np.stack(
            [
                bw.values(chrom, start, end, numpy=True)
                for bw in self.bigwigs
                for chrom, start, end in zip(chromosomes, starts, ends)
            ]
        )


class PytorchPyBigWigDataset(IterableDataset):

    """
    Pytorch IterableDataset over FASTA files and BigWig profiles.
    Args:
        regions_of_interest: pandas dataframe with intervals within the chromosomes to
            sample from. These intervals are not intervals of sequence_length
            but regions of the chromosome that were selected up front. For
            instance regions in which called peaks are located or a threshold
            value is exceeded.
        bigwig_path: path to BigWig Directory or list of BigWig files.
        reference_genome_path: path to fasta file containing the reference genome.
        sequence_length: number of base pairs in input sequence
        center_bin_to_predict: if given, only do prediction on a central window. Should be
            smaller than or equal to sequence_length. If not given will be the same as
            sequence_length.
        window_size: used to down sample the resolution of the target from sequence_length
        batch_size: batch size
        batches_per_epoch: because the length of an epoch is slightly arbitrary here,
            the number of batches can be set by hand. If not the number of batches per
            epoch will be (totol number of bases in combined intervals) // sequence_length // batch_size
        maximum_unknown_bases_fraction: maximum number of bases in an input sequence that
            is unknown.
        sequence_encoder: encoder to apply to the sequence. Default: bigwig_loader.util.onehot_sequences
    """

    def __init__(
        self,
        regions_of_interest: pd.DataFrame,
        bigwig_path: Path,
        reference_genome_path: Path,
        sequence_length: int = 1000,
        center_bin_to_predict: Optional[int] = 200,
        window_size: int = 1,
        batch_size: int = 256,
        batches_per_epoch: Optional[int] = None,
        maximum_unknown_bases_fraction: float = 0.1,
        sequence_encoder: Optional[Union[Callable, str]] = "onehot",
    ):
        super().__init__()
        self._dataset = PyBigWigDataset(
            regions_of_interest=regions_of_interest,
            bigwig_path=bigwig_path,
            reference_genome_path=reference_genome_path,
            sequence_length=sequence_length,
            center_bin_to_predict=center_bin_to_predict,
            batch_size=batch_size,
            batches_per_epoch=batches_per_epoch,
            maximum_unknown_bases_fraction=maximum_unknown_bases_fraction,
            sequence_encoder=sequence_encoder,
        )
        if window_size and window_size > 1:
            self.average_pool_target = torch.nn.AvgPool1d(
                kernel_size=window_size, stride=window_size
            )
        else:
            self.average_pool_target = None

    def __iter__(self):
        iter(self._dataset)
        return self

    def __next__(self):
        sequences, target = next(self._dataset)
        target = torch.from_numpy(target).float()
        sequences = torch.FloatTensor(sequences)
        if self.average_pool_target:
            target = self.average_pool_target(target)
        return sequences, target


def run():
    train_regions = pd.read_csv("train_regions.tsv", sep="\t")
    print("Loading from:", config.bigwig_dir)

    dataset = PytorchPyBigWigDataset(
        regions_of_interest=train_regions,
        bigwig_path=config.bigwig_dir,
        reference_genome_path=config.reference_genome,
        sequence_length=1000,
        center_bin_to_predict=1000,
        window_size=1,
        batch_size=256,
        batches_per_epoch=20,
        maximum_unknown_bases_fraction=0.1,
        sequence_encoder="onehot",
    )

    loader = DataLoader(dataset, batch_size=None, num_workers=8)

    elapsed = []
    for i, (X, Y) in enumerate(loader):
        if i == 0:
            start = time.perf_counter()
            continue
        end = time.perf_counter()
        elapsed.append(end - start)
        start = end
        # if not i % 100:
        print("Seconds per batch:", sum(elapsed) / i)
        print(Y.shape)


if __name__ == "__main__":
    run()

benchmark/05_benchmark_entire_loader_pybigwig_multiple_workers.py
pfizer-opensource/bigwig-loader
__author__ = ["Joren Retel"]
__email__ = ["joren.retel@pfizer.com"]

try:
    from bigwig_loader._version import version as __version__
except ImportError:
    # this protects against an edge case where a user tries to import
    # the module without installing it, by adding it manually to
    # sys.path or trying to run it directly from the git checkout.
    __version__ = "not-installed"

from bigwig_loader.settings import Settings as _Settings

config = _Settings()

bigwig_loader/__init__.py
pfizer-opensource/bigwig-loader
import zlib
from pathlib import Path
from random import sample
from typing import BinaryIO
from typing import Iterable
from typing import Literal
from typing import Optional
from typing import Sequence
from typing import Union

import cupy as cp
import numpy as np
import numpy.typing as npt
import pandas as pd
from ncls import NCLS

from bigwig_loader.gpu_decompressor import Decoder
from bigwig_loader.memory_bank import MemoryBank
from bigwig_loader.merge_intervals import merge_intervals
from bigwig_loader.parser import BBIHeader
from bigwig_loader.parser import ChromosomeTreeHeader
from bigwig_loader.parser import ChromosomeTreeNode
from bigwig_loader.parser import RTreeIndexHeader
from bigwig_loader.parser import RTreeNode
from bigwig_loader.parser import TotalSummary
from bigwig_loader.parser import WIGSectionHeader
from bigwig_loader.parser import ZoomHeader
from bigwig_loader.store import BigWigStore
from bigwig_loader.util import get_standard_chromosomes


class BigWig:
    def __init__(self, path: Path, id: Optional[int] = None):
        """
        Create a BigWig object representing one BigWig file.
        Args:
            path: path to BigWig file
            id: integer representing the file in a collection
                when part of a collection.
        """

        with open(path, "rb") as bigwig:
            self.path = path
            self.id = id
            self.bbi_header = BBIHeader.from_file(bigwig)
            self.zoom_headers = ZoomHeader.all(
                bigwig, first_offset=64, n_zoom_levels=self.bbi_header.zoom_levels
            )
            self.total_summary = TotalSummary.from_file_and_offset(
                bigwig, self.bbi_header.total_summary_offset
            )
            self.chromosome_tree_header = ChromosomeTreeHeader.from_file_and_offset(
                bigwig, self.bbi_header.chromosome_tree_offset
            )

            self.chromosome_head_node = ChromosomeTreeNode.from_file_and_offset(
                bigwig,
                key_size=self.chromosome_tree_header.key_size,
                offset=self.bbi_header.chromosome_tree_offset + 32,
            )
            self.chromosome_sizes = {
                chrom.key: chrom.chrom_size for chrom in self.chromosome_head_node.items
            }
            self.rtree_index_header = RTreeIndexHeader.from_file_and_offset(
                bigwig, self.bbi_header.full_index_offset
            )
            self.rtree_head_node = RTreeNode.from_file_and_offset(
                file_object=bigwig,
                start_chrom_ix=self.rtree_index_header.start_chrom_ix,
                start_base=self.rtree_index_header.start_base,
                end_chrom_ix=self.rtree_index_header.end_chrom_ix,
                end_base=self.rtree_index_header.end_base,
                offset=None,
            )

            self.rtree_leaf_nodes = self.rtree_head_node.get_leaf_nodes()
            self.max_rows_per_chunk = self._guess_max_rows_per_chunk(bigwig)

        self.chrom_to_chrom_id: dict[str, int] = {
            item.key: item.chrom_id for item in self.chromosome_head_node.items  # type: ignore
        }
        self._chrom_id_to_chrom = self._create_chrom_id_to_chrom_key()

        self.chromosome_offsets: npt.NDArray[np.int64] = None  # type: ignore
        self.store: BigWigStore = None  # type: ignore
        self.ncls_index: NCLS = None
        self.reference_df: pd.DataFrame = None  # type: ignore

    def run_indexing(self, chromosome_offsets: npt.NDArray[np.int64]) -> None:
        """Run NCLS indexing of BigWig file. The bigwig file has
        an index itself as well. But we prefer to recalculate
        an index here.
        """

        self.chromosome_offsets = chromosome_offsets
        self.ncls_index, self.reference_df = self.build_ncls_index()
        self.store = BigWigStore(
            self.path,
            chunk_sizes=self.reference_df["data_size"].values,  # type: ignore
            chunk_offsets=self.reference_df["data_offset"].values,  # type: ignore
        )

    def chromosomes(
        self,
        include_chromosomes: Union[Literal["all", "standard"], Sequence[str]] = "all",
        exclude_chromosomes: Optional[Sequence[str]] = None,
    ) -> list[str]:
        """
        Get chromosome keys ("chr1", "chr2"...) that are present in
        the BigWig file. Subsets of all chromosome keys found in this file
        can be made by supplying include_chromosomes and exclude_chromosomes.
        Args:
            include_chromosomes: chromosome keys you want to include.
                Can alternatively be "all" or "standard".
            exclude_chromosomes: Optional set of chromosomes to exclude.

        Returns: list of chromosome keys

        """
        exclude_chromosomes = exclude_chromosomes or []
        chromosomes_in_file = list(self.chrom_to_chrom_id.keys())
        if include_chromosomes == "all":
            include_chromosomes = chromosomes_in_file
        elif include_chromosomes == "standard":
            include_chromosomes = [
                chrom
                for chrom in get_standard_chromosomes(exclude=exclude_chromosomes)
                if chrom in chromosomes_in_file
            ]
        elif isinstance(include_chromosomes, str):
            include_chromosomes = [include_chromosomes]
        return [
            chrom for chrom in include_chromosomes if chrom not in exclude_chromosomes
        ]

    def intervals(
        self,
        include_chromosomes: Union[Literal["all", "standard"], Sequence[str]] = "all",
        exclude_chromosomes: Optional[Sequence[str]] = None,
        threshold: Optional[float] = None,
        merge: bool = False,
        merge_allow_gap: int = 0,
        memory_bank: Optional[MemoryBank] = None,
        decoder: Optional[Decoder] = None,
        batch_size: int = 4096,
    ) -> pd.DataFrame:
        """
        Does not give back data in the chromosome order you asked for necessarily
        Args:
            include_chromosomes: list of chromosome, "standard" or "all" (default).
            exclude_chromosomes: list of chromosomes you want to exclude
            threshold: only return intervals of which the value exceeds
                this threshold.
            merge: whether to merge intervals that are directly following
                eachother. The value will be the max value of the merged
                intervals.
            merge_allow_gap: default 0. Allow intervals seperated by size
                merge_allow_gap bases to still be merged.
        Returns: pandas dataframe of intervals (chrom, start, end, value)

        """
        if memory_bank is None:
            memory_bank = MemoryBank(elastic=True)
        if decoder is None:
            decoder = Decoder(
                max_num_chunks=10000 * 100,  # 512 * len(self.bigwigs) * 10,
                max_rows_per_chunk=self.max_rows_per_chunk,
                max_uncompressed_chunk_size=self.max_rows_per_chunk * 12 + 24,
                chromosome_offsets=None,
            )

        chromosome_keys = self.chromosomes(
            include_chromosomes=include_chromosomes,
            exclude_chromosomes=exclude_chromosomes,
        )

        # Doing the sort here so that start_chrom_ids are sorted. This prevents
        # a sort down the line when merging intervals. Because of the different
        # order and mapping between chrom_ids and chrom_keys in different bigwig
        # files, it does not work to sort the resulting chunk_ids, as their order
        # is determined by the consensus order that was created over all bigwig
        # files.
        selected_df = self.reference_df[
            self.reference_df["start_chrom_key"].isin(chromosome_keys)
        ].sort_values(by=["start_chrom_ix", "start_base"])

        offsets = selected_df["data_offset"].values
        sizes = selected_df["data_size"].values

        chrom_ids = []
        starts = []
        ends = []
        values = []
        for i in range(0, len(offsets), batch_size):
            memory_bank.reset()
            memory_bank.add_many(
                self.store._fh,
                offsets[i : i + batch_size],
                sizes[i : i + batch_size],
                skip_bytes=2,
            )
            gpu_byte_array, comp_chunks, compressed_chunk_sizes = memory_bank.to_gpu()
            chrom_id, start, end, value, n_rows_for_chunks = decoder.decode(
                gpu_byte_array,
                comp_chunks,
                compressed_chunk_sizes,
            )
            chrom_id = cp.repeat(chrom_id, n_rows_for_chunks.get().tolist())
            if threshold:
                mask = value > threshold
                chrom_id = chrom_id[mask]
                start = start[mask]
                end = end[mask]
                value = value[mask]
            # bringing everything back to CPU
            chrom_ids.append(chrom_id.get())
            starts.append(start.get())
            ends.append(end.get())
            values.append(value.get())
        chrom_ids = np.concatenate(chrom_ids)
        starts = np.concatenate(starts)
        ends = np.concatenate(ends)
        values = np.concatenate(values)

        if merge:
            chrom_ids, starts, ends, values = merge_intervals(
                chrom_ids,
                starts,
                ends,
                values,
                is_sorted=True,
                allow_gap=merge_allow_gap,
            )
        chrom_key = self._chrom_id_to_chrom[chrom_ids]
        data = pd.DataFrame(
            {"chrom": chrom_key, "start": starts, "end": ends, "value": values}
        )
        return data

    def get_batch_offsets_and_sizes_with_global_positions(
        self,
        global_starts: Union[Iterable[int], npt.NDArray[np.int64]],
        global_ends: Union[Iterable[int], npt.NDArray[np.int64]],
    ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
        """
        Get the offsets and sizes in bytes for the compressed chunks that need
        to be loaded to be able to extract the data belonging to the intervals.

        Args:
            global_starts: start positions that are running over all chromosomes
            global_ends: end positions that are running over all chromosomes

        Returns:
            Tuple of np.ndarray offsets and np.ndarray of sizes in Bytes

        """
        _, right_index = self.search_ncls_with_global_positions(
            global_starts, global_ends
        )
        return self.store.get_offsets_and_sizes(np.sort(np.unique(right_index)))

    def build_ncls_index(self) -> tuple[NCLS, pd.DataFrame]:
        leaf_nodes = self.rtree_leaf_nodes

        df = pd.DataFrame(
            {
                "start_chrom_ix": [
                    leaf_node.start_chrom_ix for leaf_node in leaf_nodes
                ],
                "start_base": [leaf_node.start_base for leaf_node in leaf_nodes],
                "end_chrom_ix": [leaf_node.end_chrom_ix for leaf_node in leaf_nodes],
                "end_base": [leaf_node.end_base for leaf_node in leaf_nodes],
                "data_offset": [leaf_node.data_offset for leaf_node in leaf_nodes],
                "data_size": [leaf_node.data_size for leaf_node in leaf_nodes],
            }
        )
        df["start_chrom_key"] = self._chrom_id_to_chrom[df["start_chrom_ix"].values]  # type: ignore
        df["start_abs"] = self.make_positions_global(
            df["start_chrom_ix"].values, df["start_base"].values  # type: ignore
        )
        df["end_abs"] = self.make_positions_global(
            df["end_chrom_ix"].values, df["end_base"].values  # type: ignore
        )
        df = df.sort_values(by="start_abs").reset_index(drop=True)
        ncls = NCLS(df["start_abs"].values, df["end_abs"].values, df.index.values)
        return ncls, df

    def search_ncls(
        self, query_df: pd.DataFrame, use_key: bool = False
    ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
        """
        Use NCLS to search for overlapping intervals. The right key gives the
        chunk id that needs to be loaded from disk to find the interval information
        needed.

        Args:
            query_df: pandas Dataframe with columns "chrom", "start", "end"
            use_key: when use_key is True, the "chrom" column is assumed to
                have chrom keys (i.e. "chr1", "chr2"...). If False, the chrom
                column is assumed to have integer chrom_ids.

        Returns:
            tuple of left indexes and right indexes of overlapping intervals

        """

        start: npt.NDArray[np.int64] = query_df["start"].values  # type: ignore
        end: npt.NDArray[np.int64] = query_df["end"].values  # type: ignore

        if use_key:
            chrom_keys: Iterable[str] = query_df["chrom"].values
            start = np.array(
                self.make_positions_global_with_chromosome_keys(chrom_keys, start)
            )
            end = np.array(
                self.make_positions_global_with_chromosome_keys(chrom_keys, end)
            )
        else:
            chrom_ids: npt.NDArray[np.int64] = query_df["chrom"].values  # type: ignore
            start = np.array(self.make_positions_global(chrom_ids, start))
            end = np.array(self.make_positions_global(chrom_ids, end))
        return self.ncls_index.all_overlaps_both(start, end, query_df.index.values)  # type: ignore

    def search_ncls_with_global_positions(
        self,
        global_starts: Union[Iterable[int], npt.NDArray[np.int64]],
        global_ends: Union[Iterable[int], npt.NDArray[np.int64]],
    ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
        """
        Use NCLS to search for overlapping intervals. The right key gives the
        chunk id that needs to be loaded from disk to find the interval information
        needed.

        Args:
            global_starts: start positions that are running over all chromosomes
            global_ends: end positions that are running over all chromosomes
        Returns:
            tuple of left indexes and right indexes of overlapping intervals

        """

        index = np.arange(len(global_starts))  # type: ignore
        return self.ncls_index.all_overlaps_both(global_starts, global_ends, index)  # type: ignore

    def make_positions_global_with_chromosome_keys(
        self,
        chromosome_keys: Iterable[str],
        positions: npt.NDArray[np.int64],
    ) -> npt.NDArray[np.int64]:
        chromosomes = [self.chrom_to_chrom_id[key] for key in chromosome_keys]
        return self.make_positions_global(chromosomes, positions)

    def make_positions_global(
        self,
        chromosomes: Union[Sequence[int], npt.NDArray[np.int64]],
        positions: npt.NDArray[np.int64],
    ) -> npt.NDArray[np.int64]:
        return positions + self.chromosome_offsets[chromosomes]

    def _create_chrom_id_to_chrom_key(self) -> npt.NDArray[np.generic]:
        """Create a mapping from chrom_ids (5, 1, 3) to chrom
        keys ("chr1", "chr2"...) in the form of a numpy
        array, for quick conversion.

        Returns:
            numpy array with chrom keys
        """

        chromosomes = [
            (chromosome.chrom_id, chromosome.key)
            for chromosome in self.chromosome_head_node.items
        ]
        largest_id = max(chromosomes)[0]
        mapping = np.empty(largest_id + 1, dtype=object)  # type: ignore
        for chrom_id, key in chromosomes:
            mapping[chrom_id] = key
        return mapping  # type: ignore

    def _guess_max_rows_per_chunk(
        self, file_object: BinaryIO, sample_size: int = 300
    ) -> int:
        """
        Randomly samples some chunks and looks in the header
        to see if in these chunks, the number of rows in the
        decompressed data is always the same. This helps with
        the ReferenceFileSystem.

        Args:
            file_object: BigWig file opened as bytes.
        Returns:
            Number of rows in a chunk
        """

        rows_for_chunks = []
        sample_leaf_nodes = sample(self.rtree_leaf_nodes, sample_size)
        for leaf_node in sample_leaf_nodes:
            file_object.seek(leaf_node.data_offset, 0)  # type: ignore
            decoded = zlib.decompress(file_object.read(leaf_node.data_size))  # type: ignore
            header = WIGSectionHeader.from_bytes(decoded[:24])
            rows_for_chunks.append(header.item_count)
        return max(rows_for_chunks)

bigwig_loader/bigwig.py
pfizer-opensource/bigwig-loader
import logging
import os
from pathlib import Path
from typing import Any
from typing import Iterable
from typing import Literal
from typing import Optional
from typing import Sequence
from typing import Union

import cupy as cp
import numpy as np
import numpy.typing as npt
import pandas as pd
from upath import UPath

from bigwig_loader.bigwig import BigWig
from bigwig_loader.gpu_decompressor import Decoder
from bigwig_loader.intervals_to_values_gpu import intervals_to_values
from bigwig_loader.memory_bank import MemoryBank
from bigwig_loader.merge_intervals import merge_interval_dataframe
from bigwig_loader.util import chromosome_sort


class BigWigCollection:
    """
    Wrapper around pyBigWig wrapping a collection of BigWig files.
    args:
        bigwig_path: path to BigWig Directory or list of BigWig files.
        file_extensions: only used when looking in directories for BigWig files.
            All files with these extensions are assumed to be BigWIg files:
            default: (".bigWig", ".bw").
        walk: to walk the directory tree or not. If False, subdirectories are
            ignored.
        first_n_files: Optional, only consider the first n files (after sorting).
            Handy for debugging.
    """

    def __init__(
        self,
        bigwig_path: Union[str, Sequence[str], Path, Sequence[Path]],
        file_extensions: Sequence[str] = (".bigWig", ".bw"),
        crawl: bool = True,
        first_n_files: Optional[int] = None,
        pinned_memory_size: int = 10000,
    ):
        self.bigwig_paths = sorted(
            interpret_path(bigwig_path, file_extensions=file_extensions, crawl=crawl)
        )[:first_n_files]
        self.bigwigs = [BigWig(path, id=i) for i, path in enumerate(self.bigwig_paths)]

        (
            self.chromosome_offset_dict,
            self.local_chrom_ids_to_offset_matrix,
        ) = self.create_global_position_system()

        max_rows_per_chunk = max([bigwig.max_rows_per_chunk for bigwig in self.bigwigs])
        self.pinned_memory_size = pinned_memory_size
        self._memory_bank: Optional[MemoryBank] = None
        self._decoder = Decoder(
            max_num_chunks=10000 * 100,  # 512 * len(self.bigwigs) * 10,
            max_rows_per_chunk=max_rows_per_chunk,
            max_uncompressed_chunk_size=max_rows_per_chunk * 12 + 24,
            chromosome_offsets=self.local_chrom_ids_to_offset_matrix,
        )
        self._out: cp.ndarray = cp.zeros((len(self), 1, 1), dtype=cp.float32)

        self.run_indexing()

    def run_indexing(self) -> None:
        for bigwig in self.bigwigs:
            bigwig.run_indexing(
                chromosome_offsets=self.local_chrom_ids_to_offset_matrix[bigwig.id]
            )

    def __len__(self) -> int:
        return len(self.bigwigs)

    def _get_memory_bank(self) -> MemoryBank:
        if self._memory_bank:
            return self._memory_bank
        self._memory_bank = MemoryBank(nbytes=self.pinned_memory_size, elastic=True)
        return self._memory_bank

    def _get_out_tensor(self, batch_size: int, sequence_length: int) -> cp.ndarray:
        shape = (len(self), batch_size, sequence_length)
        if self._out.shape != shape:
            self._out = cp.zeros(shape, dtype=cp.float32)
        return self._out

    def get_batch(
        self,
        chromosomes: Union[Sequence[str], npt.NDArray[np.generic]],
        start: Union[Sequence[int], npt.NDArray[np.int64]],
        end: Union[Sequence[int], npt.NDArray[np.int64]],
        out: Optional[cp.ndarray] = None,
    ) -> cp.ndarray:
        memory_bank = self._get_memory_bank()
        memory_bank.reset()

        if out is None:
            sequence_length = end[0] - start[0]
            out = self._get_out_tensor(len(start), sequence_length)

        abs_start = self.make_positions_global(chromosomes, start)
        abs_end = self.make_positions_global(chromosomes, end)

        n_chunks_per_bigwig = []
        bigwig_ids = []
        for bigwig in self.bigwigs:
            offsets, sizes = bigwig.get_batch_offsets_and_sizes_with_global_positions(
                abs_start, abs_end
            )
            n_chunks_per_bigwig.append(len(offsets))
            bigwig_ids.extend([bigwig.id] * len(offsets))
            # read chunks into preallocated memory
            memory_bank.add_many(
                bigwig.store._fh,
                offsets,
                sizes,
                skip_bytes=2,
            )

        # bring the gpu
        bigwig_ids = cp.asarray(bigwig_ids, dtype=cp.uint32)
        gpu_byte_array, comp_chunks, compressed_chunk_sizes = memory_bank.to_gpu()
        _, start, end, value, n_rows_for_chunks = self._decoder.decode(
            gpu_byte_array, comp_chunks, compressed_chunk_sizes, bigwig_ids=bigwig_ids
        )

        chunk_row_numbers = cp.pad(cp.cumsum(n_rows_for_chunks), (1, 0))
        i = 0
        for n_chunks, partial_out in zip(n_chunks_per_bigwig, out):
            bigwig_end = i + n_chunks

            row_number_start = chunk_row_numbers[i]
            row_number_end = chunk_row_numbers[bigwig_end]

            intervals_to_values(
                track_starts=start[row_number_start:row_number_end],
                track_ends=end[row_number_start:row_number_end],
                track_values=value[row_number_start:row_number_end],
                query_starts=cp.asarray(abs_start, dtype=cp.uint32),
                query_ends=cp.asarray(abs_end, dtype=cp.uint32),
                out=partial_out,
            )
            i = bigwig_end
        return cp.transpose(out, (1, 0, 2))

    def make_positions_global(
        self,
        chromosomes: Union[Sequence[str], npt.NDArray[np.generic]],
        positions: Union[Sequence[int], npt.NDArray[np.int64]],
    ) -> npt.NDArray[np.int64]:
        """
        Take a set of positions and corresponding chromosomes
        and transform those positions into a coordinates that
        run over all chromosomes by adding specific offsets for
        chromosomes.

        Args:
            chromosomes: chromosome keys like "chr1, chr2...."
            positions: numpy array of integer positions on the chromosomes

        Returns:
            numpy array with "global" positions

        """
        offsets = np.array(
            [self.chromosome_offset_dict[chrom] for chrom in chromosomes]
        )
        return positions + offsets

    def intervals(
        self,
        include_chromosomes: Union[Literal["all", "standard"], Sequence[str]] = "all",
        exclude_chromosomes: Optional[Sequence[str]] = None,
        merge: bool = False,
        threshold: float = 0.0,
        merge_allow_gap: int = 0,
        batch_size: int = 4096,
    ) -> pd.DataFrame:
        intervals = pd.concat(
            [
                bw.intervals(
                    include_chromosomes=include_chromosomes,
                    exclude_chromosomes=exclude_chromosomes,
                    threshold=threshold,
                    merge=merge,
                    merge_allow_gap=merge_allow_gap,
                    memory_bank=self._get_memory_bank(),
                    decoder=self._decoder,
                    batch_size=batch_size,
                )
                for bw in self.bigwigs
            ]
        )
        if merge:
            intervals = merge_interval_dataframe(
                intervals, is_sorted=False, allow_gap=merge_allow_gap
            )
        return intervals

    def chromosomes(
        self,
        include_chromosomes: Union[Literal["all", "standard"], Sequence[str]] = "all",
        exclude_chromosomes: Optional[Sequence[str]] = None,
    ) -> set[str]:
        """
        Get a list of all the chromosomes present in files.

        Args:
            include_chromosomes: chromosome keys you want to include.
                Can alternatively be "all" or "standard".
            exclude_chromosomes: Optional set of chromosomes to exclude.

        Returns: a list of chromosomes over all BigWig files in the format:
            chr1_kidney, chr3_heart etc...

        """
        return {
            f"{chrom}_{path.stem}"
            for bigwig, path in zip(self.bigwigs, self.bigwig_paths)
            for chrom in bigwig.chromosomes(
                include_chromosomes=include_chromosomes,
                exclude_chromosomes=exclude_chromosomes,
            )
        }

    def get_chromosomes_present_in_all_files(
        self,
        include_chromosomes: Union[Literal["all", "standard"], Sequence[str]] = "all",
        exclude_chromosomes: Optional[Sequence[str]] = None,
    ) -> list[str]:
        """
        Get the subset of chromosome keys (i.e. "chr1", "chr4"...) that
        is present in every single BigWig file. The intersection of sets
        of chromosome keys over all the files.
        Args:
            include_chromosomes: chromosome keys you want to include.
                Can alternatively be "all" or "standard".
            exclude_chromosomes: Optional set of chromosomes to exclude.

        Returns:
            A list of chromosomes in a nicely sorted order.

        """
        return chromosome_sort(
            set.intersection(
                *[
                    set(
                        bigwig.chromosomes(
                            include_chromosomes=include_chromosomes,
                            exclude_chromosomes=exclude_chromosomes,
                        )
                    )
                    for bigwig in self.bigwigs
                ]
            )
        )

    def create_global_position_system(
        self,
    ) -> tuple[dict[str, int], npt.NDArray[np.int64]]:
        """
        Bigwig files internally use chrom_ids which are integers mapping to the different chromosomes.
        The problem is that this mapping to the actual chromosomes can actually be different from file
        to file. For performance reasons we want to decompress chunks from different files at the same
        time and while we are at it convert combinations of (chrom_id, position) to a positions system
        of a single integer that runs over all the chromosomes. For this we need to create a matrix
        local_chrom_ids_to_offset of size (n_files x n_chrom_ids) that we can use later on to add
        offsets to positions after a quick indexing operation:

            offsets = local_chrom_ids_to_offset[file_ids, chrom_ids]

        Returns:
            Tuple(Dict mapping chromosome keys to offsets like {"chr1": 200, "chr2": 250, },
                  np.ndarray size (n_files x n_chrom_ids) as described.

        """
        # find all chromosome keys (i.e. chr1, chr2...)
        unique_chromosomes = chromosome_sort(
            {
                chromosome
                for bigwig in self.bigwigs
                for chromosome in bigwig.chromosomes()
            }
        )
        global_chromosome_ids = {key: i for i, key in enumerate(unique_chromosomes)}
        max_local_chromosome_id = max(
            [
                chrom_id
                for bigwig in self.bigwigs
                for chrom_id in bigwig.chrom_to_chrom_id.values()
            ]
        )
        chrom_id_mapping_matrix = np.zeros(
            shape=(len(self.bigwigs), max_local_chromosome_id + 1), dtype=np.uint32
        )

        chromosome_sizes = np.array(
            [
                max(  # type: ignore
                    [
                        bigwig.chromosome_sizes.get(chromosome, 0)
                        for bigwig in self.bigwigs
                    ]
                )
                for chromosome in unique_chromosomes
            ]
        )
        chromosome_offsets = np.pad(np.cumsum(chromosome_sizes), (1, 0))
        chromosome_offset_dict = {
            key: size for key, size in zip(unique_chromosomes, chromosome_offsets)
        }
        if chromosome_offsets[-1] > 2**32:
            # exceeding this would be a problem as all the positions
            # in bigwig files are encoded in 32 bits. Making a global
            # base position that exceeds this number would be a problem.
            logging.warning(
                "The sum of sizes of the chromosomes exceeds the size of uint32"
            )
        for bigwig in self.bigwigs:
            for key, chrom_id in bigwig.chrom_to_chrom_id.items():
                chrom_id_mapping_matrix[bigwig.id, chrom_id] = global_chromosome_ids[
                    key
                ]
        local_chrom_ids_to_offset = chromosome_offsets[chrom_id_mapping_matrix]
        return chromosome_offset_dict, local_chrom_ids_to_offset


def get_bigwig_files_from_path(
    path: Path, file_extensions: Iterable[str] = (".bigWig", ".bw"), crawl: bool = True
) -> list[Path]:
    """
    For a path object, get all bigwig files. If the path is directly pointing
    to a bigwig file, a list with just that file is returned. In case of a
    directory all files are gathered with file extensions that are part of
    file_extensions.
    Args:
        path: or upath.Path object. Either directory or BigWig file
        file_extensions: used to find all BigWig files in a directory.
            Default: (".bigWig", ".bw")
        crawl: whether to find BigWig files in subdirectories. Default: True.

    Returns: list of paths to BigWig files.

    """
    if not path.exists():
        raise FileNotFoundError(f"No such file or directory: {path}")
    elif path.is_dir():
        if crawl:
            pattern = "**/*"
        else:
            pattern = "*"
        return [
            file
            for extension in file_extensions
            for file in path.glob(f"{pattern}{extension}")
        ]
    return [path]


def interpret_path(
    bigwig_path: Union[
        Union[str, "os.PathLike[Any]"], Iterable[Union[str, "os.PathLike[Any]"]]
    ],
    file_extensions: Iterable[str] = (".bigWig", ".bw"),
    crawl: bool = True,
) -> list[Path]:
    """
    Get all bigwig files for a path. Also excepts strings.
    If the path is directly pointing to a bigwig file, a list with just that
    file is returned. In case of a directory all files are gathered with file
    extensions that are part of file_extensions.
    Args:
        bigwig_path: str, Path or upath.Path object. Either directory or BigWig file
        file_extensions: used to find all BigWig files in a directory.
            Default: (".bigWig", ".bw")
        crawl: whether to find BigWig files in subdirectories. Default: True.

    Returns: list of paths to BigWig files.

    """
    if isinstance(bigwig_path, str) or isinstance(bigwig_path, os.PathLike):
        return get_bigwig_files_from_path(
            UPath(bigwig_path), file_extensions=file_extensions, crawl=crawl
        )

    elif isinstance(bigwig_path, Iterable):
        return [
            path
            for element in bigwig_path
            for path in interpret_path(
                element, file_extensions=file_extensions, crawl=crawl
            )
        ]
    raise ValueError(
        f"Can not interpret {bigwig_path} as path or a collection of paths."
    )

bigwig_loader/collection.py
pfizer-opensource/bigwig-loader
import math
from pathlib import Path
from typing import Any
from typing import Callable
from typing import Iterator
from typing import Literal
from typing import Optional
from typing import Sequence
from typing import Union

import cupy as cp
import numpy as np
import pandas as pd

from bigwig_loader.collection import BigWigCollection
from bigwig_loader.genome import Genome
from bigwig_loader.position_sampler import PositionSampler


class BigWigDataset:
    """
    Dataset over FASTA files and BigWig profiles.
    Args:
        regions_of_interest: pandas dataframe with intervals within the chromosomes to
            sample from. These intervals are not intervals of sequence_length
            but regions of the chromosome that were selected up front. For
            instance regions in which called peaks are located or a threshold
            value is exceeded.
        collection: instance of bigwig_loader.BigWigCollection or path to BigWig
            Directory or list of BigWig files.
        reference_genome_path: path to fasta file containing the reference genome.
        sequence_length: number of base pairs in input sequence
        center_bin_to_predict: if given, only do prediction on a central window. Should be
            smaller than or equal to sequence_length. If not given will be the same as
            sequence_length.
        window_size: used to down sample the resolution of the target from sequence_length
        batch_size: batch size returned by the dataset
        super_batch_size: batch size that is used in the background to load data from
            bigwig files. Should be larger than batch_size. If None, it will be equal to
            batch_size.
        batches_per_epoch: because the length of an epoch is slightly arbitrary here,
            the number of batches can be set by hand. If not the number of batches per
            epoch will be (totol number of bases in combined intervals) // sequence_length // batch_size
        maximum_unknown_bases_fraction: maximum number of bases in an input sequence that
            is unknown.
        sequence_encoder: encoder to apply to the sequence. Default: bigwig_loader.util.onehot_sequences
    """

    def __init__(
        self,
        regions_of_interest: pd.DataFrame,
        collection: Union[str, Sequence[str], Path, Sequence[Path], BigWigCollection],
        reference_genome_path: Path,
        sequence_length: int = 1000,
        center_bin_to_predict: Optional[int] = None,
        batch_size: int = 256,
        super_batch_size: Optional[int] = None,
        batches_per_epoch: Optional[int] = None,
        maximum_unknown_bases_fraction: float = 0.1,
        sequence_encoder: Optional[
            Union[Callable[[Sequence[str]], Any], Literal["onehot"]]
        ] = "onehot",
        file_extensions: Sequence[str] = (".bigWig", ".bw"),
        crawl: bool = True,
        first_n_files: Optional[int] = None,
        position_sampler_buffer_size: int = 100000,
    ):
        self.batch_size = batch_size
        super_batch_size = super_batch_size or batch_size
        if super_batch_size < batch_size:
            raise AttributeError(
                f"super_batch_size {batch_size} can not be smaller than batch_size {batch_size}"
            )
        self.batches_per_epoch = (
            batches_per_epoch
            or (regions_of_interest["end"] - regions_of_interest["start"]).sum()
            // sequence_length
            // batch_size
        )

        super_batches_per_epoch = math.ceil(
            batch_size * self.batches_per_epoch / super_batch_size
        )

        self._super_dataset = BigWigSuperDataset(
            regions_of_interest=regions_of_interest,
            collection=collection,
            reference_genome_path=reference_genome_path,
            sequence_length=sequence_length,
            center_bin_to_predict=center_bin_to_predict,
            batch_size=super_batch_size,
            batches_per_epoch=super_batches_per_epoch,
            maximum_unknown_bases_fraction=maximum_unknown_bases_fraction,
            sequence_encoder=sequence_encoder,
            file_extensions=file_extensions,
            crawl=crawl,
            first_n_files=first_n_files,
            position_sampler_buffer_size=position_sampler_buffer_size,
        )
        self._super_batch_sequences: cp.ndarray = None
        self._super_batch_targets: cp.ndarray = None
        self._n = 0
        self._offset = 0

    def __iter__(self) -> Iterator[tuple[Any, cp.ndarray]]:
        self._n = 0
        self._offset = 0
        iter(self._super_dataset)
        return self

    def __next__(self) -> tuple[Any, cp.ndarray]:
        if self._n < self.batches_per_epoch:
            self._n += 1
            if (
                self._super_batch_sequences is None
                or self._super_batch_targets is None
                or self._offset + self.batch_size > self._super_dataset.batch_size
            ):
                self._super_batch_sequences, self._super_batch_targets = next(
                    self._super_dataset
                )
                self._offset = 0

            sequences = self._super_batch_sequences[
                self._offset : self._offset + self.batch_size
            ]
            target = self._super_batch_targets[
                self._offset : self._offset + self.batch_size
            ]
            self._offset += self.batch_size
            return sequences, target

        raise StopIteration


class BigWigSuperDataset:
    """
    Dataset over FASTA files and BigWig profiles.
    Args:
        regions_of_interest: pandas dataframe with intervals within the chromosomes to
            sample from. These intervals are not intervals of sequence_length
            but regions of the chromosome that were selected up front. For
            instance regions in which called peaks are located or a threshold
            value is exceeded.
        collection: instance of bigwig_loader.BigWigCollection or path to BigWig
            Directory or list of BigWig files.
        reference_genome_path: path to fasta file containing the reference genome.
        sequence_length: number of base pairs in input sequence
        center_bin_to_predict: if given, only do prediction on a central window. Should be
            smaller than or equal to sequence_length. If not given will be the same as
            sequence_length.
        window_size: used to down sample the resolution of the target from sequence_length
        batch_size: batch size
        batches_per_epoch: because the length of an epoch is slightly arbitrary here,
            the number of batches can be set by hand. If not the number of batches per
            epoch will be (totol number of bases in combined intervals) // sequence_length // batch_size
        maximum_unknown_bases_fraction: maximum number of bases in an input sequence that
            is unknown.
        sequence_encoder: encoder to apply to the sequence. Default: bigwig_loader.util.onehot_sequences
        position_samples_buffer_size: number of intervals picked up front by the position sampler.
            When all intervals are used, new intervals are picked.

    """

    def __init__(
        self,
        regions_of_interest: pd.DataFrame,
        collection: Union[str, Sequence[str], Path, Sequence[Path], BigWigCollection],
        reference_genome_path: Path,
        sequence_length: int = 1000,
        center_bin_to_predict: Optional[int] = None,
        batch_size: int = 256,
        batches_per_epoch: Optional[int] = None,
        maximum_unknown_bases_fraction: float = 0.1,
        sequence_encoder: Optional[
            Union[Callable[[Sequence[str]], Any], Literal["onehot"]]
        ] = "onehot",
        file_extensions: Sequence[str] = (".bigWig", ".bw"),
        crawl: bool = True,
        first_n_files: Optional[int] = None,
        position_sampler_buffer_size: int = 100000,
    ):
        super().__init__()

        self.regions_of_interest = regions_of_interest
        if isinstance(collection, BigWigCollection):
            self._bigwig_collection: Optional[BigWigCollection] = collection
            self._bigwig_path: Optional[
                Union[str, Sequence[str], Path, Sequence[Path]]
            ] = None
        else:
            self._bigwig_collection = None
            self._bigwig_path = collection
        self.reference_genome_path = reference_genome_path
        self.sequence_length = sequence_length
        if center_bin_to_predict:
            self.center_bin_to_predict = center_bin_to_predict
        else:
            self.center_bin_to_predict = sequence_length
        self.batch_size = batch_size
        self.batches_per_epoch = (
            batches_per_epoch
            or (regions_of_interest["end"] - regions_of_interest["start"]).sum()
            // sequence_length
            // batch_size
        )
        self._n = 0
        self.maximum_unknown_bases_fraction = maximum_unknown_bases_fraction
        self.sequence_encoder = sequence_encoder
        self._first_n_files = first_n_files
        self._file_extensions = file_extensions
        self._crawl = crawl
        self._genome: Optional[Genome] = None
        self._prepared_out: Optional[cp.ndarray] = None
        self._position_sampler_buffer_size = position_sampler_buffer_size

    @property
    def genome(self) -> Genome:
        """
        Setup genome object to get sequences from the reference genome.
        """
        if not self._genome:
            position_sampler = PositionSampler(
                regions_of_interest=self.regions_of_interest,
                buffer_size=self._position_sampler_buffer_size,
            )
            self._genome = Genome(
                self.reference_genome_path,
                sequence_length=self.sequence_length,
                position_sampler=position_sampler,
                batch_size=self.batch_size,
                maximum_unknown_bases_fraction=self.maximum_unknown_bases_fraction,
                encoder=self.sequence_encoder,
            )
        return self._genome

    @property
    def bigwig_collection(self) -> BigWigCollection:
        """
        Setup the BigWigCollection to get values from all bigwig
        files from
        """

        if self._bigwig_collection is not None:
            return self._bigwig_collection

        elif self._bigwig_path is not None:
            self._bigwig_collection = BigWigCollection(
                self._bigwig_path,
                file_extensions=self._file_extensions,
                crawl=self._crawl,
                first_n_files=self._first_n_files,
            )
            return self._bigwig_collection
        else:
            raise RuntimeError(
                f"{self}._bigwig_collection and {self}._bigwig_path are bot None. At least one should be set."
            )

    @property
    def _out(self) -> cp.ndarray:
        if self._prepared_out is None:
            self._prepared_out = cp.zeros(
                (
                    len(self.bigwig_collection),
                    self.batch_size,
                    self.center_bin_to_predict,
                ),
                dtype=cp.float32,
            )
        return self._prepared_out

    def __iter__(self) -> Iterator[tuple[Any, cp.ndarray]]:
        self._n = 0
        return self

    def __next__(self) -> tuple[Any, cp.ndarray]:
        if self._n < self.batches_per_epoch:
            self._n += 1
            positions, sequences = self.genome.get_batch()
            chromosomes, center = zip(*positions)
            start = np.array(center) - (self.center_bin_to_predict // 2)
            end = start + self.center_bin_to_predict
            target = self.bigwig_collection.get_batch(
                chromosomes,
                start,
                end,
                out=self._out,
            )
            return sequences, target
        raise StopIteration

bigwig_loader/dataset.py
pfizer-opensource/bigwig-loader
import hashlib
import logging
import subprocess
import urllib.request
from pathlib import Path
from typing import BinaryIO

from bigwig_loader import config
from bigwig_loader.collection import get_bigwig_files_from_path

LOGGER = logging.getLogger(__name__)


def download_example_data() -> None:
    get_reference_genome()
    get_example_bigwigs_files()


def get_reference_genome(reference_genome_path: Path = config.reference_genome) -> Path:
    compressed_file = reference_genome_path.with_suffix(".fasta.gz")
    if reference_genome_path.exists():
        return reference_genome_path
    elif compressed_file.exists():
        subprocess.run(["bgzip", "-d", compressed_file])
    else:
        LOGGER.info("Need reference genome for tests. Downloading it from ENCODE.")
        url = "https://www.encodeproject.org/files/GRCh38_no_alt_analysis_set_GCA_000001405.15/@@download/GRCh38_no_alt_analysis_set_GCA_000001405.15.fasta.gz"
        urllib.request.urlretrieve(url, compressed_file)
        subprocess.run(["bgzip", "-d", compressed_file])
    return reference_genome_path


EXAMPLE_FILES = {
    "ENCFF270YCY.bigWig": (
        "https://www.encodeproject.org/files/ENCFF270YCY/@@download/ENCFF270YCY.bigWig",
        "64886b78f4cd70531fceffbaeafd4adb",
    ),
    "ENCFF270ISQ.bigWig": (
        "https://www.encodeproject.org/files/ENCFF270ISQ/@@download/ENCFF270ISQ.bigWig",
        "5fef60f9f1e43b9a17075c352650865e",
    ),
}


def checksum_md5(f: BinaryIO, *, chunk_size: int = 10 * 1024 * 1024) -> str:
    """return the md5sum"""
    m = hashlib.md5(b"", usedforsecurity=False)
    for data in iter(lambda: f.read(chunk_size), b""):
        m.update(data)
    return m.hexdigest()


def get_example_bigwigs_files(bigwig_dir: Path = config.bigwig_dir) -> Path:
    available_files = [pth.name for pth in get_bigwig_files_from_path(bigwig_dir)]
    if len(available_files) < 2:
        for fn, (url, md5) in EXAMPLE_FILES.items():
            file = bigwig_dir / fn
            if not file.exists():
                urllib.request.urlretrieve(url, file)
            with file.open(mode="rb") as f:
                if checksum_md5(f) != md5:
                    raise RuntimeError(f"{fn} has incorrect checksum!")
    return bigwig_dir

bigwig_loader/download_example_data.py
pfizer-opensource/bigwig-loader
from pathlib import Path
from typing import Any
from typing import Callable
from typing import Iterator
from typing import Literal
from typing import Optional
from typing import Sequence
from typing import Union

import fsspec
from pyfaidx import Fasta

from bigwig_loader.util import fraction_non_standard
from bigwig_loader.util import onehot_sequences

ENCODERS: dict[str, Callable[[Sequence[str]], Any]] = {"onehot": onehot_sequences}


class Genome:
    def __init__(
        self,
        reference_genome_path: Path,
        sequence_length: int,
        position_sampler: Iterator[tuple[str, int]],
        batch_size: int,
        maximum_unknown_bases_fraction: float = 0.1,
        encoder: Optional[
            Union[Callable[[Sequence[str]], Any], Literal["onehot"]]
        ] = "onehot",
    ):
        self.reference_genome_path = reference_genome_path
        self.sequence_length = sequence_length
        self.genome = Fasta(fsspec.open(self.reference_genome_path))
        self.position_sampler = position_sampler
        self.batch_size = batch_size
        self.maximum_unknown_bases_fraction = maximum_unknown_bases_fraction
        self.encoder = self._select_encoder(encoder)

    def get_batch(self) -> tuple[list[tuple[str, int]], Any]:
        selected_positions: list[tuple[str, int]] = []
        sequences = []
        while len(selected_positions) < self.batch_size:
            chromosome, center = next(self.position_sampler)
            start = center - (self.sequence_length // 2)
            end = start + self.sequence_length
            sequence = self._get_sequence(chromosome, start, end)
            if not sequence:
                continue
            if fraction_non_standard(sequence) > self.maximum_unknown_bases_fraction:
                continue
            sequences.append(sequence)
            selected_positions.append((chromosome, center))
        if self.encoder is not None:
            sequences = self.encoder(sequences)
        return selected_positions, sequences

    def _get_sequence(
        self, chrom: str, start: int, end: int, strand: Literal["+", "-"] = "+"
    ) -> Any:
        return self.genome[chrom][start:end].seq

    def _select_encoder(
        self,
        encoder: Optional[Union[Callable[[Sequence[str]], Any], Literal["onehot"]]],
    ) -> Optional[Callable[[Sequence[str]], Any]]:
        if encoder is None:
            return None
        if callable(encoder):
            return encoder
        if isinstance(encoder, str) and encoder in ENCODERS:
            return ENCODERS[encoder]
        raise ValueError(
            f"{encoder} is not a valid encoder. It is not listed in {list(ENCODERS.keys())} and also not a Callable."
        )


def register_encoder(encoder_id: str, encoder: Callable[[Sequence[str]], Any]) -> None:
    ENCODERS[encoder_id] = encoder

bigwig_loader/genome.py
pfizer-opensource/bigwig-loader
from typing import Optional
from typing import Sequence
from typing import Union

import cupy as cp
import numcodecs
import numpy as np
import numpy.typing as npt
from numcodecs.compat import ensure_contiguous_ndarray_like

codec = numcodecs.registry.get_codec(
    {
        "id": "nvcomp_batch",
        "algorithm": "Deflate",
    }
)


class Decoder:
    # based on: `kvikio.nvcomp_codec.NvCompBatchCodec`

    def __init__(
        self,
        max_num_chunks: int,
        max_rows_per_chunk: int,
        max_uncompressed_chunk_size: int = (2048 * 12),
        chromosome_offsets: Optional[
            Union[Sequence[int], npt.NDArray[np.int64]]
        ] = None,
    ):
        self.max_uncompressed_chunk_size = max_uncompressed_chunk_size

        self.uncomp_chunks = [
            cp.empty(max_uncompressed_chunk_size, dtype=cp.uint8)
            for _ in range(max_num_chunks)
        ]
        self.uncomp_chunk_ptrs = cp.array(
            [c.data.ptr for c in self.uncomp_chunks], dtype=cp.uint64
        )
        self.uncomp_chunk_sizes = cp.array(
            [max_uncompressed_chunk_size] * max_num_chunks, dtype=cp.uint64
        )
        self.actual_uncomp_chunk_sizes = cp.empty(max_num_chunks, dtype=cp.uint64)
        self.statuses = cp.empty(max_num_chunks, dtype=cp.int32)
        self.max_rows_per_chunk = max_rows_per_chunk

        self.range_of_indexes = cp.arange(max_rows_per_chunk)
        if chromosome_offsets is not None:
            self.chromosome_offsets = cp.asarray(chromosome_offsets, dtype="uint32")
        else:
            self.chromosome_offsets = None

    def decode(
        self,
        gpu_array: cp.ndarray,
        comp_chunks: cp.ndarray,
        comp_chunk_sizes: cp.ndarray,
        bigwig_ids: Optional[cp.ndarray] = None,
    ) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray]:
        num_chunks = len(comp_chunks)

        max_chunk_size = int(self.max_uncompressed_chunk_size)

        # Get temp buffer size.
        temp_size = codec._algo.get_decompress_temp_size(num_chunks, max_chunk_size)

        temp_buf = cp.empty(temp_size, dtype=cp.uint8)

        uncomp_chunks = self.uncomp_chunks[:num_chunks]
        uncomp_chunk_ptrs = self.uncomp_chunk_ptrs[:num_chunks]
        uncomp_chunk_sizes = self.uncomp_chunk_sizes[:num_chunks]

        actual_uncomp_chunk_sizes = self.actual_uncomp_chunk_sizes[:num_chunks]

        statuses = self.statuses[:num_chunks]

        codec._algo.decompress(
            comp_chunks,  # const void* const*
            comp_chunk_sizes,  # const size_t*
            num_chunks,  # ??
            temp_buf,  #
            uncomp_chunk_ptrs,
            uncomp_chunk_sizes,
            actual_uncomp_chunk_sizes,
            statuses,
            codec._stream,
        )

        chrom_ids, start, end, value, n_rows_for_chunks = self.post_process(
            cp.concatenate(uncomp_chunks), bigwig_ids
        )
        return chrom_ids, start, end, value, n_rows_for_chunks

    def post_process(
        self, arr: cp.ndarray, bigwig_ids: Optional[cp.ndarray] = None
    ) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray]:
        max_rows_per_chunk = self.max_rows_per_chunk
        n_rows_for_chunks = arr.view(dtype="u2")[
            11 :: (max_rows_per_chunk * 12 + 24) // 2
        ]

        chrom_ids = arr.view(dtype="<u4")[0 :: (max_rows_per_chunk * 12 + 24) // 4]

        start = (
            arr.view(dtype="<u4")
            .reshape((-1, 3))[:, 0]
            .reshape((-1, max_rows_per_chunk + 2))[:, 2:]
        )
        end = (
            arr.view(dtype="<u4")
            .reshape((-1, 3))[:, 1]
            .reshape((-1, max_rows_per_chunk + 2))[:, 2:]
        )
        value = (
            arr.view(dtype="<f4")
            .reshape((-1, 3))[:, 2]
            .reshape((-1, max_rows_per_chunk + 2))[:, 2:]
        )

        if self.chromosome_offsets is not None and bigwig_ids is not None:
            chrom_offsets = self.chromosome_offsets[bigwig_ids, chrom_ids]
            chrom_offsets = chrom_offsets[..., cp.newaxis]
            start += chrom_offsets
            end += chrom_offsets

        if (n_rows_for_chunks == max_rows_per_chunk).all():
            start = cp.ravel(start)
            end = cp.ravel(end)
            value = cp.ravel(value)
        else:
            mask = cp.greater.outer(n_rows_for_chunks, self.range_of_indexes)
            start = start[mask]
            end = end[mask]
            value = value[mask]

        return chrom_ids, start, end, value, n_rows_for_chunks

bigwig_loader/gpu_decompressor.py
pfizer-opensource/bigwig-loader
import cupy as cp

_zero = cp.asarray(0.0, dtype=cp.float32).item()

_cuda_kernel = """
extern "C" __global__
void intervals_to_values(
        const int* query_starts,
        const int* query_ends,
        const int* found_starts,
        const int* found_ends,
        const int* track_starts,
        const int* track_ends,
        const float* track_values,
        const int batch_size,
        const int sequence_length,
        const int max_number_intervals,
        const int max_interval_length,
        float* out
) {

    //printf("batch_size %d", batch_size);
    //printf("sequence_length %d", sequence_length);
    //printf("max_number_intervals %d", max_number_intervals);
    //printf("max_interval_length %d", max_interval_length);

    int thread = blockIdx.x * blockDim.x + threadIdx.x;

    int i = thread % batch_size;
    int j = (thread / batch_size)%max_number_intervals;
    int k = thread / (batch_size*max_number_intervals);

    if (i <  batch_size){

        int found_start_index = found_starts[i];
        int found_end_index = found_ends[i];
        int query_start = query_starts[i];
        int query_end = query_ends[i];

        int cursor = found_start_index + j;

        if (cursor < found_end_index){
            int interval_start = track_starts[cursor];
            int interval_end = track_ends[cursor];
            int start_index = max(interval_start - query_start, 0);
            int end_index = (i * sequence_length) + min(interval_end, query_end) - query_start;
            int position = (i * sequence_length) + start_index + k;

            if (position < end_index){
                out[position] = track_values[cursor];
            }
        }
    }
}
"""

cuda_kernel = cp.RawKernel(_cuda_kernel, "intervals_to_values")
cuda_kernel.compile()


def intervals_to_values(
    track_starts: cp.ndarray,
    track_ends: cp.ndarray,
    track_values: cp.ndarray,
    query_starts: cp.ndarray,
    query_ends: cp.ndarray,
    out: cp.ndarray,
) -> cp.ndarray:
    out *= _zero
    found_starts = cp.searchsorted(track_ends, query_starts, side="right").astype(
        dtype=cp.int32
    )
    found_ends = cp.searchsorted(track_starts, query_ends, side="left").astype(
        dtype=cp.int32
    )

    sequence_length = (query_ends[0] - query_starts[0]).item()

    max_number_intervals = min(
        sequence_length, (found_ends - found_starts).max().item()
    )
    max_interval_length = min(sequence_length, (track_ends - track_starts).max().item())
    batch_size = query_starts.shape[0]
    n_threads_needed = batch_size * max_interval_length * max_number_intervals
    grid_size, block_size = get_grid_and_block_size(n_threads_needed)

    cuda_kernel(
        (grid_size,),
        (block_size,),
        (
            query_starts,
            query_ends,
            found_starts,
            found_ends,
            track_starts,
            track_ends,
            track_values,
            batch_size,
            sequence_length,
            max_number_intervals,
            max_interval_length,
            out,
        ),
    )
    return out


def get_grid_and_block_size(n_threads: int) -> tuple[int, int]:
    n_blocks_needed = cp.ceil(n_threads / 512).astype(dtype=cp.int32).item()
    if n_blocks_needed == 1:
        threads_per_block = n_threads
    else:
        threads_per_block = 512
    return n_blocks_needed, threads_per_block


def kernel_in_python(
    grid_size: int,
    block_size: int,
    args: tuple[
        cp.ndarray,
        cp.ndarray,
        cp.ndarray,
        cp.ndarray,
        cp.ndarray,
        cp.ndarray,
        cp.ndarray,
        int,
        int,
        int,
        int,
        cp.ndarray,
    ],
) -> list[float]:
    """Equivalent in python to cuda_kernel. Just for debugging."""

    (
        query_starts,
        query_ends,
        found_starts,
        found_ends,
        track_starts,
        track_ends,
        track_values,
        batch_size,
        sequence_length,
        max_number_intervals,
        max_interval_length,
        _,
    ) = args

    query_starts = query_starts.get().tolist()
    query_ends = query_ends.get().tolist()

    found_starts = found_starts.get().tolist()
    found_ends = found_ends.get().tolist()
    track_starts = track_starts.get().tolist()
    track_ends = track_ends.get().tolist()
    track_values = track_values.get().tolist()

    n_threads = grid_size * block_size

    out = [0.0] * 4 * batch_size

    for thread in range(n_threads):
        i = thread % batch_size
        j = (thread // batch_size) % max_number_intervals
        k = thread // (batch_size * max_number_intervals)
        print("---")
        print(i, j, k)

        if i < batch_size:
            found_start_index = found_starts[i]
            found_end_index = found_ends[i]
            query_start = query_starts[i]
            query_end = query_ends[i]

            cursor = found_start_index + j
            print("cursor", cursor)

            if cursor < found_end_index:
                interval_start = track_starts[cursor]
                interval_end = track_ends[cursor]
                start_index = max(interval_start - query_start, 0)
                end_index = (
                    (i * sequence_length) + min(interval_end, query_end) - query_start
                )
                position = (i * sequence_length) + start_index + k
                # position = start_index + k
                print("position", position)

                if position < end_index:
                    out[position] = track_values[cursor]
        print(out)

    print(out)
    return out

bigwig_loader/intervals_to_values_gpu.py
pfizer-opensource/bigwig-loader
import logging
from typing import BinaryIO
from typing import Iterable
from typing import Optional

import cupy as cp
import numpy as np


class MemoryBank:
    def __init__(self, nbytes: int = 10000, elastic: bool = True):
        self._mem = cp.cuda.alloc_pinned_memory(nbytes)
        self._view = memoryview(self._mem)
        # self._cpu_array =
        self._current_file_handle: Optional[BinaryIO] = None
        self.n_chunks = 0
        self.compressed_chunk_sizes: list[int] = []
        self.compressed_chunk_offsets: list[int] = [0]
        self.cumulative_n_chunks_per_file: list[int] = []
        self.elastic = elastic

    def reset(self) -> None:
        """
        Reset all the containing data to empty values. Can be used
        before a new batch comes in.
        """
        self._current_file_handle = None
        self.n_chunks = 0
        self.compressed_chunk_sizes = []
        self.compressed_chunk_offsets = [0]
        self.cumulative_n_chunks_per_file = []

    def increase_memory_size(self, new_size: int, generosity: float = 1.2) -> None:
        """
        Increase the size of the pinned memory when needed.
        Args:
            new_size: new size of pinned memory needed
            generosity: when increasing the memory, increase more
                by this factor to prevent having to update too often.

        """
        if new_size > self._mem.size():
            new_size = int(new_size * generosity)
            logging.debug(f"Increasing size of pinned memory to {new_size}")
            new_mem = cp.cuda.alloc_pinned_memory(new_size)
            new_view = memoryview(new_mem)
            new_view[: self._view.nbytes] = self._view
            self._mem = new_mem
            self._view = new_view

    def add(
        self, file_handle: BinaryIO, offset: int, size: int, skip_bytes: int = 2
    ) -> None:
        """
        Add one compressed chunk to the pinned memory.
        Args:
            file_handle: opened file
            offset: byte offset to start of compressed chunk
            size: byte size of compressed chunk
            skip_bytes: skips this number of bytes from offset
            increase_memory_size: whether to elastically increase the
                size of the pinned memory if the new chunk doesn't
                fit in anymore. (default: True)
        Returns:

        """
        offset += skip_bytes
        size -= skip_bytes
        if file_handle is not self._current_file_handle:
            self.cumulative_n_chunks_per_file.append(self.n_chunks)
            self._current_file_handle = file_handle
        self.n_chunks += 1
        current_offset = self.compressed_chunk_offsets[-1]
        new_offset = current_offset + size
        file_handle.seek(offset)
        # file_handle.readinto(self._view[current_offset:new_offset])
        if self.elastic:
            self.increase_memory_size(new_offset)
        file_handle.readinto(self._view[current_offset:new_offset])  # type: ignore
        self.compressed_chunk_offsets.append(new_offset)
        self.compressed_chunk_sizes.append(size)

    def add_many(
        self,
        file_handle: BinaryIO,
        offsets: Iterable[int],
        sizes: Iterable[int],
        skip_bytes: int = 2,
    ) -> None:
        """
        Add many compressed chunks from a single file.
        Args:
            file_handle: opened file
            offsets: byte offsets to starts of compressed chunks
            sizes: sizes of compressed chunks
            skip_bytes: skips this number of bytes from offset
            increase_memory_size: whether to elastically increase the
                size of the pinned memory if the new chunks don't
                fit in anymore. (default: True)

        """
        new_size = self.compressed_chunk_offsets[-1] + sum(sizes)
        if self.elastic:
            # better increase once for an entire batch than separately for each chunk
            self.increase_memory_size(new_size)
        for offset, size in zip(offsets, sizes):
            self.add(
                file_handle,
                offset,
                size,
                skip_bytes=skip_bytes,
            )

    def to_gpu(self) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray]:
        """
        Brings the data from pinned memory to GPU.
        Returns: everything needed for decompression on GPU:
            Tuple(gpu byte array, pointers to starts of compressed
                  chunks, sizes of compressed chunks)
        """
        compressed_chunk_sizes = cp.array(self.compressed_chunk_sizes, dtype=cp.uint64)
        # offsets = cp.array(self.compressed_chunk_offsets[:-1], dtype=cp.uint64)
        offsets = cp.pad(cp.cumsum(compressed_chunk_sizes), (1, 0))[:-1]

        cpu_array = np.frombuffer(
            self._mem, dtype="byte", count=self.compressed_chunk_offsets[-1]
        )

        logging.debug(f"CPU Array {cpu_array}")

        # gpu_byte_array = cp.asarray(self._mem)[:self.compressed_chunk_offsets[-1]]
        gpu_byte_array = cp.asarray(cpu_array)
        comp_chunks = offsets + gpu_byte_array.data.ptr
        return gpu_byte_array, comp_chunks, compressed_chunk_sizes


if __name__ == "__main__":
    bank = MemoryBank()
    file_handle = open("look.txt", "rb", buffering=0)

    bank.add(file_handle, offset=0, size=16)
    bank.add(file_handle, offset=16, size=16)

    file_handle = open("new.prof", "rb", buffering=0)

    bank.add(file_handle, offset=0, size=16)
    bank.add(file_handle, offset=16, size=16)

    print(np.frombuffer(bank._mem, dtype="int", count=100))
    print(bank.to_gpu())

    print(bank.compressed_chunk_offsets)
    print(bank.cumulative_n_chunks_per_file)

    bank.reset()

    print("reset----")

    file_handle = open("look.txt", "rb", buffering=0)

    bank.add(file_handle, offset=0, size=16)
    bank.add(file_handle, offset=16, size=16)

    file_handle = open("new.prof", "rb", buffering=0)

    bank.add(file_handle, offset=0, size=16)
    bank.add(file_handle, offset=16, size=16)

    print(np.frombuffer(bank._mem, dtype="byte", count=100))
    print("to gpu:")
    print(bank.to_gpu())

    print(bank.compressed_chunk_offsets)
    print(bank.cumulative_n_chunks_per_file)

bigwig_loader/memory_bank.py
pfizer-opensource/bigwig-loader
from __future__ import annotations

import struct
from dataclasses import dataclass
from typing import BinaryIO
from typing import Optional

BBI_HEADER_ENCODING = "<LHHQQQHHQQLQ"
ZOOM_HEADER_ENCODING = "<LLQQ"
TOTAL_SUMMARY_ENCODING = "<Qdddd"
CHROMOSOME_TREE_HEADER_ENCODING = "<LLLLQQ"
CHROMOSOME_TREE_NODE_ENCODING = "<??H"
RTREE_INDEX_HEADER_ENCODING = "<LLQLLLLQLL"
RTREE_NODE_ENCODING = "<??H"
WIG_SECTION_HEADER_ENCODING = "<LLLLLBBH"


@dataclass
class BBIHeader:
    magic: int
    version: int
    zoom_levels: int
    chromosome_tree_offset: int
    full_data_offset: int
    full_index_offset: int
    field_count: int
    defined_field_count: int
    auto_sql_offset: int
    total_summary_offset: int
    uncompress_buff_size: int
    reserved: int

    @classmethod
    def from_file(cls, file_object: BinaryIO) -> BBIHeader:
        return cls.from_bytes(file_object.read(64))

    @classmethod
    def from_bytes(cls, data: bytes) -> BBIHeader:
        return cls(*struct.unpack(BBI_HEADER_ENCODING, data))


@dataclass
class ZoomHeader:
    reduction_level: int
    reserved: int
    data_offset: int
    index_offset: int

    @classmethod
    def all(
        cls, file_object: BinaryIO, first_offset: int, n_zoom_levels: int
    ) -> list[ZoomHeader]:
        zoom_headers = []
        next_level_offset = first_offset
        for level in range(n_zoom_levels):
            zoom_header = cls.from_file_and_offset(file_object, next_level_offset)
            zoom_headers.append(zoom_header)
            next_level_offset = file_object.tell()
        return zoom_headers

    @classmethod
    def from_file_and_offset(cls, file_object: BinaryIO, offset: int) -> ZoomHeader:
        file_object.seek(offset, 0)
        return cls.from_bytes(file_object.read(24))

    @classmethod
    def from_bytes(cls, data: bytes) -> ZoomHeader:
        return cls(*struct.unpack(ZOOM_HEADER_ENCODING, data))


@dataclass
class TotalSummary:
    bases_covered: int
    min_val: float
    max_val: float
    sum_data: float
    sum_squares: float

    @classmethod
    def from_file_and_offset(cls, file_object: BinaryIO, offset: int) -> TotalSummary:
        file_object.seek(offset, 0)
        return cls.from_bytes(file_object.read(40))

    @classmethod
    def from_bytes(cls, data: bytes) -> TotalSummary:
        return cls(*struct.unpack(TOTAL_SUMMARY_ENCODING, data))


@dataclass
class ChromosomeTreeHeader:
    magic: int
    block_size: int
    key_size: int
    val_size: int
    item_count: int
    reserved: int

    @classmethod
    def from_file_and_offset(
        cls, file_object: BinaryIO, offset: int
    ) -> ChromosomeTreeHeader:
        file_object.seek(offset, 0)
        return cls.from_bytes(file_object.read(32))

    @classmethod
    def from_bytes(cls, data: bytes) -> ChromosomeTreeHeader:
        return cls(*struct.unpack(CHROMOSOME_TREE_HEADER_ENCODING, data))


@dataclass
class ChromosomeTreeLeafItem:
    key: str
    chrom_id: Optional[int]
    chrom_size: Optional[int]
    child_offset: Optional[int]
    child_node: Optional[object]


@dataclass
class ChromosomeTreeNode:
    is_leaf: bool
    reserved: bool
    count: int
    items: list[ChromosomeTreeLeafItem]

    @classmethod
    def from_file_and_offset(
        cls, file_object: BinaryIO, key_size: int, offset: Optional[int] = None
    ) -> ChromosomeTreeNode:
        if offset is not None:
            file_object.seek(offset, 0)
        is_leaf, reserved, count = struct.unpack(
            CHROMOSOME_TREE_NODE_ENCODING, file_object.read(4)
        )

        items = []
        location = file_object.tell()
        for i in range(count):
            location = file_object.seek(location, 0)
            chrom_id = None
            chrom_size = None
            child_offset = None
            child_node = None
            bits = file_object.read(key_size + 8)
            if is_leaf:
                key, chrom_id, chrom_size = struct.unpack(f"<{key_size}sLL", bits)
                location = file_object.tell()
            else:
                key, child_offset = struct.unpack(f"<{key_size}Q", bits)
                location = file_object.tell()
                child_node = ChromosomeTreeNode.from_file_and_offset(
                    file_object, key_size=key_size, offset=child_offset
                )

            key = key.decode("utf-8").strip("\x00")
            items.append(
                ChromosomeTreeLeafItem(
                    key, chrom_id, chrom_size, child_offset, child_node=child_node
                )
            )
        return cls(is_leaf, reserved, count, items)


@dataclass
class RTreeIndexHeader:
    magic: int
    block_size: int
    item_count: int
    start_chrom_ix: int
    start_base: int
    end_chrom_ix: int
    end_base: int
    end_file_offset: int
    items_per_slot: int
    reserved: int

    @classmethod
    def from_file_and_offset(
        cls, file_object: BinaryIO, offset: int
    ) -> RTreeIndexHeader:
        file_object.seek(offset, 0)
        return cls.from_bytes(file_object.read(48))

    @classmethod
    def from_bytes(cls, data: bytes) -> RTreeIndexHeader:
        return cls(*struct.unpack(RTREE_INDEX_HEADER_ENCODING, data))


@dataclass
class WIGSectionHeader:
    chrom_id: int
    chrom_start: int
    chrom_end: int
    item_step: int
    item_span: int
    type: str
    item_count: int

    @classmethod
    def from_bytes(cls, data: bytes) -> WIGSectionHeader:
        (
            chrom_id,
            chrom_start,
            chrom_end,
            item_step,
            item_span,
            section_type,
            reserved,
            item_count,
        ) = struct.unpack(WIG_SECTION_HEADER_ENCODING, data)

        section_type = {1: "bed_graph", 2: "variable_step", 3: "fixed_step"}[
            section_type
        ]
        return cls(
            chrom_id,
            chrom_start,
            chrom_end,
            item_step,
            item_span,
            section_type,
            item_count,
        )


@dataclass
class ZoomRecord:
    chrom_id: int
    chrom_start: int
    chrom_end: int
    valid_count: int
    min_val: float
    max_val: float
    sum_data: float
    sum_of_squares: float

    @classmethod
    def from_bytes(cls, data: bytes) -> ZoomRecord:
        return cls(*struct.unpack("<LLLLffff", data))


@dataclass
class ZoomLevel:
    zoom_count: int
    zoom_records: list[ZoomRecord]

    @classmethod
    def from_file_and_offset(
        cls, file_object: BinaryIO, offset: Optional[int] = None
    ) -> ZoomLevel:
        if offset:
            file_object.seek(offset, 0)
        (zoom_count,) = struct.unpack("<L", file_object.read(4))
        print("zoom count:", zoom_count)
        zoom_records = []
        for i in range(zoom_count):
            record = ZoomRecord.from_bytes(file_object.read(32))
            if record.min_val > record.max_val:
                raise ValueError("ZoomRecord min_val is larger than max_val.")
            if record.chrom_start > record.chrom_end:
                raise ValueError("ZoomRecord chrom_start is larger than chrom_end.")
            zoom_records.append(record)
        return cls(zoom_count=zoom_count, zoom_records=zoom_records)


@dataclass
class RTreeNode:
    is_leaf: bool
    children: list[RTreeNode]
    start_chrom_ix: int
    start_base: int
    end_chrom_ix: int
    end_base: int
    data_offset: Optional[int] = None
    data_size: Optional[int] = None

    def get_leaf_nodes(self) -> list[RTreeNode]:
        if self.is_leaf:
            return [self]
        leaf_nodes = []
        for child in self.children:
            leaf_nodes.extend(child.get_leaf_nodes())
        return leaf_nodes

    @classmethod
    def from_file_and_offset(
        cls,
        file_object: BinaryIO,
        start_chrom_ix: int,
        start_base: int,
        end_chrom_ix: int,
        end_base: int,
        offset: Optional[int] = None,
    ) -> RTreeNode:
        current_start_chrom_ix = start_chrom_ix
        current_start_base = start_base
        current_end_chrom_ix = end_chrom_ix
        current_end_base = end_base

        if offset is not None:
            file_object.seek(offset, 0)
        is_leaf, reserved, count = struct.unpack(
            RTREE_NODE_ENCODING, file_object.read(4)
        )

        children = []
        position = file_object.tell()
        for i in range(count):
            file_object.seek(position)
            if is_leaf:
                (
                    start_chrom_ix,
                    start_base,
                    end_chrom_ix,
                    end_base,
                    data_offset,
                    data_size,
                ) = struct.unpack("<LLLLQQ", file_object.read(32))
                position = file_object.tell()

                child_node = RTreeNode(
                    is_leaf=True,
                    children=[],
                    start_chrom_ix=start_chrom_ix,
                    start_base=start_base,
                    end_chrom_ix=end_chrom_ix,
                    end_base=end_base,
                    data_offset=data_offset,
                    data_size=data_size,
                )
            else:
                (
                    start_chrom_ix,
                    start_base,
                    end_chrom_ix,
                    end_base,
                    data_offset,
                ) = struct.unpack("<LLLLQ", file_object.read(24))
                position = file_object.tell()

                child_node = RTreeNode.from_file_and_offset(
                    file_object=file_object,
                    start_chrom_ix=start_chrom_ix,
                    start_base=start_base,
                    end_chrom_ix=end_chrom_ix,
                    end_base=end_base,
                    offset=data_offset,
                )
            children.append(child_node)

        return cls(
            is_leaf=False,
            children=children,
            start_chrom_ix=current_start_chrom_ix,
            start_base=current_start_base,
            end_chrom_ix=current_end_chrom_ix,
            end_base=current_end_base,
            data_offset=None,
            data_size=None,
        )

bigwig_loader/parser.py
pfizer-opensource/bigwig-loader
from typing import Iterator

import numpy as np
import pandas as pd

from bigwig_loader.util import make_cumulative_index_intervals


class PositionSampler:
    def __init__(
        self, regions_of_interest: pd.DataFrame, buffer_size: int = 100000
    ) -> None:
        self.regions_of_interest = make_cumulative_index_intervals(regions_of_interest)
        self.buffer_size = buffer_size
        self._max_index = self.regions_of_interest.index[-1]
        self._chromosomes: list[str] = []
        self._centers: list[int] = []

    def __iter__(self) -> Iterator[tuple[str, int]]:
        return self

    def __next__(self) -> tuple[str, int]:
        if not self._chromosomes:
            self._refresh_buffer()
        return self._chromosomes.pop(), self._centers.pop()

    def _refresh_buffer(self) -> None:
        batch_rand = np.random.randint(
            low=0, high=self._max_index, size=self.buffer_size
        )
        containing_intervals = self.regions_of_interest.iloc[
            self.regions_of_interest.index.searchsorted(batch_rand) - 1  # type: ignore
        ]
        self._centers = list(
            containing_intervals["start"] + (batch_rand - containing_intervals.index)
        )
        self._chromosomes = list(containing_intervals["chrom"])

bigwig_loader/position_sampler.py
pfizer-opensource/bigwig-loader
from pathlib import Path
from typing import Any
from typing import Callable
from typing import Iterator
from typing import Literal
from typing import Optional
from typing import Sequence
from typing import Union

import pandas as pd
import torch
from torch.utils.data import IterableDataset

from bigwig_loader.collection import BigWigCollection
from bigwig_loader.dataset import BigWigDataset


class PytorchBigWigDataset(
    IterableDataset[tuple[torch.FloatTensor, torch.FloatTensor]]
):

    """
    Pytorch IterableDataset over FASTA files and BigWig profiles.
    Args:
        regions_of_interest: pandas dataframe with intervals within the chromosomes to
            sample from. These intervals are not intervals of sequence_length
            but regions of the chromosome that were selected up front. For
            instance regions in which called peaks are located or a threshold
            value is exceeded.
        collection: instance of bigwig_loader.BigWigCollection or path to BigWig
            Directory or list of BigWig files.
        reference_genome_path: path to fasta file containing the reference genome.
        sequence_length: number of base pairs in input sequence
        center_bin_to_predict: if given, only do prediction on a central window. Should be
            smaller than or equal to sequence_length. If not given will be the same as
            sequence_length.
        window_size: used to down sample the resolution of the target from sequence_length
        batch_size: batch size
        super_batch_size: batch size that is used in the background to load data from
            bigwig files. Should be larger than batch_size. If None, it will be equal to
            batch_size.
        batches_per_epoch: because the length of an epoch is slightly arbitrary here,
            the number of batches can be set by hand. If not the number of batches per
            epoch will be (totol number of bases in combined intervals) // sequence_length // batch_size
        maximum_unknown_bases_fraction: maximum number of bases in an input sequence that
            is unknown.
        sequence_encoder: encoder to apply to the sequence. Default: bigwig_loader.util.onehot_sequences
        position_samples_buffer_size: number of intervals picked up front by the position sampler.
            When all intervals are used, new intervals are picked.
    """

    def __init__(
        self,
        regions_of_interest: pd.DataFrame,
        collection: Union[str, Sequence[str], Path, Sequence[Path], BigWigCollection],
        reference_genome_path: Path,
        sequence_length: int = 1000,
        center_bin_to_predict: Optional[int] = 200,
        window_size: int = 1,
        batch_size: int = 256,
        super_batch_size: Optional[int] = None,
        batches_per_epoch: Optional[int] = None,
        maximum_unknown_bases_fraction: float = 0.1,
        sequence_encoder: Optional[
            Union[Callable[[Sequence[str]], Any], Literal["onehot"]]
        ] = "onehot",
        file_extensions: Sequence[str] = (".bigWig", ".bw"),
        crawl: bool = True,
        first_n_files: Optional[int] = None,
        position_sampler_buffer_size: int = 100000,
    ):
        super().__init__()
        self._dataset = BigWigDataset(
            regions_of_interest=regions_of_interest,
            collection=collection,
            reference_genome_path=reference_genome_path,
            sequence_length=sequence_length,
            center_bin_to_predict=center_bin_to_predict,
            batch_size=batch_size,
            super_batch_size=super_batch_size,
            batches_per_epoch=batches_per_epoch,
            maximum_unknown_bases_fraction=maximum_unknown_bases_fraction,
            sequence_encoder=sequence_encoder,
            file_extensions=file_extensions,
            crawl=crawl,
            first_n_files=first_n_files,
            position_sampler_buffer_size=position_sampler_buffer_size,
        )
        if window_size and window_size > 1:
            self.average_pool_target: Optional[torch.nn.AvgPool1d] = torch.nn.AvgPool1d(
                kernel_size=window_size, stride=window_size
            )
        else:
            self.average_pool_target = None

    def __iter__(self) -> Iterator[tuple[torch.FloatTensor, torch.FloatTensor]]:
        iter(self._dataset)
        return self

    def __next__(self) -> tuple[torch.FloatTensor, torch.FloatTensor]:
        sequences, target = next(self._dataset)
        target = torch.as_tensor(target, device="cuda")
        sequences = torch.FloatTensor(sequences)
        if self.average_pool_target:
            target = self.average_pool_target(target)
        return sequences, target

bigwig_loader/pytorch.py
pfizer-opensource/bigwig-loader
import os

from dotenv import find_dotenv
from pydantic.functional_validators import BeforeValidator
from pydantic_settings import BaseSettings
from pydantic_settings import SettingsConfigDict
from typing_extensions import Annotated
from upath import UPath

ENCODE_BASE_URL: str = "https://www.encodeproject.org/"


def string_to_path(v: str) -> UPath:
    return UPath(os.path.expandvars(v))


Path = Annotated[UPath, BeforeValidator(string_to_path)]


class Settings(BaseSettings):
    model_config = SettingsConfigDict(
        env_prefix="bigwig_loader_",
        env_file=find_dotenv(".env"),
        env_file_encoding="utf-8",
        extra="ignore",
    )

    example_data_dir: Path = Path(__file__).parent.parent / "example_data"
    example_positions: Path = example_data_dir / "some_positions.tsv"
    reference_genome_url: str = (
        ENCODE_BASE_URL
        + "files/GRCh38_no_alt_analysis_set_GCA_000001405.15/@@download/GRCh38_no_alt_analysis_set_GCA_000001405.15.fasta.gz"
    )
    reference_genome: Path = (
        example_data_dir / "GRCh38_no_alt_analysis_set_GCA_000001405.15.fasta"
    )
    bigwig_dir: Path = example_data_dir / "bigwig"

    def __str__(self) -> str:
        return super().__str__().replace(" ", "\n")

bigwig_loader/settings.py
pfizer-opensource/bigwig-loader
from pathlib import Path
from typing import Union

import numpy as np
import numpy.typing as npt


class BigWigStore:
    def __init__(
        self,
        fn: Union[str, Path],
        chunk_offsets: npt.NDArray[np.int64],
        chunk_sizes: npt.NDArray[np.int64],
    ) -> None:
        self._filename = fn
        self._fh = open(fn, "rb", buffering=0)
        self.chunk_offsets = chunk_offsets
        self.chunk_sizes = chunk_sizes

    def get_offsets_and_sizes(
        self, keys: npt.NDArray[np.int64]
    ) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
        return self.chunk_offsets[keys], self.chunk_sizes[keys]

bigwig_loader/store.py
pfizer-opensource/bigwig-loader
from typing import Iterable
from typing import Optional
from typing import Sequence

import pandas as pd
from natsort import natsort_keygen
from natsort import natsorted


def get_standard_chromosomes(
    exclude: Optional[Iterable[str]] = ("chrX", "chrY", "chrM")
) -> list[str]:
    exclude = exclude or set()
    chromosomes = [f"chr{i}" for i in range(1, 23)] + ["chrX", "chrY", "chrM"]
    # Doing like this instead of with set to preserve nice order.
    return [chrom for chrom in chromosomes if chrom not in exclude]


def sort_intervals(intervals: pd.DataFrame, inplace: bool = False) -> pd.DataFrame:
    if inplace:
        intervals.sort_values(by=["chrom", "start"], key=natsort_keygen(), inplace=True)
        return intervals
    else:
        return intervals.sort_values(
            by=["chrom", "start"], key=natsort_keygen(), inplace=False
        )


def make_cumulative_index_intervals(intervals: pd.DataFrame) -> pd.DataFrame:
    intervals.reset_index(drop=True, inplace=True)
    intervals.index = (
        (intervals["end"] - intervals["start"]).cumsum().shift().fillna(0).astype(int)  # type: ignore
    )
    return intervals


_string_to_encoding = {
    "A": [1.0, 0.0, 0.0, 0.0],
    "C": [0.0, 1.0, 0.0, 0.0],
    "G": [0.0, 0.0, 1.0, 0.0],
    "T": [0.0, 0.0, 0.0, 1.0],
    "R": [0.5, 0, 0.5, 0],
    "Y": [0, 0.5, 0, 0.5],
    "K": [0, 0, 0.5, 0.5],
    "M": [0.5, 0.5, 0, 0],
    "S": [0, 0.5, 0.5, 0],
    "W": [0.5, 0, 0, 0.5],
    "B": [0, 0.333, 0.333, 0.333],
    "D": [0.333, 0, 0.333, 0.333],
    "H": [0.333, 0.333, 0, 0.333],
    "V": [0.333, 0.333, 0.333, 0],
    "N": [0.25, 0.25, 0.25, 0.25],
}

_standard_bases = {"A", "C", "G", "T"}


def fraction_non_standard(sequence: str) -> float:
    standard_bases = _standard_bases
    n_non_standard = 0
    for base in sequence:
        if base not in standard_bases:
            n_non_standard += 1
    return n_non_standard / len(sequence)


def onehot_sequences(sequences: Sequence[str]) -> list[list[list[float]]]:
    return [[_string_to_encoding[base] for base in sequence] for sequence in sequences]


def chromosome_sort(chromosomes: Iterable[str]) -> list[str]:
    chromosomes = set(chromosomes)
    standard = set(get_standard_chromosomes(exclude=None))
    standard_present = chromosomes.intersection(standard)
    rest = chromosomes - standard_present
    return natsorted(standard_present) + natsorted(rest)  # type: ignore


if __name__ == "__main__":
    onehot = onehot_sequences(["AAAAATTTTACGT", "CAGAATTGTACGT"])
    print(onehot)

bigwig_loader/util.py
pfizer-opensource/bigwig-loader
from bigwig_loader import config
from bigwig_loader.collection import BigWigCollection


def create_train_val_test(
    val_chromosomes=("chr6", "chr7"),
    test_chromosomes=("chr8", "chr9"),
    threshold=2,
    merge_allow_gap=1000,
):
    bigwig_path = config.bigwig_dir
    print(bigwig_path)
    collection = BigWigCollection(bigwig_path)
    print(collection)
    merged_intervals = collection.intervals(
        "standard",
        exclude_chromosomes=val_chromosomes + test_chromosomes,
        merge=True,
        threshold=threshold,
        merge_allow_gap=merge_allow_gap,
    )
    print("Train intervals:", merged_intervals)
    merged_intervals.to_csv("train_regions.tsv", sep="\t", index=None)
    merged_intervals = collection.intervals(
        include_chromosomes=val_chromosomes,
        merge=True,
        threshold=threshold,
        merge_allow_gap=merge_allow_gap,
    )
    print("Validation intervals:", merged_intervals)
    merged_intervals.to_csv("validation_regions.tsv", sep="\t", index=None)
    merged_intervals = collection.intervals(
        include_chromosomes=test_chromosomes,
        merge=True,
        threshold=threshold,
        merge_allow_gap=merge_allow_gap,
    )
    print("Test intervals:", merged_intervals)
    merged_intervals.to_csv("test_regions.tsv", sep="\t", index=None)


if __name__ == "__main__":
    create_train_val_test()

examples/create_train_val_test_intervals.py
pfizer-opensource/bigwig-loader
import pandas as pd

from bigwig_loader import config
from bigwig_loader.dataset import BigWigDataset
from bigwig_loader.download_example_data import download_example_data


def run():
    # Download some data to play with
    download_example_data()

    # created by running examples/create_train_val_test_intervals.py
    train_regions = pd.read_csv("train_regions.tsv", sep="\t")

    # now there is some example data here
    bigwig_dir = config.bigwig_dir
    reference_genome = config.reference_genome
    print("Loading from:", bigwig_dir)

    dataset = BigWigDataset(
        regions_of_interest=train_regions,
        collection=bigwig_dir,
        reference_genome_path=reference_genome,
        sequence_length=1000,
        center_bin_to_predict=1000,
        batch_size=256,
        batches_per_epoch=20,
        maximum_unknown_bases_fraction=0.1,
        sequence_encoder="onehot",
    )

    for input, target in dataset:
        print(input)
        print(target)


if __name__ == "__main__":
    run()

examples/framework_agnostic_example.py
pfizer-opensource/bigwig-loader
import pandas as pd
from torch.utils.data import DataLoader

from bigwig_loader import config
from bigwig_loader.download_example_data import download_example_data
from bigwig_loader.pytorch import PytorchBigWigDataset


def run():
    # Download some data to play with
    download_example_data()

    # created by running examples/create_train_val_test_intervals.py
    train_regions = pd.read_csv("train_regions.tsv", sep="\t")

    # now there is some example data here
    bigwig_dir = config.bigwig_dir
    reference_genome = config.reference_genome
    print("Loading from:", bigwig_dir)

    dataset = PytorchBigWigDataset(
        regions_of_interest=train_regions,
        collection=bigwig_dir,
        reference_genome_path=reference_genome,
        sequence_length=1000,
        center_bin_to_predict=1000,
        window_size=1,
        batch_size=256,
        batches_per_epoch=20,
        maximum_unknown_bases_fraction=0.1,
        sequence_encoder="onehot",
    )

    loader = DataLoader(dataset, batch_size=None, num_workers=0)

    for input, target in loader:
        print(input)
        print(target)


if __name__ == "__main__":
    run()

examples/pytorch_example.py
pfizer-opensource/bigwig-loader
import numpy
from Cython.Build import cythonize
from setuptools import setup
from setuptools.extension import Extension

setup(
    name="bigwig_loader",
    ext_modules=cythonize(
        [
            Extension(
                "bigwig_loader.merge_intervals",
                sources=["bigwig_loader/merge_intervals.pyx"],
                include_dirs=[
                    numpy.get_include(),
                ],
                library_dirs=[],
            ),
        ],
        annotate=True,
    ),
)

setup.py
pfizer-opensource/bigwig-loader
import logging

import pytest

from bigwig_loader import config
from bigwig_loader.collection import BigWigCollection
from bigwig_loader.download_example_data import get_example_bigwigs_files
from bigwig_loader.download_example_data import get_reference_genome

LOGGER = logging.getLogger(__name__)


@pytest.fixture
def example_data_dir():
    return config.example_data_dir


@pytest.fixture
def reference_genome_path():
    return get_reference_genome()


@pytest.fixture(scope="session")
def bigwig_path():
    return get_example_bigwigs_files()


@pytest.fixture
def collection(bigwig_path):
    collection = BigWigCollection(bigwig_path, first_n_files=2)
    return collection


@pytest.fixture
def merged_intervals(collection):
    return collection.intervals(
        ["chr3", "chr4", "chr5"], exclude_chromosomes=["chr4"], merge=True, threshold=2
    )

tests/conftest.py
pfizer-opensource/bigwig-loader
import numpy as np
import pandas as pd
import pyBigWig

from bigwig_loader import config
from bigwig_loader.collection import BigWigCollection
from bigwig_loader.collection import interpret_path


class PyBigWigCollection:
    def __init__(self, bigwig_path, first_n_files=None):
        bigwig_paths = sorted(interpret_path(bigwig_path))
        self.bigwigs = [pyBigWig.open(str(bw)) for bw in bigwig_paths][:first_n_files]

    def get_batch(self, chromosomes, starts, ends):
        out = np.stack(
            [
                bw.values(chrom, start, end, numpy=True)
                for chrom, start, end in zip(chromosomes, starts, ends)
                for bw in self.bigwigs
            ]
        )
        return out.reshape(
            (out.shape[0] // len(self.bigwigs), len(self.bigwigs), out.shape[1])
        )


def test_same_output(bigwig_path):
    pybigwig_collection = PyBigWigCollection(bigwig_path, first_n_files=2)
    collection = BigWigCollection(bigwig_path, first_n_files=2)

    df = pd.read_csv(config.example_positions, sep="\t")
    df = df[df["chr"].isin(collection.get_chromosomes_present_in_all_files())]
    chromosomes, starts, ends = (
        list(df["chr"]),
        list(df["center"] - 500),
        list(df["center"] + 500),
    )

    pybigwig_batch = pybigwig_collection.get_batch(chromosomes, starts, ends)
    np.nan_to_num(pybigwig_batch, copy=False, nan=0.0)
    this_batch = collection.get_batch(chromosomes, starts, ends).get()
    print("PyBigWig:")
    print(pybigwig_batch)
    print(type(this_batch), "shape:", pybigwig_batch.shape)
    print("This Library:")
    print(this_batch)
    print(type(this_batch), "shape:", this_batch.shape)
    print(this_batch[pybigwig_batch != this_batch])
    print(pybigwig_batch[pybigwig_batch != this_batch])
    assert (pybigwig_batch == this_batch).all()

tests/test_against_pybigwig.py
pfizer-opensource/bigwig-loader
import pandas as pd
import pytest

from bigwig_loader import config
from bigwig_loader.collection import BigWigCollection
from bigwig_loader.dataset import BigWigDataset
from bigwig_loader.position_sampler import PositionSampler
from bigwig_loader.util import sort_intervals


def _get_example_batch_of_intervals(n_bases=1000, batch_size=256):
    batch_of_positions = pd.read_csv(
        config.example_data_dir / "some_positions.tsv", sep="\t"
    )

    batch_of_positions = batch_of_positions.sample(batch_size)

    center = batch_of_positions["center"]
    batch_of_intervals = pd.DataFrame(
        {"chrom": batch_of_positions["chr"], "start": center - (n_bases // 2)}
    )
    batch_of_intervals["end"] = batch_of_intervals["start"] + n_bases
    return batch_of_intervals


@pytest.fixture
def bigwig_dataset(bigwig_path, reference_genome_path, merged_intervals):
    dataset = BigWigDataset(
        regions_of_interest=merged_intervals,
        collection=bigwig_path,
        reference_genome_path=reference_genome_path,
        sequence_length=1000,
        center_bin_to_predict=1000,
        batch_size=256,
        batches_per_epoch=10,
        maximum_unknown_bases_fraction=0.1,
        first_n_files=2,
    )
    return dataset


def test_dataset(bigwig_dataset):
    for sequence, target in bigwig_dataset:
        assert target.shape == (256, 2, 1000)


def test_get_batch(collection):
    intervals = _get_example_batch_of_intervals(n_bases=1000, batch_size=256)
    n_files = len(collection.bigwig_paths)
    batch = collection.get_batch(
        intervals["chrom"].values,
        intervals["start"].values,
        intervals["end"].values,
    )
    assert batch.shape == (256, n_files, 1000)


def test_exclude_intervals(collection):
    intervals = collection.intervals(
        ["chr3", "chr4", "chr5"], exclude_chromosomes=["chr4"]
    )
    present_chromosomes = set(intervals["chrom"])
    assert "chr5" in present_chromosomes and "chr4" not in present_chromosomes


def test_merge_intervals(collection, merged_intervals):
    unmerged_intervals = collection.intervals(
        ["chr3", "chr4", "chr5"], exclude_chromosomes=["chr4"], merge=False, threshold=2
    )
    assert len(merged_intervals) < len(unmerged_intervals)


def test_sort_intervals(example_data_dir):
    intervals = pd.read_csv(example_data_dir / "some_intervals.tsv", sep="\t")
    intervals.rename(columns={"chr": "chrom"}, inplace=True)
    df = sort_intervals(intervals)
    df.reset_index(drop=True, inplace=True)
    assert df[df["chrom"] == "chr2"].index[0] < df[df["chrom"] == "chr10"].index[0]


def test_position_sampler(merged_intervals):
    sampler = PositionSampler(merged_intervals)
    sample = next(sampler)
    assert sample[0].startswith("chr")


def test_chromosomes_present_in_all_files(collection: BigWigCollection):
    chromosomes = collection.get_chromosomes_present_in_all_files("standard")
    assert len(chromosomes) > 0
    assert all(isinstance(c, str) and c.strip() for c in chromosomes)

tests/test_collection.py
pfizer-opensource/bigwig-loader
import pytest

from bigwig_loader.collection import interpret_path


@pytest.fixture(scope="session")
def directory_structure(tmp_path_factory):
    path = tmp_path_factory.mktemp("root_one")
    sub_dir_one = path / "one"
    sub_dir_two = path / "two"
    sub_dir_one.mkdir()
    sub_dir_two.mkdir()
    (path / "heart.bw").write_text("")
    (sub_dir_one / "kidney.bw").write_text("")
    (sub_dir_one / "lung.bw").write_text("")
    (sub_dir_two / "eye.bigWig").write_text("")
    (sub_dir_two / "misspelled.biggwigg").write_text("")
    return path


def test_interpret_path_string(directory_structure):
    assert set(interpret_path(str(directory_structure))) == {
        directory_structure / "heart.bw",
        directory_structure / "one" / "kidney.bw",
        directory_structure / "one" / "lung.bw",
        directory_structure / "two" / "eye.bigWig",
    }


def test_interpret_path(directory_structure):
    assert set(interpret_path(directory_structure)) == {
        directory_structure / "heart.bw",
        directory_structure / "one" / "kidney.bw",
        directory_structure / "one" / "lung.bw",
        directory_structure / "two" / "eye.bigWig",
    }


def test_interpret_path_no_walk(directory_structure):
    assert interpret_path(directory_structure, crawl=False) == [
        directory_structure / "heart.bw"
    ]


def test_interpret_path_custom_extension(directory_structure):
    assert interpret_path(directory_structure, file_extensions=(".biggwigg",)) == [
        directory_structure / "two" / "misspelled.biggwigg",
    ]


def test_interpret_sequence_of_pathlib(directory_structure):
    assert set(
        interpret_path(
            [
                directory_structure / "heart.bw",
                directory_structure / "one",
                directory_structure / "two",
            ]
        )
    ) == {
        directory_structure / "heart.bw",
        directory_structure / "one" / "kidney.bw",
        directory_structure / "one" / "lung.bw",
        directory_structure / "two" / "eye.bigWig",
    }


def test_interpret_sequence_of_pathlib_2(directory_structure):
    assert set(
        interpret_path([directory_structure / "one", directory_structure / "two"])
    ) == {
        directory_structure / "one" / "kidney.bw",
        directory_structure / "one" / "lung.bw",
        directory_structure / "two" / "eye.bigWig",
    }

tests/test_collection_file_path_interpretation.py
pfizer-opensource/bigwig-loader
import cupy as cp

from bigwig_loader.intervals_to_values_gpu import intervals_to_values


def test_get_values_from_intervals() -> None:
    """Probably most frequent situation."""
    track_starts = cp.asarray([1, 3, 10, 12, 16], dtype=cp.int32)
    track_ends = cp.asarray([3, 10, 12, 16, 20], dtype=cp.int32)
    track_values = cp.asarray([20.0, 15.0, 30.0, 40.0, 50.0], dtype=cp.dtype("f4"))
    query_starts = cp.asarray([2], dtype=cp.int32)
    query_ends = cp.asarray([17], dtype=cp.int32)
    reserved = cp.zeros((1, 15), dtype=cp.float32)
    values = intervals_to_values(
        track_starts, track_ends, track_values, query_starts, query_ends, reserved
    )
    assert (
        values
        == cp.asarray([[20, 15, 15, 15, 15, 15, 15, 15, 30, 30, 40, 40, 40, 40, 50]])
    ).all()


def test_get_values_from_intervals_edge_case_1() -> None:
    """Query start is somewhere in a "gap"."""
    track_starts = cp.asarray([1, 10, 12, 16], dtype=cp.int32)
    track_ends = cp.asarray([3, 12, 16, 20], dtype=cp.int32)
    track_values = cp.asarray([20.0, 30.0, 40.0, 50.0], dtype=cp.dtype("f4"))
    query_starts = cp.asarray([6], dtype=cp.int32)
    query_ends = cp.asarray([18], dtype=cp.int32)
    reserved = cp.zeros((1, 12), dtype=cp.dtype("<f4"))
    values = intervals_to_values(
        track_starts, track_ends, track_values, query_starts, query_ends, reserved
    )
    expected = cp.asarray([[0, 0, 0, 0, 30, 30, 40, 40, 40, 40, 50, 50]])
    assert (values == expected).all() and expected.shape[-1] == query_ends[
        0
    ] - query_starts[0]


def test_get_values_from_intervals_edge_case_2() -> None:
    """Query start is exactly at start index after "gap"."""
    track_starts = cp.asarray([1, 10, 12, 16], dtype=cp.int32)
    track_ends = cp.asarray([3, 12, 16, 20], dtype=cp.int32)
    track_values = cp.asarray([20.0, 30.0, 40.0, 50.0], dtype=cp.dtype("f4"))
    query_starts = cp.asarray([10], dtype=cp.int32)
    query_ends = cp.asarray([18], dtype=cp.int32)
    reserved = cp.zeros((1, 8), dtype=cp.dtype("<f4"))
    values = intervals_to_values(
        track_starts, track_ends, track_values, query_starts, query_ends, reserved
    )
    expected = cp.asarray([[30, 30, 40, 40, 40, 40, 50, 50]])
    assert (values == expected).all() and expected.shape[-1] == query_ends[
        0
    ] - query_starts[0]


def test_get_values_from_intervals_edge_case_3() -> None:
    """Query end is somewhere in a "gap"."""
    track_starts = cp.asarray([5, 10, 12, 18], dtype=cp.int32)
    track_ends = cp.asarray([10, 12, 14, 20], dtype=cp.int32)
    track_values = cp.asarray([20.0, 30.0, 40.0, 50.0], dtype=cp.dtype("f4"))
    query_starts = cp.asarray([9], dtype=cp.int32)
    query_ends = cp.asarray([16], dtype=cp.int32)
    reserved = cp.zeros((1, 7), dtype=cp.dtype("<f4"))
    values = intervals_to_values(
        track_starts, track_ends, track_values, query_starts, query_ends, reserved
    )
    expected = cp.asarray([[20, 30, 30, 40, 40, 0, 0]])
    assert (values == expected).all() and expected.shape[-1] == query_ends[
        0
    ] - query_starts[0]


def test_get_values_from_intervals_edge_case_4() -> None:
    """Query end is exactly at end index before "gap"."""
    track_starts = cp.asarray([5, 10, 12, 18], dtype=cp.int32)
    track_ends = cp.asarray([10, 12, 14, 20], dtype=cp.int32)
    track_values = cp.asarray([20.0, 30.0, 40.0, 50.0], dtype=cp.dtype("f4"))
    query_starts = cp.asarray([9], dtype=cp.int32)
    query_ends = cp.asarray([14], dtype=cp.int32)
    reserved = cp.zeros((1, 5), dtype=cp.dtype("<f4"))
    values = intervals_to_values(
        track_starts, track_ends, track_values, query_starts, query_ends, reserved
    )
    expected = cp.asarray([[20, 30, 30, 40, 40]])
    assert (values == expected).all() and expected.shape[-1] == query_ends[
        0
    ] - query_starts[0]


def test_get_values_from_intervals_edge_case_5() -> None:
    """Query end is exactly at end index before "gap"."""
    track_starts = cp.asarray([5, 10, 12, 18], dtype=cp.uint32)
    track_ends = cp.asarray([10, 12, 14, 20], dtype=cp.uint32)
    track_values = cp.asarray([20.0, 30.0, 40.0, 50.0], dtype=cp.dtype("f4"))
    query_starts = cp.asarray([9], dtype=cp.uint32)
    query_ends = cp.asarray([20], dtype=cp.uint32)
    reserved = cp.zeros((1, 11), dtype=cp.dtype("<f4"))
    values = intervals_to_values(
        track_starts, track_ends, track_values, query_starts, query_ends, reserved
    )
    expected = cp.asarray([[20, 30, 30, 40, 40, 0, 0, 0, 0, 50, 50]])
    assert (values == expected).all() and expected.shape[-1] == query_ends[
        0
    ] - query_starts[0]


def test_get_values_from_intervals_batch_of_2() -> None:
    """Query end is exactly at end index before "gap"."""
    track_starts = cp.asarray([5, 10, 12, 18], dtype=cp.int32)
    track_ends = cp.asarray([10, 12, 14, 20], dtype=cp.int32)
    track_values = cp.asarray([20.0, 30.0, 40.0, 50.0], dtype=cp.dtype("f4"))
    query_starts = cp.asarray([7, 9], dtype=cp.int32)
    query_ends = cp.asarray([18, 20], dtype=cp.int32)
    reserved = cp.zeros([2, 11], dtype=cp.dtype("<f4"))
    values = intervals_to_values(
        track_starts, track_ends, track_values, query_starts, query_ends, reserved
    )
    expected = cp.asarray(
        [
            [20.0, 20.0, 20.0, 30.0, 30.0, 40.0, 40.0, 0.0, 0.0, 0.0, 0.0],
            [20.0, 30.0, 30.0, 40.0, 40.0, 0.0, 0.0, 0.0, 0.0, 50.0, 50.0],
        ]
    )
    print(expected)
    print(values)
    assert (values == expected).all()

tests/test_intervals_to_values_gpu.py
pfizer-opensource/bigwig-loader
import numpy as np
import pandas as pd

from bigwig_loader.merge_intervals import merge_interval_dataframe
from bigwig_loader.merge_intervals import merge_intervals
from bigwig_loader.merge_intervals import merge_intervals_with_chrom_keys


def test_merge_intervals() -> None:
    chrom_id = np.array([1, 1, 1, 1], np.dtype("uint32"))
    start = np.array([1, 4, 12, 16], np.dtype("uint32"))
    end = np.array([4, 9, 16, 19], np.dtype("uint32"))
    value = np.array([1.0, 2.0, 3.0, 4.0], np.dtype("f4"))
    out_chrom_id, out_start, out_end, out_value = merge_intervals(
        chrom_id, start, end, value, True, 0
    )
    assert list(out_start) == [1, 12] and list(out_end) == [9, 19]


def test_merge_interval_2() -> None:
    chrom_id = np.array([1, 1, 1, 1, 1], np.dtype("uint32"))
    start = np.array([1, 4, 12, 16, 40], np.dtype("uint32"))
    end = np.array([4, 9, 16, 19, 45], np.dtype("uint32"))
    value = np.array([1.0, 2.0, 3.0, 4.0, 5.0], np.dtype("f4"))
    out_chrom_id, out_start, out_end, out_value = merge_intervals(
        chrom_id, start, end, value, True, 0
    )
    assert list(out_start) == [1, 12, 40] and list(out_end) == [9, 19, 45]


def test_merge_interval_3() -> None:
    chrom_id = np.array([1, 1, 1, 1, 1], np.dtype("uint32"))
    start = np.array([1, 4, 12, 16, 40], np.dtype("uint32"))
    end = np.array([6, 9, 16, 19, 45], np.dtype("uint32"))
    value = np.array([1.0, 2.0, 3.0, 4.0, 5.0], np.dtype("f4"))
    out_chrom_id, out_start, out_end, out_value = merge_intervals(
        chrom_id, start, end, value, True, 0
    )
    assert list(out_start) == [1, 12, 40] and list(out_end) == [9, 19, 45]


def test_merge_interval_4() -> None:
    chrom_id = np.array([1, 2, 3, 3, 3], np.dtype("uint32"))
    start = np.array([1, 4, 12, 16, 40], np.dtype("uint32"))
    end = np.array([6, 9, 16, 19, 45], np.dtype("uint32"))
    value = np.array([1.0, 2.0, 3.0, 4.0, 5.0], np.dtype("f4"))
    out_chrom_id, out_start, out_end, out_value = merge_intervals(
        chrom_id, start, end, value, True, 0
    )
    assert (
        list(out_chrom_id) == [1, 2, 3, 3]
        and list(out_start) == [1, 4, 12, 40]
        and list(out_end) == [6, 9, 19, 45]
    )


def test_merge_interval_with_chrom_keys() -> None:
    chrom_id = np.array(["chr2", "chr4", "chr12", "chr12", "chr12"], dtype=object)
    start = np.array([1, 4, 12, 16, 40], np.dtype("uint32"))
    end = np.array([6, 9, 16, 19, 45], np.dtype("uint32"))
    value = np.array([1.0, 2.0, 3.0, 4.0, 5.0], np.dtype("f4"))
    out_chrom_id, out_start, out_end, out_value = merge_intervals_with_chrom_keys(
        chrom_id, start, end, value, False, 0
    )
    assert (
        list(out_chrom_id) == ["chr2", "chr4", "chr12", "chr12"]
        and list(out_start) == [1, 4, 12, 40]
        and list(out_end) == [6, 9, 19, 45]
    )


def test_merge_intervals_dataframe() -> None:
    chrom = np.array(["chr2", "chr4", "chr12", "chr12", "chr12"])
    start = np.array([1, 4, 12, 16, 40], np.dtype("uint32"))
    end = np.array([6, 9, 16, 19, 45], np.dtype("uint32"))
    value = np.array([1.0, 2.0, 3.0, 4.0, 5.0], np.dtype("f4"))

    df = pd.DataFrame({"chrom": chrom, "start": start, "end": end, "value": value})
    df_shuffled = df.sample(frac=1)

    merged = merge_interval_dataframe(df_shuffled, False)
    assert list(merged["chrom"]) == ["chr2", "chr4", "chr12", "chr12"]
    assert list(merged["start"]) == [1, 4, 12, 40]
    assert list(merged["end"]) == [6, 9, 19, 45]


def test_merge_interval_with_allow_gap() -> None:
    chrom_id = np.array([1, 2, 3, 3, 3], np.dtype("uint32"))
    start = np.array([1, 4, 12, 20, 29], np.dtype("uint32"))
    end = np.array([6, 9, 16, 25, 32], np.dtype("uint32"))
    value = np.array([1.0, 2.0, 3.0, 4.0, 5.0], np.dtype("f4"))
    out_chrom_id, out_start, out_end, out_value = merge_intervals(
        chrom_id, start, end, value, True, 4
    )
    assert (
        list(out_chrom_id) == [1, 2, 3]
        and list(out_start) == [1, 4, 12]
        and list(out_end) == [6, 9, 32]
    )

tests/test_merge_intervals.py
pfizer-opensource/bigwig-loader
import pytest

torch = pytest.importorskip("torch")


@pytest.fixture
def pytorch_dataset(bigwig_path, reference_genome_path, merged_intervals):
    from bigwig_loader.pytorch import PytorchBigWigDataset

    dataset = PytorchBigWigDataset(
        regions_of_interest=merged_intervals,
        collection=bigwig_path,
        reference_genome_path=reference_genome_path,
        sequence_length=1000,
        center_bin_to_predict=1000,
        batch_size=256,
        batches_per_epoch=4,
        maximum_unknown_bases_fraction=0.1,
        first_n_files=2,
    )
    return dataset


def test_pytorch_dataset(pytorch_dataset):
    for sequence, target in pytorch_dataset:
        assert target.shape == (256, 5, 1000)


def test_input_and_target_is_torch_tensor(pytorch_dataset):
    sequence, target = next(iter(pytorch_dataset))
    assert isinstance(sequence, torch.Tensor)
    assert isinstance(target, torch.Tensor)

tests/test_pytorch_dataset.py
pfizer-opensource/bigwig-loader