<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
FEATURES TO IMPLEMENT
---------------------
* colon directive.
  :myfunc
      some param
      and more
  which would be the equivalent of calling
  :myfunc('some param\nandmore')
* setup string.Formatter custom namespace instead of
  passing in safe_locals
* try to speed up eval, check out mako source maybe?
  speed up eval by making only one eval call. Somehow get all
  : directives, viewing it as an overlay, remember where each
  directive is located, and produce a controlled result of all
  directives with a single eval call. Then layer this back into
  the original document. It's crazy, but it just might work...
  re.findall(':(.*\=.*$|.*\))', f, re.M)
* reserve keywords for safe_locals such as "encoding", "doctype", etc

IDEAS
-----
* static html instead of quick tags
* line escaping for text using \
* attribute syntax (attr=val) can stretch multiple lines
  this one can be done by seeing it doesn't end, buffer it,
  and mark the next iteration a continuation, maybe?
* tag/ self closing tags? does it matter? lxml etree handling this
* whitespace control with > and < does this matter? relevant? etree again
  handling pre elements could be trouble and relevant
* declare a doctype at top of document and charset, pass on to etree.tostring
* forward slash / for html comments, or how about just writing html comments...
  hmm but wrapping indented sections of code, now we're talkin
* /[if IE] oh yeah, need this
* comments specific to the doc that are not rendered anywhere, and indented
  sections are included as part of the comment
* not sure how best to handle this,
  - for x in range(4)
      %p x
  and then an unindent would end this bit
* whitespace preservation?
* :javascript directive?
  :overlay main.hdae namethis
  would process this document then create a dict(namethis=processed) passed to the formatter
  to overlay specified maybe? sounds good for one, but what about multiple includes?
* escaping/unescaping html
* pass namespace of controller method to formatter, that would kick fucking ass
  @publisher(_locals=True)
  def index(self):
      cat = 'meow'
  ---
  %p {cat}
  ---
* conditionals inside of attribute tags (attr=val) somehow something nice.. would probably
  fit in with - for syntax
  :r = range(4)
  - for i in r
      %li({if i+1 == len(r): class=last}) i

BUGFIXES
--------
look for the fixme's atm
FIXME process_plntxt after parse "for-loop" finishes for plain text located

"""
import __builtin__
from lxml import etree
from time import time
import re

# this bit would create seperate groups which might be useful
# for establish what should happen, but then it might not be
# relevant with compile(... 'exec')
# re.findall(':(.*\=.*$)|:(.*\))', f, re.M)

p = re.compile(':(.*\=.*$|.*\))', re.M) # parse out : directives
p2 = re.compile('(\w.*?)=(\w.*?)\ ?,')

safe_locals = {'len': __builtin__.len, 'locals': __builtin__.locals}

def is_func_call(x):
    """
    determines if string is a function call.
    This differentiates between item assignment
    and a function call on the assumption that there is no space
    between func_name and the first instance of (

    # FIXME this should account for the following

        user='(Daniel)'

      by checking for = between ^ and ( index
    """
    if '(' in x and ' ' not in x[:x.index('(')]:
        return True
    return False

def func_to_locals(x):
    """
    This parses a string to be eval'd so as to add function
    calls to the locals() index, ie:

      greeting('Daniel')

    will get transformed to:

      locals()['''greeting('Daniel')'''] = greeting('Daniel')

    while calls such as:

      user = 'Daniel'

    will be left intact.
    """
    if is_func_call(x):
        return '''locals()["""{0}"""]={0}'''.format(x)
    return x

def safe_eval(s):
    return eval(s, {'__builtins__': None}, safe_locals)

def parse_eval(f):
    l = p.findall(f) # [':greeting = lambda x: x', ':greeting("meh")']
    m = [func_to_locals(x) for x in l]
    c = compile(';'.join(m), '<string>', 'exec')
    safe_eval(c) # populates safe_locals
    for x in l:
        if x in safe_locals:
            f = f.replace(':'+x, safe_locals[x])
        else:
            f = f.replace(':'+x, '')
    return f

def parse_doc(f):
    r = []
    plntxt = []

    for l in f.splitlines():
        l = l.rstrip() # remove line break endings
        if l == '':
            continue # skip blank lines

        # inspect directive, determine if plain text
        directive = l.lstrip()[0]
        if directive == '%':
            pass
        elif directive == '#':
            l = l.replace('#', '%#', 1)
        elif directive == '.':
            l = l.replace('.', '%.', 1)
        else:
            plntxt.append(l)
            continue

        # check plntxt queue
        # everything in queue should be same indention width
        # since nowhere in a doc should there be plain text indented to plain text
        for x in plntxt:
            text = x.strip()
            ws = ' '*(len(x)-len(text))
            j = -1

            while j:
                if ws > r[j][0]:
                    r[j][1].text += ' '+text
                    break
                elif ws == r[j][0]:
                    if r[j][1].tail is None: # faster than init'ing tail on element creation everytime
                        r[j][1].tail = text
                    else:
                        r[j][1].tail += ' '+text
                    break
                j -= 1

        plntxt = []

        #t = t.format(**safe_locals)
        l = l.partition('%') # ('    ', '%', 'tag#id.class(attr=val) content')

        # determine tag attributes
        # TODO setup attributes to span multiple lines
        attr = None
        if '(' in l[2]:
            op_i = l[2].index('(')
            if ' ' not in l[2][:op_i]: # there should be no spaces in %tag#id.class(...
                cp_i = l[2].index(')')+1
                attr = l[2][op_i:cp_i]
                u = l[2].replace(attr, '').partition(' ')
            else:
                u = l[2].partition(' ')
        else:
            u = l[2].partition(' ')

        # this is pretty sweet, basically it will always return the following format
        #   [('tag', '', 'class'), ('#', '', ''), ('id', '.', 'class')]
        # no matter the string order of elements, for example
        #   tag#id.class.class
        #   tag.class#id.class
        #   tag.class.class#id
        #   #id.class.class
        #   #id
        # for all of these examples, tag will always be [0][0], id will always be [2][0]
        # and class will always be [0][2]+[2][2]
        # and this should be effin fast comparitively (to regex, other builtin string operations)
        tag = [x.partition('.') for x in u[0].partition('#')]

        #
        e = etree.Element(tag[0][0] or 'div')
        e.text = u[2]

        if tag[2][0] != '':
            e.attrib['id'] = tag[2][0]

        _class = tag[0][2] + ' ' + tag[2][2]
        if _class != ' ':
            e.attrib['class'] = _class.replace('.', ' ').strip()

        if attr is not None:
            for x in attr[1:-1].split(','):
                k, v = x.split('=')
                e.attrib[k.strip()] = v.strip()

        r.append((l[0], e)) # ('    ', etree.Element)
    return r

def heuristic_test(parsed):
    """
    This returns a step-count based on indention. The step-count represents
    the number of steps upward an element needs to move to find its parent.
    This is only accurate for counts of more and same. Less can not be divined
    """
    m = []
    more = 0
    same = 0

    prev = parsed[0][0]
    for d in parsed[1:]:
        if d[0] > prev:
            m.append(1)
            more += 1
            same = 0
        elif d[0] == prev:
            m.append(2+same) # (m->s) == 2
            same += 1
        elif d[0] < prev:
            m.append(more+same+1)
            more = 0
            same = 0
        prev = d[0]
    return m

def hr_build(parsed):
    """ FIXME something is busted here as can be seen with bench/hdae/template.html
    This combines the heuristic and relative build methods. The heuristic method
    is able to figure out position of positively indented elements faster but is unable to
    tell position when indention decreases. At this point, we use a relative position
    method to locate this node's parent. The speed increase is minimal, but every bit helps.

    When a section unindents, I think its safe to assume that it will fall on a pre-existing
    indention level. All this is hardly my forte so this may be wrong. By recording a map of
    all the latest elements on an indention level, we can access the last element on the same
    indention level, get its parent, and append the new element to it. The one catch with this
    method is we need to remove any mappings to larger indentions of the current unindent so
    future elements do not get appended to incorrect items if variable indention is used.
    """
    r = parsed
    more = 0
    same = 0
    m = {}

    prev = parsed[0][0]
    for i, d in enumerate(parsed[1:]):
        i += 1
        if d[0] > prev:
            r[i-1][1].append(r[i][1])   # ('    ', Element)[1].append(...)
            more += 1
            same = 0
            m[d[0]] = i         # m['    '] = 1
        elif d[0] == prev:
            r[i-(2+same)][1].append(r[i][1])    # (more_indent->same_indent) == 2
            same += 1
        elif d[0] < prev:
            j = m[d[0]]     # j = [i]
            m[d[0]] = i         # m['    '] = 1
            #print etree.tostring(r[0][1])
            parent = r[j][1].getparent()
            parent.append(d[1])

            # clear step counts
            more = 0
            same = 0

            # purge mapping of larger indents then this unindent
            for k in m.keys():
                if k > d[0]:
                    m.pop(k)

        prev = d[0]
    return r

def relative_build(parsed):
    """
    stable parsing of document, but slightly slower then heuristic (which is still questionable atm)
    """
    r = [x for x in parsed if not isinstance(x, str)]
    m = {}
    prev = parsed[0][0]
    for i, d in enumerate(parsed[1:]):
        i += 1
        k, v = d
        if k in m:
            j = m[k]
            parent = r[j][1].getparent()
            child = r[i][1]
            parent.append(child)
            m[k] = i
        else:
            m[k] = i
            parent = r[i-1][1]
            child = r[i][1]
            parent.append(child)
        if d[0] < prev:
            for k in m.keys():
                if k > d[0]:
                    m.pop(k)
        prev = d[0]
    return r

###################################

def haml(f, t='r'):
    _f = open(f).read()
    # process eval stuff first
    _f = parse_eval(_f)

    #parse document next
    l = parse_doc(_f)
    if t == 'r':
        b = relative_build(l)
    elif t == 'hr':
        b = hr_build(l)
    elif t == 'h':
        b = heuristic_test(l)
    return etree.tostring(b[0][1])

def test(func):
    from time import time
    times = []
    for x in range(20):
        a = time()
        func()
        times.append(time()-a)
    print(min(times))

if __name__ == '__main__':
    import sys
    f = sys.argv[1]
    t = sys.argv[2]

    def render(f, t):
        def _render():
            return haml(f, t)
        return _render

    if t == 'r':
        test(render(f, t))
    if t == 'hr':
        test(render(f, t))
    if t == 'haml':
        print haml(f)

<commit_msg>fixed hr_build, hopefully for good<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
FEATURES TO IMPLEMENT
---------------------
* colon directive.
  :myfunc
      some param
      and more
  which would be the equivalent of calling
  :myfunc('some param\nandmore')
* setup string.Formatter custom namespace instead of
  passing in safe_locals
* try to speed up eval, check out mako source maybe?
  speed up eval by making only one eval call. Somehow get all
  : directives, viewing it as an overlay, remember where each
  directive is located, and produce a controlled result of all
  directives with a single eval call. Then layer this back into
  the original document. It's crazy, but it just might work...
  re.findall(':(.*\=.*$|.*\))', f, re.M)
* reserve keywords for safe_locals such as "encoding", "doctype", etc

IDEAS
-----
* static html instead of quick tags
* line escaping for text using \
* attribute syntax (attr=val) can stretch multiple lines
  this one can be done by seeing it doesn't end, buffer it,
  and mark the next iteration a continuation, maybe?
* tag/ self closing tags? does it matter? lxml etree handling this
* whitespace control with > and < does this matter? relevant? etree again
  handling pre elements could be trouble and relevant
* declare a doctype at top of document and charset, pass on to etree.tostring
* forward slash / for html comments, or how about just writing html comments...
  hmm but wrapping indented sections of code, now we're talkin
* /[if IE] oh yeah, need this
* comments specific to the doc that are not rendered anywhere, and indented
  sections are included as part of the comment
* not sure how best to handle this,
  - for x in range(4)
      %p x
  and then an unindent would end this bit
* whitespace preservation?
* :javascript directive?
  :overlay main.hdae namethis
  would process this document then create a dict(namethis=processed) passed to the formatter
  to overlay specified maybe? sounds good for one, but what about multiple includes?
* escaping/unescaping html
* pass namespace of controller method to formatter, that would kick fucking ass
  @publisher(_locals=True)
  def index(self):
      cat = 'meow'
  ---
  %p {cat}
  ---
* conditionals inside of attribute tags (attr=val) somehow something nice.. would probably
  fit in with - for syntax
  :r = range(4)
  - for i in r
      %li({if i+1 == len(r): class=last}) i

BUGFIXES
--------
look for the fixme's atm
FIXME process_plntxt after parse "for-loop" finishes for plain text located

"""
import __builtin__
from lxml import etree
from time import time
import re

# this bit would create seperate groups which might be useful
# for establish what should happen, but then it might not be
# relevant with compile(... 'exec')
# re.findall(':(.*\=.*$)|:(.*\))', f, re.M)

p = re.compile(':(.*\=.*$|.*\))', re.M) # parse out : directives
p2 = re.compile('(\w.*?)=(\w.*?)\ ?,')

safe_locals = {'len': __builtin__.len, 'locals': __builtin__.locals}

def is_func_call(x):
    """
    determines if string is a function call.
    This differentiates between item assignment
    and a function call on the assumption that there is no space
    between func_name and the first instance of (

    # FIXME this should account for the following

        user='(Daniel)'

      by checking for = between ^ and ( index
    """
    if '(' in x and ' ' not in x[:x.index('(')]:
        return True
    return False

def func_to_locals(x):
    """
    This parses a string to be eval'd so as to add function
    calls to the locals() index, ie:

      greeting('Daniel')

    will get transformed to:

      locals()['''greeting('Daniel')'''] = greeting('Daniel')

    while calls such as:

      user = 'Daniel'

    will be left intact.
    """
    if is_func_call(x):
        return '''locals()["""{0}"""]={0}'''.format(x)
    return x

def safe_eval(s):
    return eval(s, {'__builtins__': None}, safe_locals)

def parse_eval(f):
    l = p.findall(f) # [':greeting = lambda x: x', ':greeting("meh")']
    m = [func_to_locals(x) for x in l]
    c = compile(';'.join(m), '<string>', 'exec')
    safe_eval(c) # populates safe_locals
    for x in l:
        if x in safe_locals:
            f = f.replace(':'+x, safe_locals[x])
        else:
            f = f.replace(':'+x, '')
    return f

def parse_doc(f):
    r = []
    plntxt = []

    for l in f.splitlines():
        l = l.rstrip() # remove line break endings
        if l == '':
            continue # skip blank lines

        # inspect directive, determine if plain text
        directive = l.lstrip()[0]
        if directive == '%':
            pass
        elif directive == '#':
            l = l.replace('#', '%#', 1)
        elif directive == '.':
            l = l.replace('.', '%.', 1)
        else:
            plntxt.append(l)
            continue

        # check plntxt queue
        # everything in queue should be same indention width
        # since nowhere in a doc should there be plain text indented to plain text
        for x in plntxt:
            text = x.strip()
            ws = ' '*(len(x)-len(text))
            j = -1

            while j:
                if ws > r[j][0]:
                    r[j][1].text += ' '+text
                    break
                elif ws == r[j][0]:
                    if r[j][1].tail is None: # faster than init'ing tail on element creation everytime
                        r[j][1].tail = text
                    else:
                        r[j][1].tail += ' '+text
                    break
                j -= 1

        plntxt = []

        #t = t.format(**safe_locals)
        l = l.partition('%') # ('    ', '%', 'tag#id.class(attr=val) content')

        # determine tag attributes
        # TODO setup attributes to span multiple lines
        attr = None
        if '(' in l[2]:
            op_i = l[2].index('(')
            if ' ' not in l[2][:op_i]: # there should be no spaces in %tag#id.class(...
                cp_i = l[2].index(')')+1
                attr = l[2][op_i:cp_i]
                u = l[2].replace(attr, '').partition(' ')
            else:
                u = l[2].partition(' ')
        else:
            u = l[2].partition(' ')

        # this is pretty sweet, basically it will always return the following format
        #   [('tag', '', 'class'), ('#', '', ''), ('id', '.', 'class')]
        # no matter the string order of elements, for example
        #   tag#id.class.class
        #   tag.class#id.class
        #   tag.class.class#id
        #   #id.class.class
        #   #id
        # for all of these examples, tag will always be [0][0], id will always be [2][0]
        # and class will always be [0][2]+[2][2]
        # and this should be effin fast comparitively (to regex, other builtin string operations)
        tag = [x.partition('.') for x in u[0].partition('#')]

        #
        e = etree.Element(tag[0][0] or 'div')
        e.text = u[2]

        if tag[2][0] != '':
            e.attrib['id'] = tag[2][0]

        _class = tag[0][2] + ' ' + tag[2][2]
        if _class != ' ':
            e.attrib['class'] = _class.replace('.', ' ').strip()

        if attr is not None:
            for x in attr[1:-1].split(','):
                k, v = x.split('=')
                e.attrib[k.strip()] = v.strip()

        r.append((l[0], e)) # ('    ', etree.Element)
    return r

def heuristic_test(parsed):
    """
    This returns a step-count based on indention. The step-count represents
    the number of steps upward an element needs to move to find its parent.
    This is only accurate for counts of more and same. Less can not be divined
    """
    m = []
    more = 0
    same = 0

    prev = parsed[0][0]
    for d in parsed[1:]:
        if d[0] > prev:
            m.append(1)
            more += 1
            same = 0
        elif d[0] == prev:
            m.append(2+same) # (m->s) == 2
            same += 1
        elif d[0] < prev:
            m.append(more+same+1)
            more = 0
            same = 0
        prev = d[0]
    return m

def hr_build(parsed):
    """ FIXME something is busted here as can be seen with bench/hdae/template.html
    This combines the heuristic and relative build methods. The heuristic method
    is able to figure out position of positively indented elements faster but is unable to
    tell position when indention decreases. At this point, we use a relative position
    method to locate this node's parent. The speed increase is minimal, but every bit helps.

    When a section unindents, I think its safe to assume that it will fall on a pre-existing
    indention level. All this is hardly my forte so this may be wrong. By recording a map of
    all the latest elements on an indention level, we can access the last element on the same
    indention level, get its parent, and append the new element to it. The one catch with this
    method is we need to remove any mappings to larger indentions of the current unindent so
    future elements do not get appended to incorrect items if variable indention is used.
    """
    r = parsed
    more = 0
    same = 0
    m = {}

    prev = parsed[0][0]
    for i, d in enumerate(parsed):
        if i is 0:
            continue

        ws = d[0]

        if ws > prev:
            r[i-1][1].append(r[i][1])   # ('    ', Element)[1].append(...)
            m[ws] = i
        elif ws == prev:
            r[i-1][1].getparent().append(r[i][1])
            m[ws] = i
        elif ws < prev:
            j = m[ws]
            r[j][1].getparent().append(r[i][1])
            # purge mapping of larger indents then this unindent
            for k in m.keys():
                if k > ws:
                    m.pop(k)
            m[ws] = i
        prev = ws
    return r

def relative_build(parsed):
    """
    stable parsing of document, but slightly slower then heuristic (which is still questionable atm)
    """
    r = [x for x in parsed if not isinstance(x, str)]
    m = {}
    prev = parsed[0][0]
    for i, d in enumerate(parsed[1:]):
        i += 1
        k, v = d
        if k in m:
            j = m[k]
            parent = r[j][1].getparent()
            child = r[i][1]
            parent.append(child)
            m[k] = i
        else:
            m[k] = i
            parent = r[i-1][1]
            child = r[i][1]
            parent.append(child)
        if d[0] < prev:
            for k in m.keys():
                if k > d[0]:
                    m.pop(k)
        prev = d[0]
    return r

###################################

def haml(f, t='r'):
    _f = open(f).read()
    # process eval stuff first
    _f = parse_eval(_f)

    #parse document next
    l = parse_doc(_f)
    if t == 'r':
        b = relative_build(l)
    elif t == 'hr':
        b = hr_build(l)
    elif t == 'h':
        b = heuristic_test(l)
    return etree.tostring(b[0][1])

def test(func):
    from time import time
    times = []
    for x in range(20):
        a = time()
        func()
        times.append(time()-a)
    print(min(times))

if __name__ == '__main__':
    import sys
    f = sys.argv[1]
    t = sys.argv[2]

    def render(f, t):
        def _render():
            return haml(f, t)
        return _render

    if t == 'r':
        test(render(f, t))
    if t == 'hr':
        test(render(f, t))
    if t == 'haml':
        print haml(f)

<|endoftext|>


<commit_before>"""Internal module for parsing names in gedcom format.
"""

from __future__ import print_function, absolute_import, division

import re

_re_name = re.compile(r"""([^/]*)           # part of given name - group 1
                          (?:/([^/]*)/)?    # optional /surname/ - group 2
                          ([^/]*)           # part of given name - group 3
                       """, re.X)

def split_name(name):
    """Extracts pieces of name from full name string.

    Full name can have one of these formats:
        <NAME_TEXT> |
        /<NAME_TEXT>/ |
        <NAME_TEXT> /<NAME_TEXT>/ |
        /<NAME_TEXT>/ <NAME_TEXT> |
        <NAME_TEXT> /<NAME_TEXT>/ <NAME_TEXT>

    <NAME_TEXT> can include almost anything excluding commas, numbers,
    special characters (though some test files use numbers for the names).
    Text between slashes is considered a surname, outside slashes - given
    name.

    This method splits full name into pieces at slashes, e.g.:

        "First /Last/" -> ("First", "Last", "")
        "/Last/ First" -> ("", "Last", "First")
        "First /Last/ Jr." -> ("First", "Last", "Jr.")
        "First Jr." -> ("First Jr.", "", "")

    :param str name: Full name string.
    :return: 2-tuple `(given1, surname, given2)`, `surname` or `given` will
        be empty strings if they are not present in full string.
    """
    given1, _, rem = name.partition("/")
    surname, _, given2 = rem.partition("/")
    return given1.strip(), surname.strip(), given2.strip()
<commit_msg>Remove unused regex from detail/name<commit_after>"""Internal module for parsing names in gedcom format.
"""

from __future__ import print_function, absolute_import, division

def split_name(name):
    """Extracts pieces of name from full name string.

    Full name can have one of these formats:
        <NAME_TEXT> |
        /<NAME_TEXT>/ |
        <NAME_TEXT> /<NAME_TEXT>/ |
        /<NAME_TEXT>/ <NAME_TEXT> |
        <NAME_TEXT> /<NAME_TEXT>/ <NAME_TEXT>

    <NAME_TEXT> can include almost anything excluding commas, numbers,
    special characters (though some test files use numbers for the names).
    Text between slashes is considered a surname, outside slashes - given
    name.

    This method splits full name into pieces at slashes, e.g.:

        "First /Last/" -> ("First", "Last", "")
        "/Last/ First" -> ("", "Last", "First")
        "First /Last/ Jr." -> ("First", "Last", "Jr.")
        "First Jr." -> ("First Jr.", "", "")

    :param str name: Full name string.
    :return: 2-tuple `(given1, surname, given2)`, `surname` or `given` will
        be empty strings if they are not present in full string.
    """
    given1, _, rem = name.partition("/")
    surname, _, given2 = rem.partition("/")
    return given1.strip(), surname.strip(), given2.strip()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
:mod:`pambox.distort` regroups various types of distortions and processings
that can be applied to signals.
"""
from __future__ import division, print_function

import numpy as np
import scipy as sp
from six.moves import zip
from scipy.io import wavfile

from pambox import utils
from pambox.utils import fftfilt

try:
    _ = np.use_fastnumpy
    from numpy.fft import fft, ifft, rfft, irfft
except AttributeError:
    from scipy.fftpack import fft, ifft
    from numpy.fft import rfft, irfft

def mix_noise(clean, noise, sent_level, snr=None):
    """Mix a signal signal noise at a given signal-to-noise ratio.

    Parameters
    ----------
    clean : ndarray
        Clean signal.
    noise : ndarray
        Noise signal.
    sent_level : float
        Sentence level, in dB SPL.
    snr :
        Signal-to-noise ratio at which to mix the signals, in dB. If snr is
        `None`,  no noise is mixed with the signal (Default value = None)

    Returns
    -------
    tuple of ndarrays
        Returns the clean signal, the mixture, and the noise.

    """

    # Pick a random section of the noise
    n_clean = len(clean)
    n_noise = len(noise)
    if n_noise > n_clean:
        start_idx = np.random.randint(n_noise - n_clean)
        noise = noise[start_idx:start_idx + n_clean]

    if snr is not None:
        # Get speech level and set noise level accordingly
        # clean_level = utils.dbspl(clean)
        # noise = utils.setdbspl(noise, clean_level - snr)
        noise = noise / utils.rms(noise) * 10 ** ((sent_level - snr) / 20)
        mix = clean + noise
    else:
        mix = clean

    return clean, mix, noise

def phase_jitter(x, a):
    """
    Apply phase jitter to a signal.

    The expression of phase jitter is:

    .. math:: y(t) = s(t) * cos(\Phi(t)),

    where :math:`\Phi(t)` is a random process uniformly distributed over
    :math:`[0, 2\pi\alpha]`. The effect of the jitter when \alpha is 0.5 or 1
    is to completely destroy the carrier signal, effictively yielding
    modulated white noise.

    Parameters
    ----------
    x : ndarray
       Signal
    a : float
        Phase jitter parameter, typically between 0 and 1, but it can be
        anything.

    Returns
    -------
    ndarray
        Processed signal of the same dimension as the input signal.

    """
    n = len(x)
    return x * np.cos(2 * np.pi * a * np.random.random_sample(n))

def reverb(x, rt):
    """
    Applies reverberation to a signal.

    Parameters
    ----------
    x : ndarray
       Input signal.
    rt : float
        Reverberation time

    Returns
    -------
    ndarray
        Processed signal.

    """
    pass

def spec_sub(x, noise, factor, w=1024 / 2., padz=1024 / 2., shift_p=0.5):
    """
    Apply spectral subtraction to a signal.

    The defaul values of the parameters are typical for a sampling frequency of
    44100 Hz. Note that (W+padz) is the final frame window and hence the fft
    length (it is normally chose as a power of 2).

    Parameters
    ----------
    x : ndarray
        Input signal
    noise :
        Input noise signal
    factor : float
        Noise subtraction factor, must be larger than 0.
    w : int
        Frame length, in samples. (Default value = 1024 / 2.)
    padz : int
        Zero padding (pad with padz/2 from the left and the right) (Default
        value = 1024 / 2.)
    shift_p : float
         Shift percentage (overlap) between each window, in fraction of the
         window size (Default value = 0.5)

    Returns
    -------
    clean_estimate : ndarray
        Estimate of the clean signal.
    noise_estimate : ndarray
        Estimate of the noisy signal.

    """
    wnd = np.hanning(w + 2)  # create hanning window with length = W
    wnd = wnd[1:-1]

    stim = np.vstack((x, noise))

    len_signal = stim.shape[-1]  # Signal length
    shift_p_indexes = np.floor(w * shift_p)
    n_segments = np.floor((len_signal - w) / shift_p_indexes + 1)
    len_segment = w + padz * 2 * shift_p
    y = np.empty((2, n_segments, len_segment))
    # Initialize arrays for spectral subtraction. Use only positive
    # frequencies.
    Y_hat = np.empty((n_segments, len_segment / 2 + 1))
    PN_hat = Y_hat.copy()

    # For each signal
    for k in range(2):
        # CUT THE APPROPRIATE SIGNAL FRAMES
        indexes = np.tile(np.arange(w), (n_segments, 1))
        index_shift = np.arange(n_segments) * shift_p_indexes
        indexes = indexes + index_shift[:, np.newaxis]
        y_tmp = stim[k]
        y_tmp = y_tmp[indexes.astype('int')] * wnd
        # PAD WITH ZEROS
        pad = np.zeros((n_segments, padz / 2))
        y_pad = np.hstack((pad, y_tmp, pad))
        y[k, :, :] = y_pad

    # FREQUENCY DOMAIN

    # signal:
    Y = fft(y[0])
    # YY = Y(1:round(end/2)+1,:); # Half window (exploit the symmetry)
    YY = Y[:, :(len_segment / 2 + 1)]  # Half window (exploit the symmetry)
    YPhase = np.angle(YY)  # Phase
    Y1 = np.abs(YY)  # Spectrum
    Y2 = Y1 ** 2  # Power Spectrum

    # noise:
    Y_N = fft(y[1])
    YY_N = Y_N[:, :(len_segment / 2 + 1)]  # Half window (exploit the symmetry)
    Y_NPhase = np.angle(YY_N)  # Phase
    Y_N1 = np.abs(YY_N)  # Spectrum
    Y_N2 = Y_N1 ** 2  # Power Spectrum

    # The noise "estimate" is simply the average of the noise power
    # spectral density in the frame:
    P_N = Y_N2.mean(axis=-1)

    Y_hat = Y2 - factor * P_N[:, np.newaxis]  # subtraction
    Y_hat = np.maximum(Y_hat, 0)  # Make the minima equal zero
    PN_hat = Y_N2 - factor * P_N[:, np.newaxis]  # subtraction for noise alone
    # PN_hat = np.maximum(PN_hat, 0)
    PN_hat[Y_hat == 0] = 0

    Y_hat[0:2, :] = 0
    PN_hat[0:2, :] = 0
    # Combining the estimated power spectrum with the original noisy phase,
    # and add the frames using an overlap-add technique
    output_Y = overlap_and_add(np.sqrt(Y_hat), YPhase, (w + padz), shift_p * w)
    output_N = overlap_and_add(np.sqrt(PN_hat.astype('complex')),
                               Y_NPhase, (w + padz), shift_p * w)

    return output_Y, output_N

def overlap_and_add(powers, phases, len_window, shift_size):
    """Reconstruct a signal with the overlap and add method.

    Parameters
    ----------
    powers : ndarray
        Magnitude of the power spectrum of the signal to reconstruct.
    phases : ndarray
        Phase of the signal to reconstruct.
    len_window : int
        Frame length, in samples.
    shift_size : int
        Shift length. For non overlapping signals, in would equal `len_window`.
        For 50% overlapping signals, it would be `len_window/2`.

    Returns
    -------
    ndarray
        Reconstructed time-domain signal.

    """
    n_frames, len_frame = powers.shape
    spectrum = powers * np.exp(1j * phases)
    signal = np.zeros(n_frames * shift_size + len_window - shift_size)

    # Create full spectrum, by joining conjugated positive spectrum
    if len_window % 2:
        # Do no duplicate the DC bin
        spectrum = np.hstack((spectrum, np.conj(np.fliplr(spectrum[:, 1:]))))
    else:
        # If odd-numbered, do not duplicated the DC ans FS/2 bins
        spectrum = np.hstack((spectrum,
                              np.conj(np.fliplr(spectrum[:, 1:-1]))))

    signal = np.zeros((n_frames - 1) * shift_size + len_window)

    for i_frame, hop in enumerate(range(0,
                                        len(signal) - int(len_window) + 1,
                                        int(shift_size))):
        signal[hop:hop + len_window] \
            += np.real(ifft(spectrum[i_frame], len_window))
    return signal

class Westermann_crm(object):

    """Applies HRTF and BRIR for a given target and masker distance."""

    def __init__(self, fs=40000):
        """

        Parameters
        ----------
        fs : int
             Samping frequenc of the process. (Default value = 40000)

        Returns
        -------

        """
        self.dist = np.asarray([0.5, 2, 5, 10])
        self.fs = fs
        self.brir = self._load_brirs()
        self.delays = self._find_delay()

    def _load_brirs(self):
        """Loads BRIRs from file."""
        brirs = {}
        for d in self.dist:
            d_str = self._normalize_fname(d)
            fname = '../stimuli/crm/brirs_40k/aud' \
                    + d_str + 'm.wav'
            wav = wavfile.read(fname)
            self.brir[d] = np.array(wav[1].astype('float') / 2. ** 15).T

    def _find_delay(self):
        """Calculates the delay of the direct sound, in samples."""
        delays = {}
        for k, v in self.brir.iteritems():
            x = np.mean(v, axis=0)
            delays[k] = np.abs(x).argmax()
        return delays

    @staticmethod
    def _normalize_fname(d):
        """

        Parameters
        ----------
        d : float

        Returns
        -------

        """
        if d > 1:
            d_str = str('%d' % d)
        else:
            d_str = str(d).replace('.', '')
        return d_str

    def _load_eqfilt(self, tdist, mdist):
        """
        Returns the equalization filter for the pair of target and masker.

        Parameters
        ----------
        tdist : float
            Target distance in meters. Must be in the set (0.5, 2, 5, 10).
        mdist :
            Masker distance in meters. Must be in the set (0.5, 2, 5, 10).

        Returns
        -------
        ndarray
            Equalization filter.

        """
        eqfilt_name = 't{}m_m{}m.mat'.format(self._normalize_fname(tdist),
                                             self._normalize_fname(mdist))
        eqfilt_path = '../stimuli/crm/eqfilts_{}/{}'.format(self.fs,
                                                            eqfilt_name)
        try:
            eqfilt = sp.io.loadmat(eqfilt_path, squeeze_me=True)
        except IOError:
            raise IOError('Cannot file file %s' % eqfilt_path)
        return eqfilt

    def apply(self, x, m, tdist, mdist, align=True):
        """Applies the "Westermann" distortion to a target and masker.

        Applies the BRIR of the required distance to the target. If the
        target and masker are not co-located, the masker is equalized before
        applying the BRIR, so that both the target and masker will have the
        same average spectrum after the BRIR filtering.

        By default, the delay introduced by the BRIR is compensated for,
        such that the maxiumum of the BRIR happen simulatenously.

        Parameters
        ----------
        x : ndarray
            Mono clean speech signal of length `N`.
        m : ndarray
            Mono masker signal of length `N`.
        tdist : float
            Target distance, in meters.
        mdist : float
            Masker distance, in meters.
        align : bool
            Compensate for the delay in the BRIRs with distance  (default is
            `True`).

        Returns
        -------
        tuple of ndarray
            Mixture and noise alone, processed by the BRIRs.

        """

        if tdist not in self.dist or mdist not in self.dist:
            raise ValueError('The distance values are incorrect.')

        n_orig = x.shape[-1]

        # Filter target with BRIR only
        out_x = np.asarray([fftfilt(b, x) for b in self.brir[tdist]])

        # Equalized masker and then apply the BRIR
        if tdist == mdist:
            m = [m, m]
        else:
            # Load the equalization filter
            eqfilt_name = 't' + self._normalize_fname(tdist) + \
                        'm_m' + self._normalize_fname(mdist) + 'm.mat'

            try:
                eqfilt = sp.io.loadmat('../stimuli/crm/eqfilts/' + eqfilt_name,
                                    squeeze_me=True)
            except IOError:
                raise IOError('Cannot file file %s' % '../stimuli/crm/eqfilts/'
                            + eqfilt_name)
            m = [fftfilt(b, m) for b in [eqfilt['br'], eqfilt['bl']]]

        out_m = np.asarray([fftfilt(b, chan) for b, chan
                            in zip(self.brir[mdist], m)])

        if align:
            i_x, i_m = self._calc_aligned_idx(tdist, mdist)
        else:
            i_x = 0
            i_m = 0

        # Pad with zeros if necessary, so that the lengths stay the same
        out_x, out_m = utils.make_same_length(out_x[:, i_x:], out_m[:, i_m:])
        return out_x, out_m

    def _calc_aligned_idx(self, tdist, mdist):
        """Calculates the index of the required delay to align the max of the
        BRIRs

        Parameters
        ----------
        tdist :
            float, distance to target, in meters
        mdist :
            float, distance to masker, in meters
            :return: tuple, index of the target and masker.

        Returns
        -------

        """
        # location of earliest peak
        m_is_shortest = np.argmin([self.delays[tdist], self.delays[mdist]])
        if m_is_shortest:
            i_x = self.delays[tdist] - self.delays[mdist]
            i_m = 0
        else:
            i_x = 0
            i_m = self.delays[mdist] - self.delays[tdist]
        return i_x, i_m

def noise_from_signal(x, fs=40000, keep_env=False):
    """Create a noise with same spectrum as the input signal.

    Parameters
    ----------
    x : array_like
        Input signal.
    fs : int
         Sampling frequency of the input signal. (Default value = 40000)
    keep_env : bool
         Apply the envelope of the original signal to the noise. (Default
         value = False)

    Returns
    -------
    ndarray
        Noise signal.

    """
    x = np.asarray(x)
    n_x = x.shape[-1]
    n_fft = utils.next_pow_2(n_x)
    X = rfft(x, utils.next_pow_2(n_fft))
    # Randomize phase.
    noise_mag = np.abs(X) * np.exp(
        2 * np.pi * 1j * np.random.random(X.shape[-1]))
    noise = np.real(irfft(noise_mag, n_fft))
    out = noise[:n_x]

    if keep_env:
        env = np.abs(sp.signal.hilbert(x))
        [bb, aa] = sp.signal.butter(6, 50 / (fs / 2))  # 50 Hz LP filter
        env = sp.signal.filtfilt(bb, aa, env)
        out *= env

    return out
<commit_msg>Make sure window_len and shift_size are integers<commit_after># -*- coding: utf-8 -*-
"""
:mod:`pambox.distort` regroups various types of distortions and processings
that can be applied to signals.
"""
from __future__ import division, print_function

import numpy as np
import scipy as sp
from six.moves import zip
from scipy.io import wavfile

from pambox import utils
from pambox.utils import fftfilt

try:
    _ = np.use_fastnumpy
    from numpy.fft import fft, ifft, rfft, irfft
except AttributeError:
    from scipy.fftpack import fft, ifft
    from numpy.fft import rfft, irfft

def mix_noise(clean, noise, sent_level, snr=None):
    """Mix a signal signal noise at a given signal-to-noise ratio.

    Parameters
    ----------
    clean : ndarray
        Clean signal.
    noise : ndarray
        Noise signal.
    sent_level : float
        Sentence level, in dB SPL.
    snr :
        Signal-to-noise ratio at which to mix the signals, in dB. If snr is
        `None`,  no noise is mixed with the signal (Default value = None)

    Returns
    -------
    tuple of ndarrays
        Returns the clean signal, the mixture, and the noise.

    """

    # Pick a random section of the noise
    n_clean = len(clean)
    n_noise = len(noise)
    if n_noise > n_clean:
        start_idx = np.random.randint(n_noise - n_clean)
        noise = noise[start_idx:start_idx + n_clean]

    if snr is not None:
        # Get speech level and set noise level accordingly
        # clean_level = utils.dbspl(clean)
        # noise = utils.setdbspl(noise, clean_level - snr)
        noise = noise / utils.rms(noise) * 10 ** ((sent_level - snr) / 20)
        mix = clean + noise
    else:
        mix = clean

    return clean, mix, noise

def phase_jitter(x, a):
    """
    Apply phase jitter to a signal.

    The expression of phase jitter is:

    .. math:: y(t) = s(t) * cos(\Phi(t)),

    where :math:`\Phi(t)` is a random process uniformly distributed over
    :math:`[0, 2\pi\alpha]`. The effect of the jitter when \alpha is 0.5 or 1
    is to completely destroy the carrier signal, effictively yielding
    modulated white noise.

    Parameters
    ----------
    x : ndarray
       Signal
    a : float
        Phase jitter parameter, typically between 0 and 1, but it can be
        anything.

    Returns
    -------
    ndarray
        Processed signal of the same dimension as the input signal.

    """
    n = len(x)
    return x * np.cos(2 * np.pi * a * np.random.random_sample(n))

def reverb(x, rt):
    """
    Applies reverberation to a signal.

    Parameters
    ----------
    x : ndarray
       Input signal.
    rt : float
        Reverberation time

    Returns
    -------
    ndarray
        Processed signal.

    """
    pass

def spec_sub(x, noise, factor, w=1024 / 2., padz=1024 / 2., shift_p=0.5):
    """
    Apply spectral subtraction to a signal.

    The defaul values of the parameters are typical for a sampling frequency of
    44100 Hz. Note that (W+padz) is the final frame window and hence the fft
    length (it is normally chose as a power of 2).

    Parameters
    ----------
    x : ndarray
        Input signal
    noise :
        Input noise signal
    factor : float
        Noise subtraction factor, must be larger than 0.
    w : int
        Frame length, in samples. (Default value = 1024 / 2.)
    padz : int
        Zero padding (pad with padz/2 from the left and the right) (Default
        value = 1024 / 2.)
    shift_p : float
         Shift percentage (overlap) between each window, in fraction of the
         window size (Default value = 0.5)

    Returns
    -------
    clean_estimate : ndarray
        Estimate of the clean signal.
    noise_estimate : ndarray
        Estimate of the noisy signal.

    """
    wnd = np.hanning(w + 2)  # create hanning window with length = W
    wnd = wnd[1:-1]

    stim = np.vstack((x, noise))

    len_signal = stim.shape[-1]  # Signal length
    shift_p_indexes = np.floor(w * shift_p)
    n_segments = np.floor((len_signal - w) / shift_p_indexes + 1)
    len_segment = w + padz * 2 * shift_p
    y = np.empty((2, n_segments, len_segment))
    # Initialize arrays for spectral subtraction. Use only positive
    # frequencies.
    Y_hat = np.empty((n_segments, len_segment / 2 + 1))
    PN_hat = Y_hat.copy()

    # For each signal
    for k in range(2):
        # CUT THE APPROPRIATE SIGNAL FRAMES
        indexes = np.tile(np.arange(w), (n_segments, 1))
        index_shift = np.arange(n_segments) * shift_p_indexes
        indexes = indexes + index_shift[:, np.newaxis]
        y_tmp = stim[k]
        y_tmp = y_tmp[indexes.astype('int')] * wnd
        # PAD WITH ZEROS
        pad = np.zeros((n_segments, padz / 2))
        y_pad = np.hstack((pad, y_tmp, pad))
        y[k, :, :] = y_pad

    # FREQUENCY DOMAIN

    # signal:
    Y = fft(y[0])
    # YY = Y(1:round(end/2)+1,:); # Half window (exploit the symmetry)
    YY = Y[:, :(len_segment / 2 + 1)]  # Half window (exploit the symmetry)
    YPhase = np.angle(YY)  # Phase
    Y1 = np.abs(YY)  # Spectrum
    Y2 = Y1 ** 2  # Power Spectrum

    # noise:
    Y_N = fft(y[1])
    YY_N = Y_N[:, :(len_segment / 2 + 1)]  # Half window (exploit the symmetry)
    Y_NPhase = np.angle(YY_N)  # Phase
    Y_N1 = np.abs(YY_N)  # Spectrum
    Y_N2 = Y_N1 ** 2  # Power Spectrum

    # The noise "estimate" is simply the average of the noise power
    # spectral density in the frame:
    P_N = Y_N2.mean(axis=-1)

    Y_hat = Y2 - factor * P_N[:, np.newaxis]  # subtraction
    Y_hat = np.maximum(Y_hat, 0)  # Make the minima equal zero
    PN_hat = Y_N2 - factor * P_N[:, np.newaxis]  # subtraction for noise alone
    # PN_hat = np.maximum(PN_hat, 0)
    PN_hat[Y_hat == 0] = 0

    Y_hat[0:2, :] = 0
    PN_hat[0:2, :] = 0
    # Combining the estimated power spectrum with the original noisy phase,
    # and add the frames using an overlap-add technique
    output_Y = overlap_and_add(np.sqrt(Y_hat), YPhase, (w + padz), shift_p * w)
    output_N = overlap_and_add(np.sqrt(PN_hat.astype('complex')),
                               Y_NPhase, (w + padz), shift_p * w)

    return output_Y, output_N

def overlap_and_add(powers, phases, len_window, shift_size):
    """Reconstruct a signal with the overlap and add method.

    Parameters
    ----------
    powers : ndarray
        Magnitude of the power spectrum of the signal to reconstruct.
    phases : ndarray
        Phase of the signal to reconstruct.
    len_window : int
        Frame length, in samples.
    shift_size : int
        Shift length. For non overlapping signals, in would equal `len_window`.
        For 50% overlapping signals, it would be `len_window/2`.

    Returns
    -------
    ndarray
        Reconstructed time-domain signal.

    """
    len_window = int(len_window)
    shift_size = int(shift_size)
    n_frames, len_frame = powers.shape
    spectrum = powers * np.exp(1j * phases)
    signal = np.zeros(n_frames * shift_size + len_window - shift_size)

    # Create full spectrum, by joining conjugated positive spectrum
    if len_window % 2:
        # Do no duplicate the DC bin
        spectrum = np.hstack((spectrum, np.conj(np.fliplr(spectrum[:, 1:]))))
    else:
        # If odd-numbered, do not duplicated the DC ans FS/2 bins
        spectrum = np.hstack((spectrum,
                              np.conj(np.fliplr(spectrum[:, 1:-1]))))

    signal = np.zeros((n_frames - 1) * shift_size + len_window)

    for i_frame, hop in enumerate(range(0,
                                        len(signal) - int(len_window) + 1,
                                        int(shift_size))):
        signal[hop:hop + len_window] \
            += np.real(ifft(spectrum[i_frame], len_window))
    return signal

class Westermann_crm(object):

    """Applies HRTF and BRIR for a given target and masker distance."""

    def __init__(self, fs=40000):
        """

        Parameters
        ----------
        fs : int
             Samping frequenc of the process. (Default value = 40000)

        Returns
        -------

        """
        self.dist = np.asarray([0.5, 2, 5, 10])
        self.fs = fs
        self.brir = self._load_brirs()
        self.delays = self._find_delay()

    def _load_brirs(self):
        """Loads BRIRs from file."""
        brirs = {}
        for d in self.dist:
            d_str = self._normalize_fname(d)
            fname = '../stimuli/crm/brirs_40k/aud' \
                    + d_str + 'm.wav'
            wav = wavfile.read(fname)
            self.brir[d] = np.array(wav[1].astype('float') / 2. ** 15).T

    def _find_delay(self):
        """Calculates the delay of the direct sound, in samples."""
        delays = {}
        for k, v in self.brir.iteritems():
            x = np.mean(v, axis=0)
            delays[k] = np.abs(x).argmax()
        return delays

    @staticmethod
    def _normalize_fname(d):
        """

        Parameters
        ----------
        d : float

        Returns
        -------

        """
        if d > 1:
            d_str = str('%d' % d)
        else:
            d_str = str(d).replace('.', '')
        return d_str

    def _load_eqfilt(self, tdist, mdist):
        """
        Returns the equalization filter for the pair of target and masker.

        Parameters
        ----------
        tdist : float
            Target distance in meters. Must be in the set (0.5, 2, 5, 10).
        mdist :
            Masker distance in meters. Must be in the set (0.5, 2, 5, 10).

        Returns
        -------
        ndarray
            Equalization filter.

        """
        eqfilt_name = 't{}m_m{}m.mat'.format(self._normalize_fname(tdist),
                                             self._normalize_fname(mdist))
        eqfilt_path = '../stimuli/crm/eqfilts_{}/{}'.format(self.fs,
                                                            eqfilt_name)
        try:
            eqfilt = sp.io.loadmat(eqfilt_path, squeeze_me=True)
        except IOError:
            raise IOError('Cannot file file %s' % eqfilt_path)
        return eqfilt

    def apply(self, x, m, tdist, mdist, align=True):
        """Applies the "Westermann" distortion to a target and masker.

        Applies the BRIR of the required distance to the target. If the
        target and masker are not co-located, the masker is equalized before
        applying the BRIR, so that both the target and masker will have the
        same average spectrum after the BRIR filtering.

        By default, the delay introduced by the BRIR is compensated for,
        such that the maxiumum of the BRIR happen simulatenously.

        Parameters
        ----------
        x : ndarray
            Mono clean speech signal of length `N`.
        m : ndarray
            Mono masker signal of length `N`.
        tdist : float
            Target distance, in meters.
        mdist : float
            Masker distance, in meters.
        align : bool
            Compensate for the delay in the BRIRs with distance  (default is
            `True`).

        Returns
        -------
        tuple of ndarray
            Mixture and noise alone, processed by the BRIRs.

        """

        if tdist not in self.dist or mdist not in self.dist:
            raise ValueError('The distance values are incorrect.')

        n_orig = x.shape[-1]

        # Filter target with BRIR only
        out_x = np.asarray([fftfilt(b, x) for b in self.brir[tdist]])

        # Equalized masker and then apply the BRIR
        if tdist == mdist:
            m = [m, m]
        else:
            # Load the equalization filter
            eqfilt_name = 't' + self._normalize_fname(tdist) + \
                        'm_m' + self._normalize_fname(mdist) + 'm.mat'

            try:
                eqfilt = sp.io.loadmat('../stimuli/crm/eqfilts/' + eqfilt_name,
                                    squeeze_me=True)
            except IOError:
                raise IOError('Cannot file file %s' % '../stimuli/crm/eqfilts/'
                            + eqfilt_name)
            m = [fftfilt(b, m) for b in [eqfilt['br'], eqfilt['bl']]]

        out_m = np.asarray([fftfilt(b, chan) for b, chan
                            in zip(self.brir[mdist], m)])

        if align:
            i_x, i_m = self._calc_aligned_idx(tdist, mdist)
        else:
            i_x = 0
            i_m = 0

        # Pad with zeros if necessary, so that the lengths stay the same
        out_x, out_m = utils.make_same_length(out_x[:, i_x:], out_m[:, i_m:])
        return out_x, out_m

    def _calc_aligned_idx(self, tdist, mdist):
        """Calculates the index of the required delay to align the max of the
        BRIRs

        Parameters
        ----------
        tdist :
            float, distance to target, in meters
        mdist :
            float, distance to masker, in meters
            :return: tuple, index of the target and masker.

        Returns
        -------

        """
        # location of earliest peak
        m_is_shortest = np.argmin([self.delays[tdist], self.delays[mdist]])
        if m_is_shortest:
            i_x = self.delays[tdist] - self.delays[mdist]
            i_m = 0
        else:
            i_x = 0
            i_m = self.delays[mdist] - self.delays[tdist]
        return i_x, i_m

def noise_from_signal(x, fs=40000, keep_env=False):
    """Create a noise with same spectrum as the input signal.

    Parameters
    ----------
    x : array_like
        Input signal.
    fs : int
         Sampling frequency of the input signal. (Default value = 40000)
    keep_env : bool
         Apply the envelope of the original signal to the noise. (Default
         value = False)

    Returns
    -------
    ndarray
        Noise signal.

    """
    x = np.asarray(x)
    n_x = x.shape[-1]
    n_fft = utils.next_pow_2(n_x)
    X = rfft(x, utils.next_pow_2(n_fft))
    # Randomize phase.
    noise_mag = np.abs(X) * np.exp(
        2 * np.pi * 1j * np.random.random(X.shape[-1]))
    noise = np.real(irfft(noise_mag, n_fft))
    out = noise[:n_x]

    if keep_env:
        env = np.abs(sp.signal.hilbert(x))
        [bb, aa] = sp.signal.butter(6, 50 / (fs / 2))  # 50 Hz LP filter
        env = sp.signal.filtfilt(bb, aa, env)
        out *= env

    return out
<|endoftext|>


<commit_before>#!/usr/bin/env python

# native Python imports
import os.path
import sys
import re
import collections
from optparse import OptionParser
import sqlite3
import numpy as np

# third-party imports
import cyvcf as vcf
import pysam

# gemini modules
from ped import pedformat
import infotag
import database
import annotations
import func_impact
import severe_impact
import popgen
from compression import pack_blob

class GeminiLoader(object):
    """
    Object for creating and populating a gemini
    database and auxillary data files.
    """
    def __init__(self, args, buffer_size = 10000):
        self.args = args

        # create the gemini database
        self._create_db()
        # create a reader for the VCF file
        self.vcf_reader = self._get_vcf_reader()
        # load sample information
        self._prepare_samples()
        self.buffer_size = buffer_size
        # initialize genotype counts for each sample
        self._init_sample_gt_counts()

    def populate_from_vcf(self):
        """
        """
        self.v_id = 1
        self.var_buffer = []
        self.var_impacts_buffer = []
        buffer_count = 0
        num_samples = len(self.samples)

        # process and load each variant in the VCF file
        for var in self.vcf_reader:
            (variant, variant_impacts) = self._prepare_variation(var)
            # add the core variant info to the variant buffer
            self.var_buffer.append(variant)
            # add each of the impact for this variant (1 per gene/transcript)
            for var_impact in variant_impacts:
                self.var_impacts_buffer.append(var_impact)

            # only infer genotypes if requested
            if not self.args.noload_genotypes and not self.args.no_genotypes:
                pass

            buffer_count += 1
            # buffer full - time to insert into DB
            if buffer_count >= self.buffer_size:
                sys.stderr.write(str(self.v_id) + " variants processed.\n")
                database.insert_variation(self.c, self.var_buffer)
                database.insert_variation_impacts(self.c, self.var_impacts_buffer)
                # binary.genotypes.append(var_buffer)
                # reset for the next batch
                self.var_buffer = []
                self.var_impacts_buffer = []
                buffer_count = 0
            self.v_id += 1
        # final load to the database
        database.insert_variation(self.c, self.var_buffer)
        database.insert_variation_impacts(self.c, self.var_impacts_buffer)
        sys.stderr.write(str(self.v_id) + " variants processed.\n")

    def build_indices_and_disconnect(self):
        """
        Create the db table indicies and close up
        db connection
        """
        # index our tables for speed
        database.create_indices(self.c)
        # commit data and close up
        database.close_and_commit(self.c, self.conn)

    def _get_vcf_reader(self):
        # open the VCF file for reading
        if self.args.vcf.endswith(".gz"):
            return vcf.VCFReader(open(self.args.vcf), 'rb', compressed=True)
        else:
            return vcf.VCFReader(open(self.args.vcf), 'rb')

    def _create_db(self):
        """
        private method to open a new DB
        and create the gemini schema.
        """
        # open up a new database
        if os.path.exists(self.args.db):
            os.remove(self.args.db)
        self.conn = sqlite3.connect(self.args.db)
        self.conn.isolation_level = None
        self.c = self.conn.cursor()
        self.c.execute('PRAGMA synchronous = OFF')
        self.c.execute('PRAGMA journal_mode=MEMORY')
        # create the gemini database tables for the new DB
        database.create_tables(self.c)

    def _prepare_variation(self, var):
        """
        private method to collect metrics for
        a single variant (var) in a VCF file.
        """
        # these metric require that genotypes are present in the file
        call_rate = None
        hwe_p_value = None
        pi_hat = None
        inbreeding_coeff = None
        hom_ref = het = hom_alt = unknown = None

        # only compute certain metrics if genoypes are available
        if not self.args.no_genotypes:
            hom_ref = var.num_hom_ref
            hom_alt = var.num_hom_alt
            het = var.num_het
            unknown = var.num_unknown
            call_rate = var.call_rate
            aaf = var.aaf
            hwe_p_value, inbreeding_coeff = popgen.get_hwe_likelihood(hom_ref, het, hom_alt, aaf)
            pi_hat = var.nucl_diversity
        else:
            aaf = extract_aaf(var)

        ############################################################
        # collect annotations from gemini's custom annotation files
        ############################################################
        cyto_band    = annotations.get_cyto_info(var)
        dbsnp_info   = annotations.get_dbsnp_info(var)
        in_dbsnp     = 0 if dbsnp_info.rs_ids is None else 1
        rmsk_hits    = annotations.get_rmsk_info(var)
        in_cpg       = annotations.get_cpg_island_info(var)
        in_segdup    = annotations.get_segdup_info(var)
        is_conserved = annotations.get_conservation_info(var)
        esp          = annotations.get_esp_info(var)
        thousandG    = annotations.get_1000G_info(var)
        recomb_rate  = annotations.get_recomb_info(var)
        gms          = annotations.get_gms(var)
        grc          = annotations.get_grc(var)
        encode_tfbs  = annotations.get_encode_tfbs(var)

        # impact is a list of impacts for this variant
        impacts = None
        severe_impacts = None
        # impact terms initialized to None for handling unannotated vcf's
        # anno_id in variants is for the transcript with the most severe impact term
        affected_gene = transcript = exon = codon_change = aa_change = aa_length = biotype = consequence = effect_severity = None
        polyphen_pred = polyphen_score = sift_pred = sift_score = anno_id = None

        if self.args.anno_type is not None:
            impacts = func_impact.interpret_impact(self.args, var)
            severe_impacts = severe_impact.interpret_severe_impact(self.args, var)
            if severe_impacts:
                affected_gene = severe_impacts.gene
                transcript = severe_impacts.transcript
                exon = severe_impacts.exon
                codon_change = severe_impacts.codon_change
                aa_change = severe_impacts.aa_change
                aa_length = severe_impacts.aa_length
                biotype = severe_impacts.biotype
                consequence = severe_impacts.consequence
                effect_severity = severe_impacts.effect_severity
                polyphen_pred = severe_impacts.polyphen_pred
                polyphen_score = severe_impacts.polyphen_score
                sift_pred = severe_impacts.sift_pred
                sift_score = severe_impacts.sift_score
                anno_id = severe_impacts.anno_id

        # construct the filter string
        filter = None
        if var.FILTER is not None and var.FILTER != ".":
            if isinstance(var.FILTER, list):
                filter = ";".join(var.FILTER)
            else:
                filter = var.FILTER

        # build up numpy arrays for the genotype information.
        # these arrays will be pickled-to-binary, compressed,
        # and loaded as SqlLite BLOB values (see compression.pack_blob)
        gt_bases  = np.array(var.gt_bases, np.str)  # 'A/G', './.'
        gt_types  = np.array(var.gt_types, np.int8) # -1, 0, 1, 2
        gt_phases = np.array(var.gt_phases, np.bool) # T F F

        # tally the genotypes
        self._update_sample_gt_counts(gt_types)

        # were functional impacts predicted by SnpEFF or VEP?
        # if so, build up a row for each of the impacts / transcript
        variant_impacts = []
        is_exonic = False
        is_coding = False
        is_lof    = False
        gene      = None

        if impacts is not None:
            for idx, impact in enumerate(impacts):
                var_impact = [self.v_id, (idx+1), impact.gene,
                              impact.transcript, impact.exonic, impact.coding,
                              impact.is_lof, impact.exon, impact.codon_change, impact.aa_change,
                              impact.aa_length, impact.biotype, impact.consequence, impact.effect_severity,
                              impact.polyphen_pred, impact.polyphen_score,
                              impact.sift_pred, impact.sift_score]
                variant_impacts.append(var_impact)
                gene = impact.gene
                if impact.exonic == True: is_exonic = True
                if impact.coding == True: is_coding = True
                if impact.is_lof == True: is_lof    = True

        # construct the core variant record.
        # 1 row per variant to VARIANTS table
        chrom = var.CHROM if var.CHROM.startswith("chr") else "chr" + var.CHROM
        variant = [chrom, var.start, var.end,
                   self.v_id, anno_id, var.REF, ','.join(var.ALT),
                   var.QUAL, filter, var.var_type,
                   var.var_subtype, pack_blob(gt_bases), pack_blob(gt_types),
                   pack_blob(gt_phases), call_rate, in_dbsnp,
                   dbsnp_info.rs_ids, dbsnp_info.in_omim, dbsnp_info.clin_sig,
                   cyto_band, rmsk_hits, in_cpg,
                   in_segdup, is_conserved, hom_ref, het,
                   hom_alt, unknown, aaf,
                   hwe_p_value, inbreeding_coeff, pi_hat,
                   recomb_rate, gene, affected_gene, transcript,
                   is_exonic, is_coding, is_lof, exon, codon_change,
                   aa_change, aa_length, biotype, consequence, effect_severity,
                   polyphen_pred, polyphen_score, sift_pred, sift_score,
                   infotag.get_depth(var), infotag.get_strand_bias(var), infotag.get_rms_map_qual(var),
                   infotag.get_homopol_run(var), infotag.get_map_qual_zero(var), infotag.get_num_of_alleles(var),
                   infotag.get_frac_dels(var), infotag.get_haplotype_score(var), infotag.get_quality_by_depth(var),
                   infotag.get_allele_count(var), infotag.get_allele_bal(var), esp.found, esp.aaf_EA, esp.aaf_AA,
                   esp.aaf_ALL, esp.exome_chip, thousandG.found, thousandG.aaf_AMR, thousandG.aaf_ASN, thousandG.aaf_AFR,
                   thousandG.aaf_EUR, thousandG.aaf_ALL, grc, gms.illumina, gms.solid, gms.iontorrent, encode_tfbs]
        return variant, variant_impacts

    def _prepare_samples(self):
        """
        private method to load sample information
        """
        if not self.args.no_genotypes:
            self.samples = self.vcf_reader.samples
            self.sample_to_id = {}
            for idx, sample in enumerate(self.samples):
                self.sample_to_id[sample] = idx + 1

        self.ped_hash = {}
        if self.args.ped_file is not None:
            for line in open(self.args.ped_file, 'r'):
                field = line.strip().split("\t")
                if len(field) > 1 and not field[0].startswith("#"):
                    ped = pedformat(field)
                    self.ped_hash[ped.name] = ped

        sample_list = []
        for sample in self.samples:
            i = self.sample_to_id[sample]
            if self.ped_hash.has_key(sample):
                ped = self.ped_hash[sample]
                sample_list = [i, sample, ped.family, ped.paternal,
                               ped.maternal, ped.sex, ped.phenotype, ped.ethnicity]
            else:
                sample_list = [i, sample, None, None, None, None, None, None]
            database.insert_sample(self.c, sample_list)

    def _init_sample_gt_counts(self):
        """
        Initialize a 2D array of counts for tabulating
        the count of each genotype type for eaxh sample.

        The first dimension is one bucket for each sample.
        The second dimension (size=4) is a count for each gt type.
           Index 0 == # of hom_ref genotypes for the sample
           Index 1 == # of het genotypes for the sample
           Index 2 == # of missing genotypes for the sample
           Index 3 == # of hom_alt genotypes for the sample
        """
        self.sample_gt_counts = np.array(np.zeros( (len(self.samples), 4) ),
                                         dtype='uint32')

    def _update_sample_gt_counts(self, gt_types):
        """
        Update the count of each gt type for each sample
        """
        for idx, gt_type in enumerate(gt_types):
            self.sample_gt_counts[idx][gt_type] += 1

    def store_sample_gt_counts(self):
        """
        Update the count of each gt type for each sample
        """
        self.c.execute("BEGIN TRANSACTION")
        for idx, gt_counts in enumerate(self.sample_gt_counts):
            self.c.execute("""insert into sample_genotype_counts values \
                            (?,?,?,?,?)""",
                            [idx,
                            int(gt_counts[0]),  # hom_ref
                            int(gt_counts[1]),  # het
                            int(gt_counts[3]),  # hom_alt
                            int(gt_counts[2])]) # missing
        self.c.execute("END")

def load(parser, args):
    if (args.db is None or args.vcf is None):
        parser.print_help()
        exit("ERROR: load needs both a VCF file and a database file\n")
    if args.anno_type not in ['snpEff', 'VEP', None]:
        parser.print_help()
        exit("\nERROR: Unsupported selection for -t\n")

    # collect of the the add'l annotation files
    annotations.load_annos()

    # create a new gemini loader and populate
    # the gemini db and files from the VCF
    gemini_loader = GeminiLoader(args)
    gemini_loader.populate_from_vcf()
    gemini_loader.build_indices_and_disconnect()
    gemini_loader.store_sample_gt_counts()

<commit_msg>changes made to incorporate database changes<commit_after>#!/usr/bin/env python

# native Python imports
import os.path
import sys
import re
import collections
from optparse import OptionParser
import sqlite3
import numpy as np

# third-party imports
import cyvcf as vcf
import pysam

# gemini modules
from ped import pedformat
import infotag
import database
import annotations
import func_impact
import severe_impact
import popgen
from compression import pack_blob

class GeminiLoader(object):
    """
    Object for creating and populating a gemini
    database and auxillary data files.
    """
    def __init__(self, args, buffer_size = 10000):
        self.args = args

        # create the gemini database
        self._create_db()
        # create a reader for the VCF file
        self.vcf_reader = self._get_vcf_reader()
        # load sample information
        self._prepare_samples()
        self.buffer_size = buffer_size
        # initialize genotype counts for each sample
        self._init_sample_gt_counts()

    def populate_from_vcf(self):
        """
        """
        self.v_id = 1
        self.var_buffer = []
        self.var_impacts_buffer = []
        buffer_count = 0
        num_samples = len(self.samples)

        # process and load each variant in the VCF file
        for var in self.vcf_reader:
            (variant, variant_impacts) = self._prepare_variation(var)
            # add the core variant info to the variant buffer
            self.var_buffer.append(variant)
            # add each of the impact for this variant (1 per gene/transcript)
            for var_impact in variant_impacts:
                self.var_impacts_buffer.append(var_impact)

            # only infer genotypes if requested
            if not self.args.noload_genotypes and not self.args.no_genotypes:
                pass

            buffer_count += 1
            # buffer full - time to insert into DB
            if buffer_count >= self.buffer_size:
                sys.stderr.write(str(self.v_id) + " variants processed.\n")
                database.insert_variation(self.c, self.var_buffer)
                database.insert_variation_impacts(self.c, self.var_impacts_buffer)
                # binary.genotypes.append(var_buffer)
                # reset for the next batch
                self.var_buffer = []
                self.var_impacts_buffer = []
                buffer_count = 0
            self.v_id += 1
        # final load to the database
        database.insert_variation(self.c, self.var_buffer)
        database.insert_variation_impacts(self.c, self.var_impacts_buffer)
        sys.stderr.write(str(self.v_id) + " variants processed.\n")

    def build_indices_and_disconnect(self):
        """
        Create the db table indicies and close up
        db connection
        """
        # index our tables for speed
        database.create_indices(self.c)
        # commit data and close up
        database.close_and_commit(self.c, self.conn)

    def _get_vcf_reader(self):
        # open the VCF file for reading
        if self.args.vcf.endswith(".gz"):
            return vcf.VCFReader(open(self.args.vcf), 'rb', compressed=True)
        else:
            return vcf.VCFReader(open(self.args.vcf), 'rb')

    def _create_db(self):
        """
        private method to open a new DB
        and create the gemini schema.
        """
        # open up a new database
        if os.path.exists(self.args.db):
            os.remove(self.args.db)
        self.conn = sqlite3.connect(self.args.db)
        self.conn.isolation_level = None
        self.c = self.conn.cursor()
        self.c.execute('PRAGMA synchronous = OFF')
        self.c.execute('PRAGMA journal_mode=MEMORY')
        # create the gemini database tables for the new DB
        database.create_tables(self.c)

    def _prepare_variation(self, var):
        """
        private method to collect metrics for
        a single variant (var) in a VCF file.
        """
        # these metric require that genotypes are present in the file
        call_rate = None
        hwe_p_value = None
        pi_hat = None
        inbreeding_coeff = None
        hom_ref = het = hom_alt = unknown = None

        # only compute certain metrics if genoypes are available
        if not self.args.no_genotypes:
            hom_ref = var.num_hom_ref
            hom_alt = var.num_hom_alt
            het = var.num_het
            unknown = var.num_unknown
            call_rate = var.call_rate
            aaf = var.aaf
            hwe_p_value, inbreeding_coeff = popgen.get_hwe_likelihood(hom_ref, het, hom_alt, aaf)
            pi_hat = var.nucl_diversity
        else:
            aaf = extract_aaf(var)

        ############################################################
        # collect annotations from gemini's custom annotation files
        ############################################################
        cyto_band    = annotations.get_cyto_info(var)
        dbsnp_info   = annotations.get_dbsnp_info(var)
        in_dbsnp     = 0 if dbsnp_info.rs_ids is None else 1
        rmsk_hits    = annotations.get_rmsk_info(var)
        in_cpg       = annotations.get_cpg_island_info(var)
        in_segdup    = annotations.get_segdup_info(var)
        is_conserved = annotations.get_conservation_info(var)
        esp          = annotations.get_esp_info(var)
        thousandG    = annotations.get_1000G_info(var)
        recomb_rate  = annotations.get_recomb_info(var)
        gms          = annotations.get_gms(var)
        grc          = annotations.get_grc(var)
        encode_tfbs  = annotations.get_encode_tfbs(var)

        # impact is a list of impacts for this variant
        impacts = None
        severe_impacts = None
        # impact terms initialized to None for handling unannotated vcf's
        # anno_id in variants is for the transcript with the most severe impact term
        gene = transcript = exon = codon_change = aa_change = aa_length = biotype = consequence = effect_severity = None
        polyphen_pred = polyphen_score = sift_pred = sift_score = anno_id = None

        if self.args.anno_type is not None:
            impacts = func_impact.interpret_impact(self.args, var)
            severe_impacts = severe_impact.interpret_severe_impact(self.args, var)
            if severe_impacts:
                gene = severe_impacts.gene
                transcript = severe_impacts.transcript
                exon = severe_impacts.exon
                codon_change = severe_impacts.codon_change
                aa_change = severe_impacts.aa_change
                aa_length = severe_impacts.aa_length
                biotype = severe_impacts.biotype
                consequence = severe_impacts.consequence
                effect_severity = severe_impacts.effect_severity
                polyphen_pred = severe_impacts.polyphen_pred
                polyphen_score = severe_impacts.polyphen_score
                sift_pred = severe_impacts.sift_pred
                sift_score = severe_impacts.sift_score
                anno_id = severe_impacts.anno_id

        # construct the filter string
        filter = None
        if var.FILTER is not None and var.FILTER != ".":
            if isinstance(var.FILTER, list):
                filter = ";".join(var.FILTER)
            else:
                filter = var.FILTER

        # build up numpy arrays for the genotype information.
        # these arrays will be pickled-to-binary, compressed,
        # and loaded as SqlLite BLOB values (see compression.pack_blob)
        gt_bases  = np.array(var.gt_bases, np.str)  # 'A/G', './.'
        gt_types  = np.array(var.gt_types, np.int8) # -1, 0, 1, 2
        gt_phases = np.array(var.gt_phases, np.bool) # T F F

        # tally the genotypes
        self._update_sample_gt_counts(gt_types)

        # were functional impacts predicted by SnpEFF or VEP?
        # if so, build up a row for each of the impacts / transcript
        variant_impacts = []
        is_exonic = False
        is_coding = False
        is_lof    = False

        if impacts is not None:
            for idx, impact in enumerate(impacts):
                var_impact = [self.v_id, (idx+1), impact.gene,
                              impact.transcript, impact.exonic, impact.coding,
                              impact.is_lof, impact.exon, impact.codon_change, impact.aa_change,
                              impact.aa_length, impact.biotype, impact.consequence, impact.effect_severity,
                              impact.polyphen_pred, impact.polyphen_score,
                              impact.sift_pred, impact.sift_score]
                variant_impacts.append(var_impact)
                if impact.exonic == True: is_exonic = True
                if impact.coding == True: is_coding = True
                if impact.is_lof == True: is_lof    = True

        # construct the core variant record.
        # 1 row per variant to VARIANTS table
        chrom = var.CHROM if var.CHROM.startswith("chr") else "chr" + var.CHROM
        variant = [chrom, var.start, var.end,
                   self.v_id, anno_id, var.REF, ','.join(var.ALT),
                   var.QUAL, filter, var.var_type,
                   var.var_subtype, pack_blob(gt_bases), pack_blob(gt_types),
                   pack_blob(gt_phases), call_rate, in_dbsnp,
                   dbsnp_info.rs_ids, dbsnp_info.in_omim, dbsnp_info.clin_sig,
                   cyto_band, rmsk_hits, in_cpg,
                   in_segdup, is_conserved, hom_ref, het,
                   hom_alt, unknown, aaf,
                   hwe_p_value, inbreeding_coeff, pi_hat,
                   recomb_rate, gene, transcript,
                   is_exonic, is_coding, is_lof, exon, codon_change,
                   aa_change, aa_length, biotype, consequence, effect_severity,
                   polyphen_pred, polyphen_score, sift_pred, sift_score,
                   infotag.get_depth(var), infotag.get_strand_bias(var), infotag.get_rms_map_qual(var),
                   infotag.get_homopol_run(var), infotag.get_map_qual_zero(var), infotag.get_num_of_alleles(var),
                   infotag.get_frac_dels(var), infotag.get_haplotype_score(var), infotag.get_quality_by_depth(var),
                   infotag.get_allele_count(var), infotag.get_allele_bal(var), esp.found, esp.aaf_EA, esp.aaf_AA,
                   esp.aaf_ALL, esp.exome_chip, thousandG.found, thousandG.aaf_AMR, thousandG.aaf_ASN, thousandG.aaf_AFR,
                   thousandG.aaf_EUR, thousandG.aaf_ALL, grc, gms.illumina, gms.solid, gms.iontorrent, encode_tfbs]
        return variant, variant_impacts

    def _prepare_samples(self):
        """
        private method to load sample information
        """
        if not self.args.no_genotypes:
            self.samples = self.vcf_reader.samples
            self.sample_to_id = {}
            for idx, sample in enumerate(self.samples):
                self.sample_to_id[sample] = idx + 1

        self.ped_hash = {}
        if self.args.ped_file is not None:
            for line in open(self.args.ped_file, 'r'):
                field = line.strip().split("\t")
                if len(field) > 1 and not field[0].startswith("#"):
                    ped = pedformat(field)
                    self.ped_hash[ped.name] = ped

        sample_list = []
        for sample in self.samples:
            i = self.sample_to_id[sample]
            if self.ped_hash.has_key(sample):
                ped = self.ped_hash[sample]
                sample_list = [i, sample, ped.family, ped.paternal,
                               ped.maternal, ped.sex, ped.phenotype, ped.ethnicity]
            else:
                sample_list = [i, sample, None, None, None, None, None, None]
            database.insert_sample(self.c, sample_list)

    def _init_sample_gt_counts(self):
        """
        Initialize a 2D array of counts for tabulating
        the count of each genotype type for eaxh sample.

        The first dimension is one bucket for each sample.
        The second dimension (size=4) is a count for each gt type.
           Index 0 == # of hom_ref genotypes for the sample
           Index 1 == # of het genotypes for the sample
           Index 2 == # of missing genotypes for the sample
           Index 3 == # of hom_alt genotypes for the sample
        """
        self.sample_gt_counts = np.array(np.zeros( (len(self.samples), 4) ),
                                         dtype='uint32')

    def _update_sample_gt_counts(self, gt_types):
        """
        Update the count of each gt type for each sample
        """
        for idx, gt_type in enumerate(gt_types):
            self.sample_gt_counts[idx][gt_type] += 1

    def store_sample_gt_counts(self):
        """
        Update the count of each gt type for each sample
        """
        self.c.execute("BEGIN TRANSACTION")
        for idx, gt_counts in enumerate(self.sample_gt_counts):
            self.c.execute("""insert into sample_genotype_counts values \
                            (?,?,?,?,?)""",
                            [idx,
                            int(gt_counts[0]),  # hom_ref
                            int(gt_counts[1]),  # het
                            int(gt_counts[3]),  # hom_alt
                            int(gt_counts[2])]) # missing
        self.c.execute("END")

def load(parser, args):
    if (args.db is None or args.vcf is None):
        parser.print_help()
        exit("ERROR: load needs both a VCF file and a database file\n")
    if args.anno_type not in ['snpEff', 'VEP', None]:
        parser.print_help()
        exit("\nERROR: Unsupported selection for -t\n")

    # collect of the the add'l annotation files
    annotations.load_annos()

    # create a new gemini loader and populate
    # the gemini db and files from the VCF
    gemini_loader = GeminiLoader(args)
    gemini_loader.populate_from_vcf()
    gemini_loader.build_indices_and_disconnect()
    gemini_loader.store_sample_gt_counts()

<|endoftext|>