View on GitHub


1 day
Test Coverage
#!/user/bin/env python
# -*- coding: utf-8 -*-
"""CSI file parser

Copyright |copy| 2018, Marcus D. Sherman

This code is part of the bamnostic distribution and governed by its
license.  Please see the LICENSE file that should have been included
as part of this package.

The Binary Alignment Map (CSI) format (defined at
allows for indexing. When the user invokes a tool to index a given BAM file, a CSI index file
can be created. Generally speaking, a CSI contains the all the virtual offsets of clusters of
reads that are associated with specific subsections of the BAM file. When the CSI is produced,
random access to the BAM is available.

This script is for parsing the binary encoded CSI file, inherently making it human readable.
Furthermore, it allows subsections of the CSI to be loaded into memory, reducing the memory
footprint and speeding up queries within a small number of references. Lastly, by parsing it
as such, random access queries directly into the associated BAM file is available to other
tools within bamnostic


from __future__ import print_function
from __future__ import absolute_import
from __future__ import division

import struct
import os
import sys
import warnings
import gzip

_PY_VERSION = sys.version_info

if _PY_VERSION[0] == 2:
    from io import open
    from functools import lru_cache

import bamnostic
from bamnostic import bai
from bamnostic.utils import *

def format_warnings(message, category, filename, lineno, file=None, line=None):
    return ' {}:{}: {}:{}'.format(filename, lineno, category.__name__, message)

warnings.formatwarning = bai.format_warnings

# Helper compiled structs
unpack_chunk = bai.unpack_chunk
unpack_bid_loffset_nchunk = struct.Struct('<IQi').unpack
unpack_unmapped = bai.unpack_unmapped

class CsiBin(object):
    __slots__ = ['loffset', 'chunks']

    def __init__(self, *args):
        self.loffset, self.chunks = args

    def __repr__(self):
        return 'CsiBin(loffset={}, chunks={})'.format(self.loffset, self.chunks)

class RefCsi(object):
    __slots__ = ['bins', 'ref_id']

    def __init__(self, *args):
        self.bins, self.ref_id = args

    def __repr__(self):
        return 'RefCsi(bins={}, ref_id={})'.format(self.bins, self.ref_id)

    def __getitem__(self, key):
        return self.bins[key]

def reg2bin_csi(rbeg, rend, min_shift=14, depth=5):
    """Finds the largest superset bin of region. Numeric values taken from hts-specs

        rbeg (int): inclusive beginning position of region
        rend (int): exclusive end position of region

        (int): distinct bin ID for largest superset bin of region

    assert 0 <= rbeg <= rend, 'Invalid region {}, {}'.format(rbeg, rend)
    total = ((1 << depth * 3) - 1) // 7
    stop -= 1
    while depth > 0:
        if start >> min_shift == stop >> min_shift:
            return total + (start >> min_shift)
        min_shift += 3
        total -= 1 << (depth * 3)
        depth -= 1
    return 0

def reg2bins_csi(rbeg, rend, min_shift=14, depth=5):
    """Generates bin ids which overlap the specified region.

        rbeg (int): inclusive beginning position of region
        rend (int): exclusive end position of region

        (int): bin IDs for overlapping bins of region

        AssertionError (Exception): if the range is malformed or invalid
    # Based off the algorithm presented in:

    assert 0 <= rbeg <= rend, 'Invalid region {}, {}'.format(rbeg, rend)
    bins = []
    shift = min_shift + depth * 3
    rend -= 1
    level = total = 0
    while level <= depth:
        beg = total + (rbeg >> shift)
        end = total + (rend >> shift)
        i = beg
        while i <= end:
            yield i
            i += 1
        shift -= 3
        total += 1 << level * 3
        level += 1

class Csi(bai.Bai):
    """ This class defines the bam CSI index file object and its interface.

    The purpose of this class is the binary parsing of the bam index file (CSI) associated with
    a given bam file. When queried, the Csi object identifies the bins of data that overlap the
    requested region and directs which parts of the bam file contain it.

    Virtual offsets are processed using the following method:
    Beginning of compressed block = coffset = virtual offset >> 16
    Position within uncompressed block = uoffset = virtual offset ^ (coffset << 16)

        _io (fileObject): opened CSI file object
        _UNMAP_BIN (int): constant for bin ID of unmapped read stats
        magic (bytes): first 4 bytes of file. Must be equal to b'CSI\x01'
        n_refs (int): number of references in CSI
        unmapped (dict): dictionary of the unmapped read stats by each reference
        current_ref (None|dict): dictionary of the current reference loaded into memory.
                                It contains the a dictionary of bin IDs and their respective
                                chunks, and a list of linear intervals.
        ref_indices (dict): dictionary of reference ids and their start/stop offsets within theCSI file
        n_no_coord (None|int): if present inCSI, is the number of reads that have no coordinates
        _last_pos (int): used for indexing, the byte position of the file head.

    __slots__ = ['_io', '_min_shift', '_UNMAP_BIN', '_magic', '_depth', 'aux', 
                'current_ref', 'ref_indices', 'n_no_coor', '_last_pos', 'n_refs']

    def __init__(self, filename):
        """Initialization method

        Generates an "index" of the index. This gives us the byte positions of each chromosome
        within the index file. Now, when a user queries over a specific chromosome, it pulls
        out just the index information for that chromosome--not the whole genome.

            filename (str): '/path/to/bam_file' that automatically adds the '.csi' suffix

            OSError (Exception): if the CSI file is not found or does not exist
            AssertionError (Exception): if CSI magic is not found
        if os.path.isfile(filename):
            self._io = gzip.GzipFile(filename)
            raise OSError('{} not found. Please change check your path or index your BAM file'.format(filename))

        # Constant for linear index window size and unmapped bin id
        self._UNMAP_BIN = 37450

        self._magic, self._min_shift, self._depth, l_aux = unpack("<4s3i", self._io)
        assert self._magic == b'CSI\x01', 'Wrong CSI magic header'

        aux_nref = unpack('<{}Bi'.format(l_aux), self._io)
        if l_aux > 0:
            self.aux = aux_nref[:-1]
            self.n_refs = aux_nref[-1]
        elif l_aux == 0:
            self.n_refs = aux_nref
            raise IOError('CSI attribute l_aux must be >= 0')

        self.unmapped = {}
        self.current_ref = None

        # Capture the offsets for each reference within the index
        self.ref_indices = {ref: self.get_ref(ref, idx=True) for ref in range(self.n_refs)}

        # Get the n_no_coor if it is present
        nnc_dat =
        self.n_no_coor = unpack('<Q', nnc_dat) if nnc_dat else None

        self._last_pos = self._io.tell()

    def get_chunks(self, n_chunks):
        """Simple generator for unpacking chunk data

        Chunks are defined as groupings of reads within a BAM file
        that share the same bin. A `Chunk` object in the context of
        this function is a `namedtuple` that contains the virtual offsets
        for the beginning and end of each of these chunks.

        Note: a special case of a chunk is in any Bin labeled as 37450.
        These bins always contain 2 chunks that provide the statistics
        of the number of reads that are unmapped to that reference.

            n_chunks (int): number of chunks to be unpacked from stream

            chunks (list): a list of Chunk objects with the attributes of
                            chunks[i] are .voffset_beg and voffset_end
        chunks = [bai.Chunk(self._io) for chunk in range(n_chunks)]
        return chunks

    def get_bins(self, n_bins, ref_id=None, idx=False):
        """Simple function that iteratively unpacks the data of a number (`n_bin`)
        of bins.

        As the function discovers the number of chunks needed for a given
        bin, it deletages work to `self.get_chunks(n_chunks)`. A bin is
        comprised of 2 parts: 1) the distinct bin ID (within a given reference). If
        no reads are associated with that bin, it is left out of the indexing process,
        and therefore not represented in theCSI file. Furthermore, while each bin
        has a bin ID, the bin IDs are only distinct within a given reference. This means
        that 2 or more references can have the same bin IDs. These bin IDs are also
        not in any order as they are essentially a hash dump. Lastly, the only reserved
        bin ID is 37450. This bin relates to 2 chunks that contain the number of
        unmapped and mapped reads for a given reference. 2) the chunk(s) of reads that
        are assigned to a given bin.

        As a secondary feature, this function will also quickly seek over regions for
        the purposes of documenting the start and stop byte offsets of a given reference
        block within the file. This is invoked by setting `idx=True`

            n_int (int): number of bins to be unpacked from stream

            bins (None | dict): None if just indexing the index file or a dictionary
                                of `bin_id: chunks` pairs
            AssertionError (Exception): if bin 37450 does not contain 2 chunks exactly
        bins = None if idx else {}

        for b in range(n_bins):
            bin_id, loffset, n_chunks = unpack_bid_loffset_nchunk(
            if idx:
                if bin_id == self._UNMAP_BIN:
                    assert n_chunks == 2, 'Bin 3740 is supposed to have 2 chunks. This has {}'.format(n_chunks)
                    unmapped = bai.Unmapped(*unpack_unmapped(
                    self.unmapped[ref_id] = unmapped
                    if not n_chunks == 0:
               * n_chunks, 1)
                chunks = self.get_chunks(n_chunks)
                bins[bin_id] = CsiBin(loffset, chunks)
            return bins

    # Cache the references to speed up queries.

    # @functools.lru_cache(maxsize=256, typed=True)
    # @lru_cache(6)
    @bai.conditional_decorator(lambda func: lru_cache(maxsize=6)(func), _PY_VERSION[0] == 2)
    def get_ref(self, ref_id=None, idx=False):
        """Iteratively unpacks all the bins, linear intervals, and chunks for a given reference

        A reference is comprised of 2 things: 1) a series of bins that reference chunks of aligned
        reads that are grouped within that bin. 2) a series of virtual offsets of the first read of a
        16384 bp window along the given reference.

        This function also serves to "index" theCSI file such that, if it is invoked by setting
        `ids=True`, will do a single pass through theCSI file and saving the start and stop
        offsets of each of the references. This is used for minimizing the memory footprint of
        storing theCSI in memory. When queried against, the appropriate reference block will be
        loaded. Because of this constant loading, `functools.lru_cache` was applied to cache recently
        used reference blocks to speed up computation. It is assumed that when querying is done, most
        users are looking and just a few references at a time.

            ref_id (None|int): used for random access or indexing theCSI
            idx (bool): Flag for setting whether or not to run an index of theCSI

            RefIdx: `namedtuple` containing the byte offsets of the reference start, stop, and number of bins
            Ref: `namedtuple` containing a dictionary of bins and list of linear intervals

            AssertionError (Exception): if, when random access is used, the current reference offset
                                        does not match indexed reference offset.

        if ref_id is not None and not idx:
                ref_start, _, _ = self.ref_indices[ref_id]
            except KeyError:
                raise KeyError('Reference is not found in header')
        ref_start = self._io.tell()

        if not idx:
            assert ref_start == self.ref_indices[ref_id].start_offset, 'ref not properly aligned'

        n_bins = unpack_int32L([0]
        bins = self.get_bins(n_bins, ref_id, idx)

        self._last_pos = self._io.tell()

        if idx:
            return bai.RefIdx(ref_start, self._last_pos, n_bins)
            return RefCsi(bins, ref_id)

    def query(self, ref_id, start, stop=-1):
        """ Main query function for determining seek offset to BAM section that
        AlignedRead objects from specified region start

            ref (int): which reference/chromosome TID
            start (int): left most bp position of region (zero-based)
            stop (int): right most bp position of region (zero-based)

            (int): the voffset_beg of the first chunk given the chunk's voffset_end
                    is greater than the voffset of the linear index that overlaps
                    the region of interest's start offset
        if stop < 0:
            end_offset = self.current_ref.intervals[-1] + 1
        assert start <= stop, 'Malformed region: start should be <= stop, you entered {}, {}'.format(start, stop)

        if self.current_ref is None:
            self.current_ref = self.get_ref(ref_id)
        elif self.current_ref.ref_id != ref_id:
            self.current_ref = self.get_ref(ref_id)

        for potential_bin in reg2bins_csi(start, stop, min_shift=self._min_shift, depth=self._depth):
                pbin = self.current_ref[potential_bin]
            except KeyError:

            reg_lin_idx = start >> self._min_shift # BAI-style linear index should match with default loffset
            for chunk in pbin.chunks:
                if pbin.loffset <= chunk.voffset_end:
                    return chunk.voffset_beg