panxpress 0.2__py3-none-any.whl

panxpress/cuckoo_filter_utils.py

@@ -0,0 +1,207 @@
+"""
+Utilities, pointer and bit magic for cuckoo filters
+"""
+
+from math import log2
+from numba import njit, uint8, uint16, uint32, uint64
+import numba as nb
+import numpy as np
+from llvmlite import ir
+from .mathutils import bitsfor
+from .lowlevel.llvm import compile_cttz
+
+
+def compute_masks(bits_per_slot, slots):
+    mask1 = ('0' * (bits_per_slot - 1) + '1') * slots
+    mask2 = ('1' + '0' * (bits_per_slot - 1)) * slots
+    return uint64(int(mask1, 2)), uint64(int(mask2, 2))
+
+
+def compute_fp_mask(bits_per_slot, fingerprint_bits, slots):
+    fp_mask = uint64(0)
+    for s in range(1, slots):
+        fp_mask |= (s << (s * bits_per_slot + fingerprint_bits))
+    return uint64(fp_mask)
+
+
+def compute_choice_fp_mask(bits_per_slot, fingerprint_bits, slots):
+    window_bits = bitsfor(slots)
+    fp_mask = uint64(0)
+    for s in range(0, slots):
+        fp_mask |= (s << (s * bits_per_slot + fingerprint_bits))
+        fp_mask |= 1 << (s * bits_per_slot + fingerprint_bits + window_bits)
+    for s in range(0, slots):
+        fp_mask |= (s << ((slots + s) * bits_per_slot + fingerprint_bits))
+    return uint64(fp_mask)
+
+
+def compute_choice_mask(bits_per_slot, fingerprint_bits, slots):
+    fp_mask = uint64(0)
+    for s in range(0, slots):
+        fp_mask |= 1 << (s * bits_per_slot + fingerprint_bits)
+    return uint64(fp_mask)
+
+
+def compile_lookup(bits_per_slot, slots, windowed=True, choice=False):
+    fingerprint_bits = uint64(bits_per_slot - int(log2(slots)) - 1 if windowed else bits_per_slot - 1)
+    if choice:
+        assert bits_per_slot * slots * 2 <= 64
+        m1, m2 = compute_masks(bits_per_slot, 2 * slots)
+        if windowed:
+            fp_mask = compute_choice_fp_mask(bits_per_slot, fingerprint_bits, slots)
+        else:
+            fp_mask = compute_choice_mask(bits_per_slot, fingerprint_bits, slots)
+    else:
+        assert bits_per_slot * slots <= 64
+        m1, m2 = compute_masks(bits_per_slot, slots)
+        fp_mask = compute_fp_mask(bits_per_slot, fingerprint_bits, slots) if windowed else uint64(0)
+
+    @njit(nogil=True)
+    def haszero(x):
+        return uint64(((x) - m1) & (~(x)) & m2)
+
+    @njit(nogil=True)
+    def hasvalue(window, fp):
+        x = haszero(window ^ ((m1 * fp) | fp_mask))
+        return x != 0
+
+    return hasvalue
+
+
+def compile_get_empty_slot(bits_per_slot, slots):
+    assert bits_per_slot * slots <= 64
+    count_trailing_zeros = compile_cttz('uint64')
+    m1, m2 = compute_masks(bits_per_slot, slots)
+
+    @njit(nogil=True)
+    def haszero(x):
+        return (((x) - m1) & (~(x)) & m2)
+
+    @njit(nogil=True, locals=dict(x=uint64, slot=uint64))
+    def get_empty_slot(window):
+        x = haszero(window)
+        slot = count_trailing_zeros(x) // bits_per_slot
+        return slot
+
+    return get_empty_slot
+
+
+def compile_load_value(nbits):
+    # fast version for multiples of 8 (cast byte counter)
+    if nbits in [8, 16, 32, 64]:
+        signatures = {8: uint8, 16: uint16, 32: uint32, 64: uint64}
+        pointer = ir.IntType(nbits).as_pointer()
+        int_type = signatures[nbits]
+
+        @nb.extending.intrinsic
+        def address_to_value(typingctx, src):
+            """ returns the value stored at a given memory address """
+            sig = int_type(src)
+            def codegen(cgctx, builder, sig, args):
+                ptr = builder.inttoptr(args[0], pointer)
+                return builder.load(ptr)
+            return sig, codegen
+
+        @njit(nogil=True, locals=dict(address=uint64))
+        def get_value(fltr, start):
+            address = fltr.ctypes.data + (uint64(start) >> 3)
+            return address_to_value(address)
+
+        return get_value
+
+    # slow version if nbits is not a multiple of 8
+    elif 0 < nbits <= 64:
+        padded_nbits = ((int(nbits) + 7) & (-8)) + 8
+        int_type = uint64
+        mask = ir.Constant(ir.IntType(64), int(nbits * '1', 2))
+        finalcast = ir.IntType(64)
+        if padded_nbits <= 64:
+            cast = ir.IntType(64)
+            pointer = ir.IntType(64).as_pointer()
+        else:
+            cast = ir.IntType(128)
+            pointer = ir.IntType(128).as_pointer()
+
+        @nb.extending.intrinsic
+        def address_to_value(typingctx, address, offset):
+            """ returns the value stored at a given memory address """
+            sig = int_type(address, offset)
+            def codegen(cgctx, builder, sig, args):
+                ptr = builder.inttoptr(args[0], pointer)
+                value = builder.load(ptr)
+                shift = builder.zext(args[1], cast)
+                value = builder.lshr(value, shift)
+                value = builder.trunc(value, finalcast)
+                return builder.and_(value, mask)
+            return sig, codegen
+
+        @njit(nogil=True, locals=dict(address=uint64, offset=uint64))
+        def get_value(array, pos):
+            address = array.ctypes.data + (uint64(pos) >> 3)
+            return address_to_value(address, (uint64(pos) & 7))
+
+        return get_value
+
+    return None
+
+def compile_store_value(nbits):
+    # fast version for multiples of 8 (cast byte counter)
+    if nbits in [8, 16, 32, 64]:
+        pointer = ir.IntType(nbits).as_pointer()
+        signatures = {8: uint8, 16: uint16, 32: uint32, 64: uint64}
+        int_type = signatures[nbits]
+
+        @nb.extending.intrinsic
+        def store_value_at_address(typingctx, address, value):
+            """returns the value stored at a given memory address """
+            sig = nb.void(nb.types.uintp, int_type)
+            def codegen(cgctx, builder, sig, args):
+                ptr = builder.inttoptr(args[0], pointer)
+                builder.store(args[1], ptr)
+            return sig, codegen
+
+        @njit(nogil=True, locals=dict(pos=uint64))
+        def store_value(array, pos, value):
+            pos = array.ctypes.data + (uint64(pos) >> 3)
+            store_value_at_address(pos, value)
+
+        return store_value
+
+    # slow version if nbits is not a multiple of 8
+    elif 0 < nbits <= 64:
+        padded_nbits = ((int(nbits) + 7) & (-8)) + 8
+        if padded_nbits <= 64:
+            cast = ir.IntType(64)
+            pointer = ir.IntType(64).as_pointer()
+            ones = ir.Constant(ir.IntType(64), int(2**64 - 1))
+            mask = ir.Constant(ir.IntType(64), 2**int(nbits) - 1)
+        elif padded_nbits <= 128:
+            cast = ir.IntType(128)
+            pointer = ir.IntType(128).as_pointer()
+            ones = ir.Constant(ir.IntType(128), 2**128 - 1)
+            mask = ir.Constant(ir.IntType(128), 2**int(nbits) - 1)
+
+        @nb.extending.intrinsic
+        def store_value_at_address(typingctx, address, value, shift):
+            """returns the value stored at a given memory address """
+            sig = nb.void(nb.types.uintp, address, value, shift)
+            def codegen(cgctx, builder, sig, args):
+                ptr = builder.inttoptr(args[0], pointer)
+                value = builder.load(ptr)
+                zero_mask = builder.shl(mask, builder.zext(args[2], cast))
+                zero_mask = builder.xor(zero_mask, ones)
+                value = builder.and_(value, zero_mask)
+                insert = builder.shl(args[1], args[2])
+                value = builder.or_(value, builder.zext(insert, cast))
+                builder.store(value, ptr)
+            return sig, codegen
+
+        @njit(nogil=True, locals=dict(address=uint64, shift=uint64, value=uint64))
+        def store_value(array, pos, value):
+            address = array.ctypes.data + (uint64(pos) >> 3)
+            shift = uint64(pos) & 7
+            store_value_at_address(address, value, shift)
+
+        return store_value
+
+    raise NotImplementedError("Only storing values <= 64 bits supported")

panxpress/dnaencode.py

@@ -0,0 +1,277 @@
+
+import numpy as np
+from numba import njit, uint8, int64, uint64, prange
+
+from .mask import Mask
+
+
+U64_MINUSONE = uint64(np.iinfo(np.uint64).max)
+
+
+# encoding DNA ###############################
+
+def _get_table_dna_to_2bits(default=4):
+    b = np.full(256, default, dtype=np.uint8)
+    b[97] = 0   # a
+    b[65] = 0   # A
+    b[99] = 1   # c
+    b[67] = 1   # C
+    b[103] = 2  # g
+    b[71] = 2   # G
+    b[116] = 3  # t
+    b[84] = 3   # T
+    b[117] = 3  # u
+    b[85] = 3   # U
+    return b
+
+
+_TABLE_DNA_TO_2BITS = _get_table_dna_to_2bits()
+
+
+@njit(nogil=True, locals=dict(i=int64))
+def _dna_to_2bits(x, table):
+    for i in range(x.size):
+        x[i] = table[x[i]]
+
+
+# this is the one to use!
+@njit(nogil=True)
+def quick_dna_to_2bits(x):
+    for i in range(len(x)):
+        x[i] = _TABLE_DNA_TO_2BITS[x[i]]
+
+
+@njit(nogil=True, parallel=True)
+def parallel_dna_to_2bits(x):
+    for i in prange(x.size):
+        x[i] = _TABLE_DNA_TO_2BITS[x[i]]
+
+
+# numba compile error
+# @njit(nogil=True, locals=dict(seq=uint8[:], table=uint8[:]))
+def dna_to_2bits(seq, table=_TABLE_DNA_TO_2BITS):
+    # we expect seq to be a bytearray
+    # xx = np.array(seq, dtype=np.uint8)
+    if isinstance(seq, bytes):
+        xx = np.frombuffer(bytearray(seq), dtype=np.uint8)
+    else:
+        xx = np.frombuffer(seq, dtype=np.uint8)
+    _dna_to_2bits(xx, table)
+    return xx
+
+
+_TABLE_BITS_TO_DNASTR = ["A", "C", "G", "T"]
+
+
+def qcode_to_dnastr(qcode, q, table=_TABLE_BITS_TO_DNASTR):
+    qc = int(qcode)
+    return "".join([table[((qc >> (2 * (q - i - 1))) & 3)] for i in range(q)])
+
+
+@njit(nogil=True, locals=dict(base=uint64))
+def write_qcode_to_buffer(qcode, q, buf, start):
+    for i in range(q):
+        base = (qcode >> (2 * (q - i - 1))) & 3
+        buf[start + i] = uint8(base)
+
+
+# no need to njit!
+def _get_table_2bits_to_dna(default=4):
+    b = np.full(256, 35, dtype=np.uint8)  # fill with b'#'
+    b[0] = 65
+    b[1] = 67
+    b[2] = 71
+    b[3] = 84
+    b[default] = 78
+    return b
+
+
+_TABLE_2BITS_TO_DNA = _get_table_2bits_to_dna()
+
+
+@njit(nogil=True)
+def twobits_to_dna_inplace(buf, start=0, end=0):
+    if end <= 0:
+        end = len(buf) - end
+    for i in range(start, end):
+        buf[i] = _TABLE_2BITS_TO_DNA[buf[i]]
+
+
+# ########## reverse complements and canonical representation ##############
+
+@njit(nogil=True,
+    locals=dict(c1=uint8, c2=uint8, n=int64, drei=uint8))
+def revcomp_inplace(seq):
+    n = seq.size
+    drei = 3
+    for i in range((n + 1) // 2):
+        j = n - 1 - i
+        c1 = seq[i]
+        c2 = seq[j]
+        seq[j] = drei - c1 if c1 < 4 else c1
+        seq[i] = drei - c2 if c2 < 4 else c2
+
+
+@njit(nogil=True, locals=dict(
+    c=uint8, n=int64, drei=uint8, rc=uint8[:]))
+def revcomp_to_buffer(seq, rc):
+    n = seq.size
+    drei = 3
+    for i in range(n):
+        c = seq[n - 1 - i]
+        rc[i] = drei - c if c < 4 else c
+
+
+@njit(nogil=True, locals=dict(rc=uint8[:]))
+def revcomp(seq):
+    rc = np.empty_like(seq, dtype=np.uint8)
+    revcomp_to_buffer(seq, rc)
+    return rc
+
+
+@njit(nogil=True, locals=dict(
+    code=uint64, drei=uint64, rc=uint64, c=uint64))
+def revcomp_code(code, q):
+    # only works for 0 <= q <= 31 !
+    # when using uints, due to a potential bug in numpy/numba,
+    # we would have to re-declare code as uint64 locally.
+    drei = uint64(3)
+    rc = 0
+    for i in range(q):
+        c = drei - (code & drei)
+        rc = (rc << 2) | c
+        code >>= 2
+    return rc
+
+
+@njit(nogil=True, locals=dict(
+    code=uint64, c=uint64))
+def _get_rctable():
+    rctable = np.zeros(256, dtype=np.uint64)
+    for c in range(256):
+        rctable[c] = revcomp_code(c, 4)
+    return rctable
+
+
+_RCTABLE = _get_rctable()
+
+
+@njit(nogil=True, locals=dict(
+    code=uint64, rc=uint64, c=uint64))
+def revcomp_code_table(code, q):
+    rc = 0
+    while q >= 4:
+        c = _RCTABLE[code & 255]
+        rc = (rc << 8) | c
+        code >>= 8
+        q -= 4
+    for i in range(q):
+        c = 3 - (code & 3)
+        rc = (rc << 2) | c
+        code >>= 2
+    return rc
+
+
+@njit(nogil=True, locals=dict(
+    code=int64, rc=int64))
+def canonical_code(code, q):
+    rc = revcomp_code(code, q)
+    return code if code <= rc else rc
+
+
+def compile_revcomp_and_canonical_code(q, rcmode):
+    """
+    return pair of functions (revcomp_code_q, canonical_code_q)
+    specialized for q-gram codes for the given value of q.
+    It is expected that LLVM optimization does loop unrolling.
+    """
+    @njit(nogil=True, locals=dict(
+        code=uint64, rc=uint64, c=uint64))
+    def _rc(code):
+        rc = 0
+        t = q // 4
+        for i in range(t):
+            c = _RCTABLE[code & 255]
+            rc = (rc << 8) | c
+            code >>= 8
+        r = q % 4
+        for i in range(r):
+            c = 3 - (code & 3)
+            rc = (rc << 2) | c
+            code >>= 2
+        return rc
+    
+    if rcmode == "min":
+        @njit(nogil=True, locals=dict(
+            code=uint64, rc=uint64))
+        def _cc(code):
+            rc = _rc(code)
+            return code if code <= rc else rc
+    elif rcmode == "max":
+        @njit(nogil=True, locals=dict(
+            code=uint64, rc=uint64))
+        def _cc(code):
+            rc = _rc(code)
+            return code if code >= rc else rc
+    elif rcmode == "r":
+        @njit(nogil=True, locals=dict(
+            code=uint64, rc=uint64))
+        def _cc(code):
+            rc = _rc(code)
+            return rc
+    else:  # 'f', 'both', ...
+        @njit(nogil=True, locals=dict(
+            code=uint64, rc=uint64))
+        def _cc(code):
+            return code
+
+    return _rc, _cc
+
+
+# translation of a DNA buffer into codes
+def compile_twobit_to_codes(tmask, rcmode, invalid=U64_MINUSONE):
+    # tmask should be a mask tuple, but might be an int in some cases
+    if isinstance(tmask, int):
+        k = w = tmask  # be safe here
+        tmask = tuple(range(k))
+    elif isinstance(tmask, Mask):
+        k = tmask.k
+        w = tmask.w
+        tmask = tmask.tuple
+    elif type(tmask) is tuple:
+        k, w = len(tmask), tmask[-1] + 1
+    else:
+        raise ValueError(f"mask type {type(tmask)} is not supported.")
+    _, ccc = compile_revcomp_and_canonical_code(k, rcmode)
+
+    @njit(nogil=True, locals=dict(code=uint64))
+    def twobit_to_codes(seq, out, start=0, n=-1):
+        """write n (or all) canonical k-mer codes from seq[start...] into out buffer"""
+        if n == -1:
+            n = len(seq) - w + 1 - start
+        for i in range(start, start + n):
+            code = 0
+            for j in tmask:
+                c = seq[i + j]
+                if c >= 4:
+                    out[i - start] = uint64(invalid)
+                    break
+                code = (code << 2) | c
+            else:
+                code = ccc(code)
+                out[i - start] = code
+
+    @njit(nogil=True, locals=dict(code=uint64))
+    def twobit_to_code(seq, start=0):
+        """return a single canonical code at seq[start...]"""
+        code = 0
+        for j in tmask:
+            c = seq[start + j]
+            if c >= 4:
+                return uint64(invalid)
+            code = (code << 2) | c
+        else:
+            code = ccc(code)
+            return code
+
+    return twobit_to_codes, twobit_to_code

panxpress/dnaencode_fast.py

@@ -0,0 +1,156 @@
+"""
+Fast DNA encoding and computation of the reverse complement A=00, C=01, T=10, G=11, U=10
+Other characters are not treated correctly, except N=7>3 and n = 6>3
+"""
+
+from numba import njit, uint64
+import numpy as np
+from .mask import Mask
+
+U64_MINUSONE = uint64(np.iinfo(np.uint64).max)
+
+
+@njit(nogil=True)
+def quick_dna_to_2bits(x):
+    for i in range(len(x)):
+        x[i] = (x[i] >> 1) & 7
+
+
+@njit(nogil=True)
+def dna_to_2bits(x, y):
+    for i in range(len(x)):
+        y[i] = (x[i] >> 1) & 7
+
+
+# no need to njit!
+def _get_table_2bits_to_dna():
+    b = np.full(256, 35, dtype=np.uint8)  # fill with b'#'
+    b[0] = 65
+    b[1] = 67
+    b[2] = 84
+    b[3] = 71
+    b[6] = 78
+    b[7] = 78
+    return b
+
+
+_TABLE_2BITS_TO_DNA = _get_table_2bits_to_dna()
+
+
+@njit(nogil=True)
+def twobits_to_dna_inplace(buf, start=0, end=0):
+    if end <= 0:
+        end = len(buf) - end
+    for i in range(start, end):
+        buf[i] = _TABLE_2BITS_TO_DNA[buf[i]]
+
+
+_TABLE_BITS_TO_DNASTR = ["A", "C", "T", "G"]
+
+
+def qcode_to_dnastr(qcode, q, table=_TABLE_BITS_TO_DNASTR):
+    qc = int(qcode)
+    return "".join([table[((qc >> (2 * (q - i - 1))) & 3)] for i in range(q)])
+
+
+def compile_revcomp_and_canonical_code(q, rcmode):
+    """
+    return pair of functions (revcomp_code_q, canonical_code_q)
+    specialized for q-gram codes for the given value of q.
+    It is expected that LLVM optimization does loop unrolling.
+    """
+
+    mask = uint64(int(q * '10', 2))
+    shift = uint64(64 - 2 * q)
+
+    @njit(nogil=True, locals=dict(value=uint64))
+    def reverse_in_pairs(value):
+        value = ((value & 0x3333333333333333) << 2) | ((value & 0xCCCCCCCCCCCCCCCC) >> 2)
+        value = ((value & 0x0F0F0F0F0F0F0F0F) << 4) | ((value >> 4) & 0x0F0F0F0F0F0F0F0F)
+        value = ((value & 0x00FF00FF00FF00FF) << 8) | ((value >> 8) & 0x00FF00FF00FF00FF)
+        value = ((value & 0x0000FFFF0000FFFF) << 16) | ((value >> 16) & 0x0000FFFF0000FFFF)
+        value = (value << 32) | (value >> 32)
+        return value
+
+    @njit(nogil=True, locals=dict(value=uint64))
+    def _rc(value):
+        value ^= mask
+        value = reverse_in_pairs(value)
+        return value >> shift
+
+    if rcmode == "min":
+        @njit(nogil=True, locals=dict(
+            code=uint64, rc=uint64))
+        def _cc(code):
+            rc = _rc(code)
+            return code if code <= rc else rc
+    elif rcmode == "max":
+        @njit(nogil=True, locals=dict(
+            code=uint64, rc=uint64))
+        def _cc(code):
+            rc = _rc(code)
+            return code if code >= rc else rc
+    elif rcmode == "r":
+        @njit(nogil=True, locals=dict(
+            code=uint64, rc=uint64))
+        def _cc(code):
+            rc = _rc(code)
+            return rc
+    else:  # 'f', 'both', ...
+        @njit(nogil=True, locals=dict(
+            code=uint64, rc=uint64))
+        def _cc(code):
+            return code
+
+    return _rc, _cc
+
+
+# translation of a DNA buffer into codes
+
+
+def compile_twobit_to_codes(tmask, rcmode, invalid=U64_MINUSONE):
+    # tmask should be a mask tuple, but might be an int in some cases
+    if isinstance(tmask, int):
+        k = w = tmask  # be safe here
+        tmask = tuple(range(k))
+    elif isinstance(tmask, Mask):
+        k = tmask.k
+        w = tmask.w
+        tmask = tmask.tuple
+    elif type(tmask) is tuple:
+        k, w = len(tmask), tmask[-1] + 1
+    else:
+        raise ValueError(f"mask type {type(tmask)} is not supported.")
+    _, ccc = compile_revcomp_and_canonical_code(k, rcmode)
+
+    @njit(nogil=True, locals=dict(code=uint64))
+    def twobit_to_codes(seq, out, start=0, n=-1):
+        """write n (or all) canonical k-mer codes from seq[start...] into out buffer"""
+        if n == -1:
+            n = len(seq) - w + 1 - start
+        for i in range(start, start + n):
+            code = 0
+            for j in tmask:
+                c = seq[i + j]
+                if c >= 4:
+                    out[i - start] = uint64(invalid)
+                    break
+                code = (code << 2) | c
+            else:
+                code = ccc(code)
+                out[i - start] = code
+
+    @njit(nogil=True, locals=dict(code=uint64))
+    def twobit_to_code(seq, start=0):
+        """return a single canonical code at seq[start...]"""
+        code = 0
+        for j in tmask:
+            c = seq[start + j]
+            if c >= 4:
+                return uint64(invalid)
+            code = (code << 2) | c
+        else:
+            code = ccc(code)
+            return code
+
+    return twobit_to_codes, twobit_to_code