<commit_before>from __future__ import division
from utils.filters import calc_lapgauss
import SimpleITK as sitk
import numpy as np
from utils.preprocess_utils import homogenize_intensity_n4
from utils.preprocess_utils import convert_positive, estimate_background_prc
from utils.preprocess_utils import resize_img
from utils.preprocess_utils import histogram_matching, wavelet_subtraction_hazen
from utils.filters import adaptive_thresh
from utils.cp_functions import align_cross_correlation, align_mutual_information
from scipy.ndimage import imread
from glob import glob
from utils.filters import interpolate_nan
from scipy.optimize import minimize
import logging
from utils.global_holder import holder
from utils.mi_align import calc_jitters_multiple, calc_crop_coordinates
from utils.shading_correction import retrieve_ff_ref

logger = logging.getLogger(__name__)
np.random.seed(0)

def gaussian_laplace(img, SIGMA=2.5, NEG=False):
    if NEG:
        img = -calc_lapgauss(img, SIGMA)
        img[img < 0 ] = 0
        return img
    return calc_lapgauss(img, SIGMA)

def curvature_anisotropic_smooth(img, NITER=10):
    fil = sitk.CurvatureAnisotropicDiffusionImageFilter()
    fil.SetNumberOfIterations(NITER)
    simg = sitk.GetImageFromArray(img.astype(np.float32))
    sres = fil.Execute(simg)
    return sitk.GetArrayFromImage(sres)

def background_subtraction_wavelet_hazen(img, THRES=100, ITER=5, WLEVEL=6, OFFSET=50):
    """Wavelet background subtraction.
    """
    back = wavelet_subtraction_hazen(img, ITER=ITER, THRES=THRES, WLEVEL=WLEVEL)
    img = img - back
    return convert_positive(img, OFFSET)

def n4_illum_correction(img, RATIO=1.5, FILTERINGSIZE=50):
    """
    Implementation of the N4 bias field correction algorithm.
    Takes some calculation time. It first calculates the background using adaptive_thesh.
    """
    bw = adaptive_thresh(img, R=RATIO, FILTERINGSIZE=FILTERINGSIZE)
    img = homogenize_intensity_n4(img, -bw)
    return img

def n4_illum_correction_downsample(img, DOWN=2, RATIO=1.05, FILTERINGSIZE=50, OFFSET=10):
    """Faster but more insensitive to local illum bias.
    """
    fil = sitk.ShrinkImageFilter()
    cc = sitk.GetArrayFromImage(fil.Execute(sitk.GetImageFromArray(img), [DOWN, DOWN]))
    bw = adaptive_thresh(cc, R=RATIO, FILTERINGSIZE=FILTERINGSIZE/DOWN)
    himg = homogenize_intensity_n4(cc, -bw)
    himg = cc - himg
    # himg[himg < 0] = 0
    bias = resize_img(himg, img.shape)
    img = img - bias
    return convert_positive(img, OFFSET)

def align(img, CROP=0.05):
    """
    CROP (float): crop images beforehand. When set to 0.05, 5% of each edges are cropped.
    """
    if not hasattr(holder, "align"):
        if isinstance(holder.inputs[0], list) or isinstance(holder.inputs[0], tuple):
            inputs = [i[0] for i in holder.inputs]
        else:
            inputs = holder.inputs

        img0 = imread(inputs[0])

        (ch, cw) = [int(CROP * i) for i in img0.shape]
        ch = None if ch == 0 else ch
        cw = None if cw == 0 else cw

        jitters = calc_jitters_multiple(inputs, ch, cw)
        holder.align = calc_crop_coordinates(jitters, img0.shape)
        logger.debug('holder.align set to {0}'.format(holder.align))
    jt = holder.align[holder.frame]
    logger.debug('Jitter: {0}'.format(jt))
    if img.ndim == 2:
        return img[jt[0]:jt[1], jt[2]:jt[3]]
    if img.ndim == 3:
        return img[jt[0]:jt[1], jt[2]:jt[3], :]

def flatfield_references(img, ff_paths=['Pos0/img00.tif', 'Pos1/img01.tif'], exp_corr=False):
    """
    Use empty images for background subtraction and illumination bias correction.
    Given multiple reference images, it will calculate median profile and use it for subtraction.
    If flatfield image has the same illumination pattern but different exposure to the img,
    turning on bg_align would calculate correction factor.

    ff_paths (str or List(str)): image path for flat fielding references.
                                 It can be single, multiple or path with wildcards.
        e.g.    ff_paths = "FF/img_000000000_YFP*"
                ff_paths = ["FF/img_01.tif", "FF/img_02.tif"]

    """
    store = []
    if isinstance(ff_paths, str):
        ff_paths = [ff_paths, ]
    for i in ff_paths:
        for ii in glob(i):
            store.append(ii)
    ff_store = []
    for path in store:
        ff_store.append(imread(path))
    ff = np.median(np.dstack(ff_store), axis=2)

    if exp_corr:
        """If a reference is taken at different exposure, or exposure is not stable over time,
        this will try to correct for it. Majority of image needs to be a backrgound.
        """
        def minimize_bg(img, ff, corr, perctile=50, weight=10):
            thres = np.percentile(img, perctile)
            res = img - corr * ff
            res = res[res < thres]
            return np.sum(res[res>0]) - weight * np.sum(res[res < 0])
        """avoid having negative values yet suppress positive values in background region.
        """
        func = lambda x: minimize_bg(img, ff, x)
        if not hasattr(holder, 'bg_corr'):
            holder.bg_corr = 1.0
        ret = minimize(func, x0=holder.bg_corr, bounds=((0, None),))
        holder.bg_corr = ret.x
        ff = ret.x * ff

    img = img - ff
    img[img < 0] = np.nan
    img = interpolate_nan(img)
    return img

def histogram_match(img, BINS=1000, QUANT=100, THRES=False):
    """
    If an optical system is not stable and shows global intensity changes over time,
    use this method to correct for it. Typically use for nuclear marker, where
    intensity and its patterns should be stable over time.
    """
    if holder.frame == 0:
        holder.first_img = img
    else:
        img = histogram_matching(img, holder.first_img, BINS, QUANT, THRES)
    return img

def shading_correction(img,
                       ch='DAPI',
                       whiteurl='http://archive.simtk.org/ktrprotocol/temp/ffref_20x3bin.npz',
                       darkurl='http://archive.simtk.org/ktrprotocol/temp/ffdarkref_20x3bin.npz'):

    """only access if it has not"""
    if hasattr(holder, 'sc_ref'):
        if ch in holder.sc_ref:
            ref = holder.sc_ref[ch]
            darkref = holder.sc_dref[ch]
        else:
            ref, darkref = retrieve_ff_ref(whiteurl, darkurl)
            holder.sc_ref[ch] = ref
            holder.sc_dref[ch] = darkref
    else:
        ref, darkref = retrieve_ff_ref(whiteurl, darkurl)
        holder.sc_ref, holder.sc_dref = {}, {}
        holder.sc_ref[ch] = ref
        holder.sc_dref[ch] = darkref

    def correct_shade(img, ref, darkref, ch):
        img = img.astype(np.float)
        d0 = img.astype(np.float) - darkref[ch]
        d1 = ref[ch] - darkref[ch]
        d1[d1 < 0] = 0
        return d1.mean() * d0/d1

    img = correct_shade(img, ref, darkref, ch)
    return img
<commit_msg>wavelet subtraction using pywt added<commit_after>from __future__ import division
from utils.filters import calc_lapgauss
import SimpleITK as sitk
import numpy as np
from utils.preprocess_utils import homogenize_intensity_n4
from utils.preprocess_utils import convert_positive, estimate_background_prc
from utils.preprocess_utils import resize_img
from utils.preprocess_utils import histogram_matching, wavelet_subtraction_hazen
from utils.filters import adaptive_thresh
from utils.cp_functions import align_cross_correlation, align_mutual_information
from scipy.ndimage import imread
from glob import glob
from utils.filters import interpolate_nan
from scipy.optimize import minimize
import logging
from utils.global_holder import holder
from utils.mi_align import calc_jitters_multiple, calc_crop_coordinates
from utils.shading_correction import retrieve_ff_ref

logger = logging.getLogger(__name__)
np.random.seed(0)

def gaussian_laplace(img, SIGMA=2.5, NEG=False):
    if NEG:
        img = -calc_lapgauss(img, SIGMA)
        img[img < 0 ] = 0
        return img
    return calc_lapgauss(img, SIGMA)

def curvature_anisotropic_smooth(img, NITER=10):
    fil = sitk.CurvatureAnisotropicDiffusionImageFilter()
    fil.SetNumberOfIterations(NITER)
    simg = sitk.GetImageFromArray(img.astype(np.float32))
    sres = fil.Execute(simg)
    return sitk.GetArrayFromImage(sres)

def background_subtraction_wavelet_hazen(img, THRES=100, ITER=5, WLEVEL=6, OFFSET=50):
    """Wavelet background subtraction.
    """
    back = wavelet_subtraction_hazen(img, ITER=ITER, THRES=THRES, WLEVEL=WLEVEL)
    img = img - back
    return convert_positive(img, OFFSET)

def n4_illum_correction(img, RATIO=1.5, FILTERINGSIZE=50):
    """
    Implementation of the N4 bias field correction algorithm.
    Takes some calculation time. It first calculates the background using adaptive_thesh.
    """
    bw = adaptive_thresh(img, R=RATIO, FILTERINGSIZE=FILTERINGSIZE)
    img = homogenize_intensity_n4(img, ~bw)
    return img

def n4_illum_correction_downsample(img, DOWN=2, RATIO=1.05, FILTERINGSIZE=50, OFFSET=10):
    """Faster but more insensitive to local illum bias.
    """
    fil = sitk.ShrinkImageFilter()
    cc = sitk.GetArrayFromImage(fil.Execute(sitk.GetImageFromArray(img), [DOWN, DOWN]))
    bw = adaptive_thresh(cc, R=RATIO, FILTERINGSIZE=FILTERINGSIZE/DOWN)
    himg = homogenize_intensity_n4(cc, ~bw)
    himg = cc - himg
    # himg[himg < 0] = 0
    bias = resize_img(himg, img.shape)
    img = img - bias
    return convert_positive(img, OFFSET)

def align(img, CROP=0.05):
    """
    CROP (float): crop images beforehand. When set to 0.05, 5% of each edges are cropped.
    """
    if not hasattr(holder, "align"):
        if isinstance(holder.inputs[0], list) or isinstance(holder.inputs[0], tuple):
            inputs = [i[0] for i in holder.inputs]
        else:
            inputs = holder.inputs

        img0 = imread(inputs[0])

        (ch, cw) = [int(CROP * i) for i in img0.shape]
        ch = None if ch == 0 else ch
        cw = None if cw == 0 else cw

        jitters = calc_jitters_multiple(inputs, ch, cw)
        holder.align = calc_crop_coordinates(jitters, img0.shape)
        logger.debug('holder.align set to {0}'.format(holder.align))
    jt = holder.align[holder.frame]
    logger.debug('Jitter: {0}'.format(jt))
    if img.ndim == 2:
        return img[jt[0]:jt[1], jt[2]:jt[3]]
    if img.ndim == 3:
        return img[jt[0]:jt[1], jt[2]:jt[3], :]

def flatfield_references(img, ff_paths=['Pos0/img00.tif', 'Pos1/img01.tif'], exp_corr=False):
    """
    Use empty images for background subtraction and illumination bias correction.
    Given multiple reference images, it will calculate median profile and use it for subtraction.
    If flatfield image has the same illumination pattern but different exposure to the img,
    turning on bg_align would calculate correction factor.

    ff_paths (str or List(str)): image path for flat fielding references.
                                 It can be single, multiple or path with wildcards.
        e.g.    ff_paths = "FF/img_000000000_YFP*"
                ff_paths = ["FF/img_01.tif", "FF/img_02.tif"]

    """
    store = []
    if isinstance(ff_paths, str):
        ff_paths = [ff_paths, ]
    for i in ff_paths:
        for ii in glob(i):
            store.append(ii)
    ff_store = []
    for path in store:
        ff_store.append(imread(path))
    ff = np.median(np.dstack(ff_store), axis=2)

    if exp_corr:
        """If a reference is taken at different exposure, or exposure is not stable over time,
        this will try to correct for it. Majority of image needs to be a backrgound.
        """
        def minimize_bg(img, ff, corr, perctile=50, weight=10):
            thres = np.percentile(img, perctile)
            res = img - corr * ff
            res = res[res < thres]
            return np.sum(res[res>0]) - weight * np.sum(res[res < 0])
        """avoid having negative values yet suppress positive values in background region.
        """
        func = lambda x: minimize_bg(img, ff, x)
        if not hasattr(holder, 'bg_corr'):
            holder.bg_corr = 1.0
        ret = minimize(func, x0=holder.bg_corr, bounds=((0, None),))
        holder.bg_corr = ret.x
        ff = ret.x * ff

    img = img - ff
    img[img < 0] = np.nan
    img = interpolate_nan(img)
    return img

def histogram_match(img, BINS=1000, QUANT=100, THRES=False):
    """
    If an optical system is not stable and shows global intensity changes over time,
    use this method to correct for it. Typically use for nuclear marker, where
    intensity and its patterns should be stable over time.
    """
    if holder.frame == 0:
        holder.first_img = img
    else:
        img = histogram_matching(img, holder.first_img, BINS, QUANT, THRES)
    return img

def shading_correction(img,
                       ch='DAPI',
                       whiteurl='http://archive.simtk.org/ktrprotocol/temp/ffref_20x3bin.npz',
                       darkurl='http://archive.simtk.org/ktrprotocol/temp/ffdarkref_20x3bin.npz'):

    """only access if it has not"""
    if hasattr(holder, 'sc_ref'):
        if ch in holder.sc_ref:
            ref = holder.sc_ref[ch]
            darkref = holder.sc_dref[ch]
        else:
            ref, darkref = retrieve_ff_ref(whiteurl, darkurl)
            holder.sc_ref[ch] = ref
            holder.sc_dref[ch] = darkref
    else:
        ref, darkref = retrieve_ff_ref(whiteurl, darkurl)
        holder.sc_ref, holder.sc_dref = {}, {}
        holder.sc_ref[ch] = ref
        holder.sc_dref[ch] = darkref

    def correct_shade(img, ref, darkref, ch):
        img = img.astype(np.float)
        d0 = img.astype(np.float) - darkref[ch]
        d1 = ref[ch] - darkref[ch]
        d1[d1 < 0] = 0
        return d1.mean() * d0/d1

    img = correct_shade(img, ref, darkref, ch)
    return img

def background_subtraction_wavelet(img, level=7, OFFSET=10):
    '''
    It might be radical but works in many cases in terms of segmentation.
    Use "background_subtraction_wavelet_hazen" for a proper implementation.
    '''
    from pywt import WaveletPacket2D
    from skimage.transform import resize
    def wavelet_subtraction(img, level):
        """6- 7 level is recommended"""
        if level == 0:
            return img
        wp = WaveletPacket2D(data=img.astype(np.uint16), wavelet='haar', mode='sym')
        back = resize(np.array(wp['a'*level].data), img.shape, order=3, mode='reflect')/(2**level)
        img = img - back
        return img
    img = wavelet_subtraction(img, level)
    return convert_positive(img, OFFSET)
<|endoftext|>


<commit_before>import heapq

from chainer import function
from chainer import variable

class DotNode(object):
    """Node of the computational graph, with utilities for dot language.

    This class represents a node of computational graph,
    with some utilities for dot language.

    """
    def __init__(self, node, attribute=None):
        """Initializes DotNode.

        Args:
            node: :class: `Variable` object or :class: `Function` object.
            attribute (dict): Attributes for the node.

        """
        assert isinstance(node, (variable.Variable, function.Function))
        self.node = node
        self.id_ = id(node)
        if isinstance(node, variable.Variable):
            self.attribute = {'shape': 'oval'}
        else:
            self.attribute = {'shape': 'box'}
        if self.attribute is not None:
            self.attribute.update(attribute)
        self.attribute['label'] = node.label

    @property
    def label(self):
        """The text that represents properties of the node.

        Returns:
            string: The text that represents the id and attributes of this
                node.
        """

        attributes = ["%s=\"%s\"" % (k, v) for (k, v)
                      in self.attribute.items()]
        return "%s [%s];" % (self.id_, ",".join(attributes))

class ComputationalGraph(object):

    """Class that represents computational graph.

    .. note::

      We assume that the computational graph is directed and acyclic.

    """
    def __init__(self, nodes, edges, variable_style=None, function_style=None):
        """Initializes computational graph.

        Args:
            nodes (list): List of nodes. Each node is either
                 :class:`Variable` object or :class:`Function` object.
            edges (list): List of edges. Each edge consists of pair of nodes.
            variable_style (dist): Dot node style for variable.
            function_style (dist): Dot node style for function.

        """
        self.nodes = nodes
        self.edges = edges
        self.variable_style = variable_style
        self.function_style = function_style

    def _to_dot(self):
        """Converts graph in dot format.

        `label` property of is used as short description of each node.
        Returns:
            str: The graph in dot format.

        """
        ret = "digraph graphname{"
        for node in self.nodes:
            assert isinstance(node, (variable.Variable, function.Function))
            if isinstance(node, variable.Variable):
                ret += DotNode(node, self.variable_style).label
            else:
                ret += DotNode(node, self.function_style).label
        for edge in self.edges:
            head, tail = edge
            if (isinstance(head, variable.Variable) and
                    isinstance(tail, function.Function)):
                head_attr = self.variable_style
                tail_attr = self.function_style
            elif (isinstance(head, function.Function) and
                  isinstance(tail, variable.Variable)):
                head_attr = self.function_style
                tail_attr = self.variable_style
            else:
                raise TypeError(
                    'head and tail should be the set of Variable and Function')
            head_node = DotNode(head, head_attr)
            tail_node = DotNode(tail, tail_attr)
            ret += "%s -> %s;" % (head_node.id_, tail_node.id_)
        ret += "}"
        return ret

    def dump(self, format='dot'):
        """Dumps graph as a text.

        Args
            format(str): The graph language name of the output.
            Currently, it must be 'dot'.

        Returns
            str: The graph in specified format.

        """
        if format == 'dot':
            return self._to_dot()
        else:
            NotImplementedError('Currently, only dot format is supported.')

def build_computational_graph(outputs, remove_split=True,
                              variable_style=None, function_style=None):
    """Builds a graph of functions and variables backward-reachable from outputs.

    Args:
        outputs(list): nodes from which the graph is constructed.
            Each element of outputs must be either :class:`Variable`
            object or :class:`Function` object.
        remove_split(bool): It must be ``True``. This argument is left for
            backward compatibility.
        variable_style(dict): Dot node style for variable.
            Possible keys are 'shape', 'color', 'fillcolor', 'style', and etc.
        function_style(dict): Dot node style for function.

    Returns:
        ComputationalGraph: A graph consisting of nodes and edges that
        are backward-reachable from at least one of ``outputs``.

        If ``unchain_backward`` was called in some variable in the
        computational graph before this function, backward step is
        stopped at this variable.

        For example, suppose that computational graph is as follows::

                |--> f ---> y
            x --+
                |--> g ---> z

        Let ``outputs = [y, z]``.
        Then the full graph is emitted.

        Next, let ``outputs = [y]``. Note that ``z`` and ``g``
        are not backward-reachable from ``y``.
        The resulting graph would be following::

            x ---> f ---> y

        See :class:`TestGraphBuilder` for details.

    """
    if not remove_split:
        raise ValueError('remove_split=False is not supported anymore')

    cands = []
    seen_edges = set()
    nodes = set()
    push_count = [0]

    # This class is for object that has not been implemented __eq__
    class HashableObject(object):

        def __init__(self, v):
            self.v = v

        def __hash__(self):
            return self.v.__hash__()

        def __eq__(self, r):
            return self.v is r.v

    def add_cand(cand):
        heapq.heappush(cands, (-cand.rank, push_count[0], cand))
        push_count[0] += 1

    for o in outputs:
        add_cand(o)
        nodes.add(HashableObject(o))

    while cands:
        _, _, cand = heapq.heappop(cands)
        if isinstance(cand, variable.Variable):
            creator = cand.creator
            if creator is not None and (creator, cand) not in seen_edges:
                add_cand(creator)
                seen_edges.add((creator, cand))
                nodes.add(HashableObject(creator))
                nodes.add(HashableObject(cand))
        elif isinstance(cand, function.Function):
            for input_ in cand.inputs:
                if input_ is not cand and (input_, cand) not in seen_edges:
                    creator = input_.creator
                    add_cand(input_)
                    seen_edges.add((input_, cand))
                    nodes.add(HashableObject(input_))
                    nodes.add(HashableObject(cand))
    return ComputationalGraph(list(i.v for i in nodes), list(seen_edges),
                              variable_style, function_style)
<commit_msg>Enable to change label attribute of dotnode with argument<commit_after>import heapq

from chainer import function
from chainer import variable

class DotNode(object):
    """Node of the computational graph, with utilities for dot language.

    This class represents a node of computational graph,
    with some utilities for dot language.

    """
    def __init__(self, node, attribute=None):
        """Initializes DotNode.

        Args:
            node: :class: `Variable` object or :class: `Function` object.
            attribute (dict): Attributes for the node.

        """
        assert isinstance(node, (variable.Variable, function.Function))
        self.node = node
        self.id_ = id(node)
        self.attribute = {'label': node.label}
        if isinstance(node, variable.Variable):
            self.attribute = {'shape': 'oval'}
        else:
            self.attribute = {'shape': 'box'}
        if attribute is not None:
            self.attribute.update(attribute)

    @property
    def label(self):
        """The text that represents properties of the node.

        Returns:
            string: The text that represents the id and attributes of this
                node.
        """

        attributes = ["%s=\"%s\"" % (k, v) for (k, v)
                      in self.attribute.items()]
        return "%s [%s];" % (self.id_, ",".join(attributes))

class ComputationalGraph(object):

    """Class that represents computational graph.

    .. note::

      We assume that the computational graph is directed and acyclic.

    """
    def __init__(self, nodes, edges, variable_style=None, function_style=None):
        """Initializes computational graph.

        Args:
            nodes (list): List of nodes. Each node is either
                 :class:`Variable` object or :class:`Function` object.
            edges (list): List of edges. Each edge consists of pair of nodes.
            variable_style (dist): Dot node style for variable.
            function_style (dist): Dot node style for function.

        """
        self.nodes = nodes
        self.edges = edges
        self.variable_style = variable_style
        self.function_style = function_style

    def _to_dot(self):
        """Converts graph in dot format.

        `label` property of is used as short description of each node.
        Returns:
            str: The graph in dot format.

        """
        ret = "digraph graphname{"
        for node in self.nodes:
            assert isinstance(node, (variable.Variable, function.Function))
            if isinstance(node, variable.Variable):
                ret += DotNode(node, self.variable_style).label
            else:
                ret += DotNode(node, self.function_style).label
        for edge in self.edges:
            head, tail = edge
            if (isinstance(head, variable.Variable) and
                    isinstance(tail, function.Function)):
                head_attr = self.variable_style
                tail_attr = self.function_style
            elif (isinstance(head, function.Function) and
                  isinstance(tail, variable.Variable)):
                head_attr = self.function_style
                tail_attr = self.variable_style
            else:
                raise TypeError(
                    'head and tail should be the set of Variable and Function')
            head_node = DotNode(head, head_attr)
            tail_node = DotNode(tail, tail_attr)
            ret += "%s -> %s;" % (head_node.id_, tail_node.id_)
        ret += "}"
        return ret

    def dump(self, format='dot'):
        """Dumps graph as a text.

        Args
            format(str): The graph language name of the output.
            Currently, it must be 'dot'.

        Returns
            str: The graph in specified format.

        """
        if format == 'dot':
            return self._to_dot()
        else:
            NotImplementedError('Currently, only dot format is supported.')

def build_computational_graph(outputs, remove_split=True,
                              variable_style=None, function_style=None):
    """Builds a graph of functions and variables backward-reachable from outputs.

    Args:
        outputs(list): nodes from which the graph is constructed.
            Each element of outputs must be either :class:`Variable`
            object or :class:`Function` object.
        remove_split(bool): It must be ``True``. This argument is left for
            backward compatibility.
        variable_style(dict): Dot node style for variable.
            Possible keys are 'shape', 'color', 'fillcolor', 'style', and etc.
        function_style(dict): Dot node style for function.

    Returns:
        ComputationalGraph: A graph consisting of nodes and edges that
        are backward-reachable from at least one of ``outputs``.

        If ``unchain_backward`` was called in some variable in the
        computational graph before this function, backward step is
        stopped at this variable.

        For example, suppose that computational graph is as follows::

                |--> f ---> y
            x --+
                |--> g ---> z

        Let ``outputs = [y, z]``.
        Then the full graph is emitted.

        Next, let ``outputs = [y]``. Note that ``z`` and ``g``
        are not backward-reachable from ``y``.
        The resulting graph would be following::

            x ---> f ---> y

        See :class:`TestGraphBuilder` for details.

    """
    if not remove_split:
        raise ValueError('remove_split=False is not supported anymore')

    cands = []
    seen_edges = set()
    nodes = set()
    push_count = [0]

    # This class is for object that has not been implemented __eq__
    class HashableObject(object):

        def __init__(self, v):
            self.v = v

        def __hash__(self):
            return self.v.__hash__()

        def __eq__(self, r):
            return self.v is r.v

    def add_cand(cand):
        heapq.heappush(cands, (-cand.rank, push_count[0], cand))
        push_count[0] += 1

    for o in outputs:
        add_cand(o)
        nodes.add(HashableObject(o))

    while cands:
        _, _, cand = heapq.heappop(cands)
        if isinstance(cand, variable.Variable):
            creator = cand.creator
            if creator is not None and (creator, cand) not in seen_edges:
                add_cand(creator)
                seen_edges.add((creator, cand))
                nodes.add(HashableObject(creator))
                nodes.add(HashableObject(cand))
        elif isinstance(cand, function.Function):
            for input_ in cand.inputs:
                if input_ is not cand and (input_, cand) not in seen_edges:
                    creator = input_.creator
                    add_cand(input_)
                    seen_edges.add((input_, cand))
                    nodes.add(HashableObject(input_))
                    nodes.add(HashableObject(cand))
    return ComputationalGraph(list(i.v for i in nodes), list(seen_edges),
                              variable_style, function_style)
<|endoftext|>


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

from __future__ import print_function
import os.path
import time
import datetime

import tifffile
import numpy

class TiffSaver(object):

    _outputdir = '/tmp'
    _basename = 'scan'
    _databroker = None

    start_time = datetime.datetime(2015, 03, 19)
    start_time = time.mktime(start_time.timetuple())
    _mtime_window = 0.05
    _output_mtime = 0
    _timetiffs = None

    def info(self):
        "Print out TiffSaver configuration."
        print("TiffSaver configuration:")
        print("  outputdir =", self.outputdir)
        print("  basename =", self.basename)
        print("Last 5 tiff files:")
        tnms = self.timetiffs.values()
        for tn in tnms[-5:]:
            print("  " + tn)
        return

    def listtiffs(self, count=None):
        """List the last count tiff files in the output directory.
        """
        print("#", self.outputdir)
        tt = self.timetiffs
        tms = tt.keys()
        if count:  tms = tms[-count:]
        for ti in tms:
            tmiso = datetime.datetime.fromtimestamp(round(ti)).isoformat()
            print(tmiso, os.path.basename(tt[ti]), sep='    ')
        return

    def saveScans(self, scanspec, basename=None):
        """Save the specified scans to the outputdir

        scanspec -- can be an integer index an array of indices or a slice
                    object like numpy.s_[_5:].  This looks up entries with
                    corresponding scan_id in the data broker.
        basename -- optional basename to be used for exporting these scans.
                    If not specified, use self.basename.

        No return value.
        """
        headers = self.findHeaders(scanspec)
        stash_basename = self.basename
        if basename is not None:
            self.basename = basename
        try:
            for h in headers:
                self.writeHeader(h)
        finally:
            self.basename = stash_basename
        return

    def saveNewScans(self):
        "Export scan data that were not yet saved to outputdir."
        lasttime = max([self.start_time] + self.timetiffs.keys())
        db = self.databroker
        headers = db.find_events(start_time=lasttime)
        for header in headers:
            self.writeHeader(header, overwrite=False)
        return

    def writeHeader(self, header, overwrite=False):
        from uiophyd.brokerutils import blank_events
        events = blank_events(header)
        for n, e in enumerate(events):
            fname = self._getOutputFilename(header=header, event=e, index=n)
            self.writeEvent(fname, e, overwrite=overwrite)
        return

    def writeEvent(self, filename, event, overwrite):
        from uiophyd.brokerutils import fill_event
        savedfile = self.findSavedEvent(event)
        if savedfile:
            if overwrite:
                os.remove(savedfile)
            else:
                return
        dd = event.data
        nlight = [k for k in dd if k.endswith('image_lightfield')]
        Alight = dd[nlight[0]]
        if isinstance(Alight, basestring):
            fill_event(event)
            Alight = event.data[nlight[0]]
        if 3 == Alight.ndim:
            Alight = Alight.sum(axis=0)
        if 0 == Alight.size:
            return
        A = Alight
        tifffile.imsave(filename, A.astype(numpy.float32))
        stinfo = os.stat(filename)
        os.utime(filename, (stinfo.st_atime, event.time))
        return

    def findSavedEvent(self, e):
        """Return full path to a saved tiff file for the specified scan event.

        e    -- event document object.

        Return string or None if the event has not been saved
        """
        import bisect
        tt = self.timetiffs
        tms = tt.keys()
        idx = bisect.bisect(tms, e.time)
        for ti in tms[idx-1:idx+1]:
            if abs(e.time - ti) < self._mtime_window:
                return tt[ti]
        return None

    def findHeaders(self, scanspec):
        """Produce header objects corresponding to scan specification.

        scanspec -- can be an integer index an array of indices,
                    slice object like numpy.s_[_5:] or string uid.

        Return a list of header objects.
        """
        db = self.databroker
        rv = []
        if isinstance(scanspec, int):
            rv.append(db[scanspec])
        elif isinstance(scanspec, slice):
            rv += db[scanspec]
        elif isinstance(scanspec, basestring):
            rv += db.find_headers(uid=scanspec)
        else:
            rv += map(self.findHeaders, scanspec)
        return rv

    # Properties -------------------------------------------------------------

    @property
    def outputdir(self):
        """Path to the tiff output directory.  When set, relative paths
        are expanded to a full physical path.
        """
        return self._outputdir

    @outputdir.setter
    def outputdir(self, value):
        if not os.path.isdir(value):
            emsg = "{!r} is not a directory.".format(value)
            raise ValueError(emsg)
        self._outputdir = os.path.abspath(value)
        return

    @property
    def basename(self):
        """Basename for the tiff files to be written.
        """
        return self._basename

    @basename.setter
    def basename(self, value):
        if '/' in value or '\\' in value:
            emsg = "Basename must not contain '/' or '\\'."
            raise ValueError(emsg)
        self._basename = value.strip().rstrip('-')
        return

    @property
    def databroker(self):
        "Return the active databroker instance."
        if self._databroker is None:
            from dataportal.broker import DataBroker
            self._databroker = DataBroker
        return self._databroker

    @property
    def timetiffs(self):
        "Ordered dictionary of mtimes and tiff files in outputdir."
        from collections import OrderedDict
        if os.path.getmtime(self.outputdir) == self._output_mtime:
            return self._timetiffs
        allfiles = [os.path.join(self.outputdir, f)
                for f in os.listdir(self.outputdir)]
        alltiffs = [f for f in allfiles
                if f.endswith('.tiff') and os.path.isfile(f)]
        tt = sorted((os.path.getmtime(f), f) for f in alltiffs)
        self._timetiffs = OrderedDict(tt)
        self._output_mtime = os.path.getmtime(self.outputdir)
        return self.timetiffs

    # Helpers ----------------------------------------------------------------

    def _getOutputFilename(self, header, event, index):
        scan_id = header.scan_id
        dd = event.data
        suffix = "{:05d}-{:03d}".format(scan_id, index)
        if 'cs700' in dd:
            tk = dd['cs700'][0]
            suffix = "{}-T{:03.1f}".format(suffix, tk)
        fmt = '{self.outputdir}/{self.basename}-{suffix}.tiff'
        fname = fmt.format(self=self, suffix=suffix)
        return fname

# class TiffSaver
<commit_msg>Add configurable suffixes to TiffSaver.<commit_after>#!/usr/bin/env python

from __future__ import print_function
import os.path
import time
import datetime

import tifffile
import numpy

class TiffSaver(object):

    _outputdir = '/tmp'
    _basename = 'scan'
    _databroker = None

    start_time = datetime.datetime(2015, 03, 19)
    start_time = time.mktime(start_time.timetuple())
    _mtime_window = 0.05
    _output_mtime = 0
    _timetiffs = None
    default_suffixes = ("{scan_id:05d}", "{index:03d}",
            "T{event.data[cs700]:03.1f}")

    def __init__(self):
        self.suffixes = list(self.default_suffixes)
        return

    def info(self):
        "Print out TiffSaver configuration."
        print("TiffSaver configuration:")
        print("  outputdir =", self.outputdir)
        print("  basename =", self.basename)
        print("  suffixes =", self.suffixes)
        print("Last 5 tiff files:")
        tnms = self.timetiffs.values()
        for tn in tnms[-5:]:
            print("  " + tn)
        return

    def listtiffs(self, count=None):
        """List the last count tiff files in the output directory.
        """
        print("#", self.outputdir)
        tt = self.timetiffs
        tms = tt.keys()
        if count:  tms = tms[-count:]
        for ti in tms:
            tmiso = datetime.datetime.fromtimestamp(round(ti)).isoformat()
            print(tmiso, os.path.basename(tt[ti]), sep='    ')
        return

    def saveScans(self, scanspec, basename=None):
        """Save the specified scans to the outputdir

        scanspec -- can be an integer index an array of indices or a slice
                    object like numpy.s_[_5:].  This looks up entries with
                    corresponding scan_id in the data broker.
        basename -- optional basename to be used for exporting these scans.
                    If not specified, use self.basename.

        No return value.
        """
        headers = self.findHeaders(scanspec)
        stash_basename = self.basename
        if basename is not None:
            self.basename = basename
        try:
            for h in headers:
                self.writeHeader(h)
        finally:
            self.basename = stash_basename
        return

    def saveNewScans(self):
        "Export scan data that were not yet saved to outputdir."
        lasttime = max([self.start_time] + self.timetiffs.keys())
        db = self.databroker
        headers = db.find_events(start_time=lasttime)
        for header in headers:
            self.writeHeader(header, overwrite=False)
        return

    def writeHeader(self, header, overwrite=False):
        from uiophyd.brokerutils import blank_events
        events = blank_events(header)
        for n, e in enumerate(events):
            fname = self._getOutputFilename(header=header, event=e, index=n)
            self.writeEvent(fname, e, overwrite=overwrite)
        return

    def writeEvent(self, filename, event, overwrite):
        from uiophyd.brokerutils import fill_event
        savedfile = self.findSavedEvent(event)
        if savedfile:
            if overwrite:
                os.remove(savedfile)
            else:
                return
        dd = event.data
        nlight = [k for k in dd if k.endswith('image_lightfield')]
        Alight = dd[nlight[0]]
        if isinstance(Alight, basestring):
            fill_event(event)
            Alight = event.data[nlight[0]]
        if 3 == Alight.ndim:
            Alight = Alight.sum(axis=0)
        if 0 == Alight.size:
            return
        A = Alight
        tifffile.imsave(filename, A.astype(numpy.float32))
        stinfo = os.stat(filename)
        os.utime(filename, (stinfo.st_atime, event.time))
        return

    def findSavedEvent(self, e):
        """Return full path to a saved tiff file for the specified scan event.

        e    -- event document object.

        Return string or None if the event has not been saved
        """
        import bisect
        tt = self.timetiffs
        tms = tt.keys()
        idx = bisect.bisect(tms, e.time)
        for ti in tms[idx-1:idx+1]:
            if abs(e.time - ti) < self._mtime_window:
                return tt[ti]
        return None

    def findHeaders(self, scanspec):
        """Produce header objects corresponding to scan specification.

        scanspec -- can be an integer index an array of indices,
                    slice object like numpy.s_[_5:] or string uid.

        Return a list of header objects.
        """
        db = self.databroker
        rv = []
        if isinstance(scanspec, int):
            rv.append(db[scanspec])
        elif isinstance(scanspec, slice):
            rv += db[scanspec]
        elif isinstance(scanspec, basestring):
            rv += db.find_headers(uid=scanspec)
        else:
            rv += map(self.findHeaders, scanspec)
        return rv

    # Properties -------------------------------------------------------------

    @property
    def outputdir(self):
        """Path to the tiff output directory.  When set, relative paths
        are expanded to a full physical path.
        """
        return self._outputdir

    @outputdir.setter
    def outputdir(self, value):
        if not os.path.isdir(value):
            emsg = "{!r} is not a directory.".format(value)
            raise ValueError(emsg)
        self._outputdir = os.path.abspath(value)
        return

    @property
    def basename(self):
        """Basename for the tiff files to be written.
        """
        return self._basename

    @basename.setter
    def basename(self, value):
        if '/' in value or '\\' in value:
            emsg = "Basename must not contain '/' or '\\'."
            raise ValueError(emsg)
        self._basename = value.strip().rstrip('-')
        return

    @property
    def databroker(self):
        "Return the active databroker instance."
        if self._databroker is None:
            from dataportal.broker import DataBroker
            self._databroker = DataBroker
        return self._databroker

    @property
    def timetiffs(self):
        "Ordered dictionary of mtimes and tiff files in outputdir."
        from collections import OrderedDict
        if os.path.getmtime(self.outputdir) == self._output_mtime:
            return self._timetiffs
        allfiles = [os.path.join(self.outputdir, f)
                for f in os.listdir(self.outputdir)]
        alltiffs = [f for f in allfiles
                if f.endswith('.tiff') and os.path.isfile(f)]
        tt = sorted((os.path.getmtime(f), f) for f in alltiffs)
        self._timetiffs = OrderedDict(tt)
        self._output_mtime = os.path.getmtime(self.outputdir)
        return self.timetiffs

    # Helpers ----------------------------------------------------------------

    def _getOutputFilename(self, header, event, index):
        sflst = []
        for s in self.suffixes:
            try:
                s2 = s.format(header=header, event=event, index=index,
                        scan_id=header.scan_id)
            except (AttributeError, KeyError):
                continue
            sflst.append(s2)
        tailname = '-'.join([self.basename] + sflst) + '.tif'
        fname = self.outputdir + '/' + tailname
        return fname

# class TiffSaver
<|endoftext|>


<commit_before>'''
Scan USB host for vendor specific device support

Usage:
    umap2vsscan -P=PHY_INFO [-q] [-d=DB_FILE] [-s=VID:PID] [-t=TIMEOUT] [-z|-b=DELAY] [-r=RESUME_FILE] [-o=OS]  [-e] [-v ...]

Options:
    -P --phy PHY_INFO           physical layer info, see list below
    -v --verbose                verbosity level
    -q --quiet                  quiet mode. only print warning/error messages
    -d --db DB_FILE             vid, pid database file (see DB_FILE below)
    -s --vid_pid VID:PID        specific VID:PID combination scan
    -t --timeout TIMEOUT        seconds to wait for host to detect each device (defualt: 3)
    -r --resume RESUME_FILE     filename to store/load scan session data
    -z --single_step            wait for keypress between each test
    -b --between DELAY          delay in seconds to wait between tests
    -o --os OS                  specify the host OS (default: Linux)
    -e --exhaustive             go over each (vid, pid) combination - do not skip device if its driver is in the supported list

Physical layer:
    fd:<serial_port>        use facedancer connected to given serial port
    gadgetfs                use gadgetfs (requires mounting of gadgetfs beforehand)

DB_FILE:
    a python file with a db member which is a list of DBEntry() objects.
    a sample can be found at: umap2/data/vid_pid_db.py

OS:
    Linux, Windows, OSX, QNX

VID:PID
    can be of the form 1234:5678 or 1234-1236:1235-1555

Examples:
    scan using a db file with 5 seconds timeout and 2 seconds delay between tries
    $ umap2vsscan -P fd:/dev/ttyUSB0 -d vid_pid_db.py -t 5 -b 2
    scan using facedancer a specific vid:pid with 5 seconds timeout
    $ umap2vsscan -P fd:/dev/ttyUSB0 -s 2058:1005 -t 5
'''
import time
import traceback
import os
import signal
import sys
import cPickle
from umap2.apps.base import Umap2App
from umap2.dev.vendor_specific import USBVendorSpecificDevice

class OS(object):
    LINUX = 'Linux'
    WINDOWS = 'Windows'
    OSX = 'OSX'
    QNX = 'QNX'

class DBEntry(object):
    '''
    DBEnrty describes a vid, pid.
    '''

    def __init__(self, vid, pid, vendor_name='', product_name='', drivers={}, constraints=[], info={}):
        self.vid = vid
        self.pid = pid
        self.vendor_name = vendor_name
        self.product_name = product_name
        self.drivers = drivers
        self.constraints = constraints
        self.info = info
        self.os = None

    def __str__(self):
        s = 'vid:pid %04x:%04x' % (self.vid, self.pid)
        if self.vendor_name:
            s += ', vendor: %s' % self.vendor_name
        if self.product_name:
            s += ', product: %s' % self.product_name
        if self.drivers:
            if self.os and self.os in self.drivers:
                s += ', driver: %s' % self.drivers[self.os]
            else:
                s += ', drivers: %s' % self.drivers
        if self.constraints:
            s += ', constraints: %s' % self.constraints
        if self.info:
            s += ', info: %s' % self.info
        return s

    def vidpid(self):
        return '%04x:%04x' % (self.vid, self.pid)

class _ScanSession(object):

    def __init__(self):
        self.timeout = 5
        self.db = []
        self.supported = []
        self.unsupported = []
        self.supported_drivers = []
        # key: device that got no response
        # value: previous device (if any)
        self.no_response = {}
        self.current = 0

class Umap2VSScanApp(Umap2App):

    def __init__(self, options):
        super(Umap2VSScanApp, self).__init__(options)
        self.current_usb_function_supported = False
        self.scan_session = _ScanSession()
        self.start_time = 0
        self.stop_signal_received = False
        self.between_delay = 5
        signal.signal(signal.SIGINT, self.signal_handler)
        timeout = self.options['--timeout']
        if timeout:
            self.scan_session.timeout = int(timeout)
        self.single_step = False
        if self.options['--single_step']:
            self.single_step = True
        elif self.options['--between']:
            self.between_delay = int(self.options['--between'])
        self.os = self.options['--os']
        if not self.os:
            self.os = OS.LINUX
        else:
            if self.os not in [getattr(OS, x) for x in dir(OS) if not x.startswith('_')]:
                self.error('Unsupported OS: %s choose a supported OS or add the new one to the OS class' % self.os)

    def get_device_info(self, device):
        info = []
        if device.endpoints:
            for e in device.endpoints:
                info.append(device.endpoints[e].get_descriptor(valid=True))
        else:
            info = ''
        if info:
            return 'num_endpoints = %d' % len(info)
        else:
            return 'device not reached set configuration state'

    def load_db_from_file(self, db_file):
        self.logger.info('loading vid_pid db file: %s' % db_file)
        dirpath, filename = os.path.split(db_file)
        modulename = filename[:-3]
        if dirpath in sys.path:
            sys.path.remove(dirpath)
        sys.path.insert(0, dirpath)
        module = __import__(modulename, globals(), locals(), [], -1)
        self.scan_session.db = module.db
        self.logger.always('loaded %d entries' % len(self.scan_session.db))

    def build_db_from_vid_pid(self, vid_pid):
        vid, pid = vid_pid.split(':')
        if '-' in vid:
            vid_start = int(vid.split('-')[0], 16)
            vid_end = int(vid.split('-')[1], 16)
            self.logger.debug('vid start=%04x, vid_end=%04x' % (vid_start, vid_end))
            vid = xrange(vid_start, vid_end)
        else:
            vid = [int(vid, 16)]
        if '-' in pid:
            pid_start = int(pid.split('-')[0], 16)
            pid_end = int(pid.split('-')[1], 16)
            self.logger.debug('pid start=%04x, pid_end=%x' % (pid_start, pid_end))
            pid = xrange(pid_start, pid_end)
        else:
            pid = [int(pid, 16)]
        for v in vid:
            for p in pid:
                self.scan_session.db.append(DBEntry(v, p))

    def build_scan_session(self):
        self.resume_file = self.options['--resume']
        if self.resume_file and os.path.exists(self.resume_file):
                self.logger.always('Resume file found. Loading scan data')
                with open(self.resume_file, 'rb') as rf:
                    self.scan_session = cPickle.load(rf)
        else:
            db_file = self.options['--db']
            vid_pid = self.options['--vid_pid']
            self.logger.always('Resume file not found. Creating new one')
            if db_file and vid_pid:
                self.logger.warning('not expecting both db file and specific vid:pid. we will use vid:pid')
            if vid_pid:
                self.build_db_from_vid_pid(vid_pid)
            elif db_file:
                self.load_db_from_file(db_file)
            else:
                self.logger.error('Must select a scan option - db (-d) or specific vid:pid (-p)')
                return

    def sync_and_increment_session(self):
        self.scan_session.current += 1
        self.sync_session()

    def sync_session(self):
        if self.resume_file:
            with open(self.resume_file, 'wb') as rf:
                cPickle.dump(self.scan_session, rf, 2)

    def print_results(self):
        num_supported = len(self.scan_session.supported)
        num_unsupported = len(self.scan_session.unsupported)
        self.logger.always('----------------------------------------')
        self.logger.always('Found %s supported device(s) (out of %s):' % (num_supported, self.scan_session.current))
        for i, db_entry in enumerate(self.scan_session.supported):
            self.logger.always('%d. %s' % (i, db_entry))
        self.logger.always('----------------------------------------')
        self.logger.always('Found %s unsupported device(s) (out of %s):' % (num_unsupported, self.scan_session.current))
        for i, db_entry in enumerate(self.scan_session.unsupported):
            self.logger.always('%d. %s' % (i, db_entry))
        self.logger.always('----------------------------------------')
        self.logger.always('Devices with no response (previous):')
        for i in sorted(self.scan_session.no_response.keys()):
            if self.scan_session.no_response[i]:
                prev = self.scan_session.no_response[i]
                pvp = self.scan_session.db[prev].vidpid()
            else:
                pvp = None
            self.logger.always('%s (%s)' % (self.scan_session.db[i], pvp))

    def run(self):
        self.build_scan_session()
        self.logger.always('Scanning host for supported vendor specific devices')
        phy = self.load_phy(self.options['--phy'])
        self.prev_index = None
        while self.scan_session.current < (len(self.scan_session.db)):
            if self.stop_signal_received:
                break
            db_entry = self.scan_session.db[self.scan_session.current]
            db_entry.os = self.os
            vid = db_entry.vid
            pid = db_entry.pid
            if not self.options['--exhaustive']:
                driver = db_entry.drivers.get(self.os, None)
                if driver and driver in self.scan_session.supported_drivers:
                    self.logger.always('skipping entry: %s' % db_entry)
                    self.sync_and_increment_session()
                    continue
            self.logger.always('Testing support for %s' % db_entry)
            self.setup_packet_received = False
            self.current_usb_function_supported = False
            self.start_time = time.time()
            device = USBVendorSpecificDevice(self, phy, vid, pid)
            try:
                device.connect()
                device.run()
            except:
                self.logger.error(traceback.format_exc())
            device.disconnect()
            if not self.is_host_alive():
                break
            if self.current_usb_function_supported:
                db_entry.info = self.get_device_info(device)
                self.scan_session.supported.append(db_entry)
                driver = db_entry.drivers.get(self.os, None)
                if driver:
                    self.scan_session.supported_drivers.append(db_entry.drivers[self.os])
            else:
                db_entry.info = self.get_device_info(device)
                self.scan_session.unsupported.append(db_entry)
            self.prev_index = self.scan_session.current
            self.sync_and_increment_session()
            if self.single_step:
                raw_input('press any key to continue')
            else:
                time.sleep(self.between_delay)
        self.print_results()

    def is_host_alive(self):
        if not self.setup_packet_received:
            self.logger.error('Host appears to have died or is simply ignoring us :(')
            current_index = self.scan_session.current
            if current_index not in self.scan_session.no_response:
                self.scan_session.no_response[current_index] = self.prev_index
            self.sync_session()
        return self.setup_packet_received

    def usb_function_supported(self):
        '''
        Callback from a USB device, notifying that the current USB device
        is supported by the host.
        '''
        self.current_usb_function_supported = True

    def signal_handler(self, signal, frame):
        self.stop_signal_received = True

    def should_stop_phy(self):
        stop_phy = False
        time_elapsed = int(time.time() - self.start_time)
        if self.current_usb_function_supported:
            stop_phy = True
        elif time_elapsed >= self.scan_session.timeout:
            self.logger.info('have been waiting long enough (over %d secs.), disconnect' % (time_elapsed))
            stop_phy = True
        return stop_phy

def main():
    app = Umap2VSScanApp(__doc__)
    app.run()

if __name__ == '__main__':
    main()
<commit_msg>[vsscan] removed logging of unsupported devices<commit_after>'''
Scan USB host for vendor specific device support

Usage:
    umap2vsscan -P=PHY_INFO [-q] [-d=DB_FILE] [-s=VID:PID] [-t=TIMEOUT] [-z|-b=DELAY] [-r=RESUME_FILE] [-o=OS]  [-e] [-v ...]

Options:
    -P --phy PHY_INFO           physical layer info, see list below
    -v --verbose                verbosity level
    -q --quiet                  quiet mode. only print warning/error messages
    -d --db DB_FILE             vid, pid database file (see DB_FILE below)
    -s --vid_pid VID:PID        specific VID:PID combination scan
    -t --timeout TIMEOUT        seconds to wait for host to detect each device (defualt: 3)
    -r --resume RESUME_FILE     filename to store/load scan session data
    -z --single_step            wait for keypress between each test
    -b --between DELAY          delay in seconds to wait between tests
    -o --os OS                  specify the host OS (default: Linux)
    -e --exhaustive             go over each (vid, pid) combination - do not skip device if its driver is in the supported list

Physical layer:
    fd:<serial_port>        use facedancer connected to given serial port
    gadgetfs                use gadgetfs (requires mounting of gadgetfs beforehand)

DB_FILE:
    a python file with a db member which is a list of DBEntry() objects.
    a sample can be found at: umap2/data/vid_pid_db.py

OS:
    Linux, Windows, OSX, QNX

VID:PID
    can be of the form 1234:5678 or 1234-1236:1235-1555

Examples:
    scan using a db file with 5 seconds timeout and 2 seconds delay between tries
    $ umap2vsscan -P fd:/dev/ttyUSB0 -d vid_pid_db.py -t 5 -b 2
    scan using facedancer a specific vid:pid with 5 seconds timeout
    $ umap2vsscan -P fd:/dev/ttyUSB0 -s 2058:1005 -t 5
'''
import time
import traceback
import os
import signal
import sys
import cPickle
from umap2.apps.base import Umap2App
from umap2.dev.vendor_specific import USBVendorSpecificDevice

class OS(object):
    LINUX = 'Linux'
    WINDOWS = 'Windows'
    OSX = 'OSX'
    QNX = 'QNX'

class DBEntry(object):
    '''
    DBEnrty describes a vid, pid.
    '''

    def __init__(self, vid, pid, vendor_name='', product_name='', drivers={}, constraints=[], info={}):
        self.vid = vid
        self.pid = pid
        self.vendor_name = vendor_name
        self.product_name = product_name
        self.drivers = drivers
        self.constraints = constraints
        self.info = info
        self.os = None

    def __str__(self):
        s = 'vid:pid %04x:%04x' % (self.vid, self.pid)
        if self.vendor_name:
            s += ', vendor: %s' % self.vendor_name
        if self.product_name:
            s += ', product: %s' % self.product_name
        if self.drivers:
            if self.os and self.os in self.drivers:
                s += ', driver: %s' % self.drivers[self.os]
            else:
                s += ', drivers: %s' % self.drivers
        if self.constraints:
            s += ', constraints: %s' % self.constraints
        if self.info:
            s += ', info: %s' % self.info
        return s

    def vidpid(self):
        return '%04x:%04x' % (self.vid, self.pid)

class _ScanSession(object):

    def __init__(self):
        self.timeout = 5
        self.db = []
        self.supported = []
        self.unsupported = []
        self.supported_drivers = []
        # key: device that got no response
        # value: previous device (if any)
        self.no_response = {}
        self.current = 0

class Umap2VSScanApp(Umap2App):

    def __init__(self, options):
        super(Umap2VSScanApp, self).__init__(options)
        self.current_usb_function_supported = False
        self.scan_session = _ScanSession()
        self.start_time = 0
        self.stop_signal_received = False
        self.between_delay = 5
        signal.signal(signal.SIGINT, self.signal_handler)
        timeout = self.options['--timeout']
        if timeout:
            self.scan_session.timeout = int(timeout)
        self.single_step = False
        if self.options['--single_step']:
            self.single_step = True
        elif self.options['--between']:
            self.between_delay = int(self.options['--between'])
        self.os = self.options['--os']
        if not self.os:
            self.os = OS.LINUX
        else:
            if self.os not in [getattr(OS, x) for x in dir(OS) if not x.startswith('_')]:
                self.error('Unsupported OS: %s choose a supported OS or add the new one to the OS class' % self.os)

    def get_device_info(self, device):
        info = []
        if device.endpoints:
            for e in device.endpoints:
                info.append(device.endpoints[e].get_descriptor(valid=True))
        else:
            info = ''
        if info:
            return 'num_endpoints = %d' % len(info)
        else:
            return 'device not reached set configuration state'

    def load_db_from_file(self, db_file):
        self.logger.info('loading vid_pid db file: %s' % db_file)
        dirpath, filename = os.path.split(db_file)
        modulename = filename[:-3]
        if dirpath in sys.path:
            sys.path.remove(dirpath)
        sys.path.insert(0, dirpath)
        module = __import__(modulename, globals(), locals(), [], -1)
        self.scan_session.db = module.db
        self.logger.always('loaded %d entries' % len(self.scan_session.db))

    def build_db_from_vid_pid(self, vid_pid):
        vid, pid = vid_pid.split(':')
        if '-' in vid:
            vid_start = int(vid.split('-')[0], 16)
            vid_end = int(vid.split('-')[1], 16)
            self.logger.debug('vid start=%04x, vid_end=%04x' % (vid_start, vid_end))
            vid = xrange(vid_start, vid_end)
        else:
            vid = [int(vid, 16)]
        if '-' in pid:
            pid_start = int(pid.split('-')[0], 16)
            pid_end = int(pid.split('-')[1], 16)
            self.logger.debug('pid start=%04x, pid_end=%x' % (pid_start, pid_end))
            pid = xrange(pid_start, pid_end)
        else:
            pid = [int(pid, 16)]
        for v in vid:
            for p in pid:
                self.scan_session.db.append(DBEntry(v, p))

    def build_scan_session(self):
        self.resume_file = self.options['--resume']
        if self.resume_file and os.path.exists(self.resume_file):
                self.logger.always('Resume file found. Loading scan data')
                with open(self.resume_file, 'rb') as rf:
                    self.scan_session = cPickle.load(rf)
        else:
            db_file = self.options['--db']
            vid_pid = self.options['--vid_pid']
            self.logger.always('Resume file not found. Creating new one')
            if db_file and vid_pid:
                self.logger.warning('not expecting both db file and specific vid:pid. we will use vid:pid')
            if vid_pid:
                self.build_db_from_vid_pid(vid_pid)
            elif db_file:
                self.load_db_from_file(db_file)
            else:
                self.logger.error('Must select a scan option - db (-d) or specific vid:pid (-p)')
                return

    def sync_and_increment_session(self):
        self.scan_session.current += 1
        self.sync_session()

    def sync_session(self):
        if self.resume_file:
            with open(self.resume_file, 'wb') as rf:
                cPickle.dump(self.scan_session, rf, 2)

    def print_results(self):
        num_supported = len(self.scan_session.supported)
        # num_unsupported = len(self.scan_session.unsupported)
        self.logger.always('----------------------------------------')
        self.logger.always('Found %s supported device(s) (out of %s):' % (num_supported, self.scan_session.current))
        for i, db_entry in enumerate(self.scan_session.supported):
            self.logger.always('%d. %s' % (i, db_entry))
        self.logger.always('----------------------------------------')
        self.logger.always('Devices with no response (previous):')
        for i in sorted(self.scan_session.no_response.keys()):
            if self.scan_session.no_response[i]:
                prev = self.scan_session.no_response[i]
                pvp = self.scan_session.db[prev].vidpid()
            else:
                pvp = None
            self.logger.always('%s (%s)' % (self.scan_session.db[i], pvp))

    def run(self):
        self.build_scan_session()
        self.logger.always('Scanning host for supported vendor specific devices')
        phy = self.load_phy(self.options['--phy'])
        self.prev_index = None
        while self.scan_session.current < (len(self.scan_session.db)):
            if self.stop_signal_received:
                break
            db_entry = self.scan_session.db[self.scan_session.current]
            db_entry.os = self.os
            vid = db_entry.vid
            pid = db_entry.pid
            if not self.options['--exhaustive']:
                driver = db_entry.drivers.get(self.os, None)
                if driver and driver in self.scan_session.supported_drivers:
                    self.logger.always('skipping entry: %s' % db_entry)
                    self.sync_and_increment_session()
                    continue
            self.logger.always('Testing support for %s' % db_entry)
            self.setup_packet_received = False
            self.current_usb_function_supported = False
            self.start_time = time.time()
            device = USBVendorSpecificDevice(self, phy, vid, pid)
            try:
                device.connect()
                device.run()
            except:
                self.logger.error(traceback.format_exc())
            device.disconnect()
            if not self.is_host_alive():
                break
            if self.current_usb_function_supported:
                db_entry.info = self.get_device_info(device)
                self.scan_session.supported.append(db_entry)
                driver = db_entry.drivers.get(self.os, None)
                if driver:
                    self.scan_session.supported_drivers.append(db_entry.drivers[self.os])
            # else:
            #     db_entry.info = self.get_device_info(device)
            #     self.scan_session.unsupported.append(db_entry)
            self.prev_index = self.scan_session.current
            self.sync_and_increment_session()
            if self.single_step:
                raw_input('press any key to continue')
            else:
                time.sleep(self.between_delay)
        self.print_results()

    def is_host_alive(self):
        if not self.setup_packet_received:
            self.logger.error('Host appears to have died or is simply ignoring us :(')
            current_index = self.scan_session.current
            if current_index not in self.scan_session.no_response:
                self.scan_session.no_response[current_index] = self.prev_index
            self.sync_session()
        return self.setup_packet_received

    def usb_function_supported(self):
        '''
        Callback from a USB device, notifying that the current USB device
        is supported by the host.
        '''
        self.current_usb_function_supported = True

    def signal_handler(self, signal, frame):
        self.stop_signal_received = True

    def should_stop_phy(self):
        stop_phy = False
        time_elapsed = int(time.time() - self.start_time)
        if self.current_usb_function_supported:
            stop_phy = True
        elif time_elapsed >= self.scan_session.timeout:
            self.logger.info('have been waiting long enough (over %d secs.), disconnect' % (time_elapsed))
            stop_phy = True
        return stop_phy

def main():
    app = Umap2VSScanApp(__doc__)
    app.run()

if __name__ == '__main__':
    main()

<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
"""
Tools to handle reads sequenced with unique molecular identifiers (UMIs).
"""
from __future__ import print_function

import os
import sys
import doctest
from re import findall
from pysam import Samfile
from toolshed import nopen
from itertools import islice, izip
from collections import Counter, defaultdict
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter

from ._version import __version__

os.environ['TERM'] = 'linux'

IUPAC = {"A": "A", "T": "T", "C": "C", "G": "G", "R": "GA", "Y": "TC",
         "M": "AC", "K": "GT", "S": "GC", "W": "AT", "H": "ACT",
         "B": "GTC", "V": "GCA", "D": "GAT", "N": "GATC"}

class Fastq(object):
    def __init__(self, args):
        self.name = args[0][1:]
        self.seq = args[1]
        self.qual = args[3]
        assert len(self.seq) == len(self.qual)

    def __repr__(self):
        return "Fastq({name})".format(name=self.name)

    def __str__(self):
        return "@{name}\n{seq}\n+\n{qual}".format(name=self.name,
                                                  seq=self.seq, qual=self.qual)

def umi_from_name(name):
    """
    extract the UMI sequence from the read name.

    >>> umi_from_name("cluster_1017333:UMI_GCCGCA")
    'GCCGCA'
    """
    return findall(r'UMI_([\w]*)', name)[0].strip()

def process_bam(args):
    """
    removes duplicate reads characterized by their UMI at any given start
    location.

    abam    input bam with potential duplicate UMIs
    bbam    output bam after removing duplicate UMIs
    """
    with Samfile(args.abam, 'rb') as in_bam, \
            Samfile(args.bbam, 'wb', template=in_bam) as out_bam:

        for chrom in in_bam.references:
            print("processing chromosome", chrom, file=sys.stderr)

            umi_idx = defaultdict(set)
            read_counts = Counter()

            for read in in_bam.fetch(chrom):
                if read.is_unmapped:
                    continue

                # get the iupac umi sequence
                umi = umi_from_name(read.qname)

                # get actual read start
                # read.pos accounts for 5' soft clipping
                if read.is_reverse:
                    # read.alen alignment length accounting for 3' soft clipping
                    # UMIs are then compared to reads with the same start
                    read_start = read.pos + read.alen
                else:
                    read_start = read.pos

                # add count for this start; counts all reads
                read_counts[read_start] += 1

                # check if UMI seen
                if umi in umi_idx[read_start]:
                    continue

                # keep track of unique UMIs - set eliminates duplicates
                umi_idx[read_start].add(umi)

                out_bam.write(read)

            # process before and after counts over chrom
            for start, before_count in sorted(read_counts.items()):
                print(chrom, start, start + 1, before_count, len(umi_idx[start]), sep="\t")

def readfq(fq):
    with nopen(fq) as fh:
        fqclean = (x.strip("\r\n") for x in fh if x.strip())
        while True:
            rd = [x for x in islice(fqclean, 4)]
            if not rd:
                raise StopIteration
            assert all(rd) and len(rd) == 4
            yield Fastq(rd)

def valid_umi(iupac, umi):
    """
    parse UMI sequence to validate against IUPAC sequence.

    >>> valid_umi("NNNV", "ACGT")
    False
    >>> valid_umi("NNNV", "ACGG")
    True
    """
    for code, base in izip(iupac, umi):
        try:
            if base not in IUPAC[code]:
                return False
        except KeyError:
            return False
    return True

def clip_umi(record, iupac_umi, n, end):
    """
    >>> fq = Fastq(["@cluster_455 2",\
                        "GGGGGAGCCACGAGGTGTGTTTTATTTTCATTATTC",\
                        "+",\
                        "C===>=B=@:<;4A;8=9?6EEC0?DDA72B@3EB4"])
    >>> clip_umi(fq, "NNNNNV", 6, "5")
    Fastq(cluster_455:UMI_GGGGGA 2)
    >>> fq = Fastq(["@cluster_455 2",\
                        "GGXXGAGCCACGAGGTGTGTTTTATTTTCATTATTC",\
                        "+",\
                        "C===>=B=@:<;4A;8=9?6EEC0?DDA72B@3EB4"])
    >>> clip_umi(fq, "NNNNNV", 6, "5")
    'GGXXGA'
    """
    if end == "5":
        umi = record.seq[:n]
        record.seq = record.seq[n:]
        record.qual = record.qual[n:]
    else:
        umi = record.seq[-n:]
        record.seq = record.seq[:-n]
        record.qual = record.qual[:-n]
    if not valid_umi(iupac_umi, umi):
        return umi
    try:
        name, pair = record.name.split(" ", 1)
        record.name = "{name}:UMI_{umi} {pair}".format(name=name,
                                                       umi=umi,
                                                       pair=pair)
    except ValueError:
        record.name = "{name}:UMI_{umi}".format(name=record.name, umi=umi)
    return record

def process_fastq(args):
    """
    for every valid umi, trim while incorporating into read name.

    args:
    fastq       reads to process
    umi         IUPAC sequence
    end         5 or 3 as a string
    top         int number of invalid sequences to output
    verbose
    """
    umi_stats = Counter()
    iupac = args.umi.upper()
    u_leng = len(iupac)
    end = args.end
    for r in readfq(args.fastq):
        r = clip_umi(r, iupac, u_leng, end)
        if type(r) is Fastq:
            print(r)
        else:
            umi_stats.update([r])
    if args.verbose:
        print("Invalid UMI Total:", sum(umi_stats.values()), file=sys.stderr)
        print("Unique UMIs Removed:", len(list(umi_stats)), file=sys.stderr)
        print("Top %d Invalid UMIs:" % args.top, file=sys.stderr)
        for umi, val in umi_stats.most_common(args.top):
            print(umi, val, sep="\t", file=sys.stderr)

def main():
    p = ArgumentParser(description=__doc__, version=__version__)
    subp = p.add_subparsers(help='commands')

    # fastq processing
    fastq = subp.add_parser('trim', description=("Trims the UMI sequence from the read, incorporating the unique "
                                                 "sequence in the read name facilitating filtering of the alignments."),
                            formatter_class=ArgumentDefaultsHelpFormatter,
                            help="trim UMI and incorporate sequence into read name")
    fastq.add_argument('fastq', metavar='FASTQ',
                       help='reads with untrimmed UMI')
    fastq.add_argument('umi', metavar='UMI',
                       help='IUPAC UMI sequence, e.g. NNNNNV')
    fastq.add_argument('--end', choices=['5', '3'], default="5",
                       help="UMI location on the read")
    fastq.add_argument('--verbose', action='store_true',
                       help="print UMI stats to stderr")
    fastq.add_argument('--top', type=int, default=10,
                       help="when verbose, print this many of the top filtered UMI sequences")
    fastq.set_defaults(func=process_fastq)

    # bam processing
    bam = subp.add_parser('rmdup', description=("Removes duplicate reads, that were previously characterized by "
                                                "their UMI, at any given start location. Coverage differences before "
                                                "and after are written to STDOUT as BED3+."),
                          formatter_class=ArgumentDefaultsHelpFormatter,
                          help="remove duplicate UMI entries from all start positions")
    bam.add_argument('abam', metavar='INPUT_BAM',
                     help='bam with UMI in read name')
    bam.add_argument('bbam', metavar='OUTPUT_BAM',
                     help='non-duplicate UMIs at any given start position')
    bam.set_defaults(func=process_bam)

    args = p.parse_args()
    args.func(args)

if __name__ == "__main__":
    if doctest.testmod(optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE).failed == 0:
        main()
<commit_msg>adds distance calculations<commit_after>#!/usr/bin/env python
# encoding: utf-8
"""
Tools to handle reads sequenced with unique molecular identifiers (UMIs).
"""
from __future__ import print_function

import doctest
import editdistance
import os
import sys
from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from collections import Counter, defaultdict
from itertools import islice, izip
from pysam import Samfile
from re import findall
from toolshed import nopen

from ._version import __version__

os.environ['TERM'] = 'linux'

IUPAC = {
    "A": "A",
    "T": "T",
    "C": "C",
    "G": "G",
    "R": "GA",
    "Y": "TC",
    "M": "AC",
    "K": "GT",
    "S": "GC",
    "W": "AT",
    "H": "ACT",
    "B": "GTC",
    "V": "GCA",
    "D": "GAT",
    "N": "GATC"
}

class Fastq(object):
    def __init__(self, args):
        self.name = args[0][1:]
        self.seq = args[1]
        self.qual = args[3]
        assert len(self.seq) == len(self.qual)

    def __repr__(self):
        return "Fastq(%s)" % self.name

    def __str__(self):
        return "@%s\n%s\n+\n%s" % (self.name, self.seq, self.qual)

def umi_from_name(name):
    """
    extract the UMI sequence from the read name.

    >>> umi_from_name("cluster_1017333:UMI_GCCGCA")
    'GCCGCA'
    """
    return findall(r'UMI_([\w]*)', name)[0].strip()

def passing_distances(subject, target_set, n):
    """
    >>> s = "ACTGA"
    >>> ts_1 = {"ACTGG"}
    >>> ts_2 = {"ACTCC", "ACTGG"}
    >>> ts_3 = {"ACTCC", "ACTTT"}
    >>> n = 1
    >>> passing_distances(s, ts_1, n)
    True
    >>> passing_distances(s, ts_2, n)
    True
    >>> passing_distances(s, ts_3, n)
    False
    """
    for target in target_set:
        if editdistance.distance(target, subject) <= n:
            return True
    return False

def process_bam(args):
    """
    removes duplicate reads characterized by their UMI at any given start
    location.

    abam        input bam with potential duplicate UMIs
    bbam        output bam after removing duplicate UMIs
    mismatches  allowable edit distance between UMIs
    """
    with Samfile(args.abam, 'rb') as in_bam, Samfile(args.bbam, 'wb', template=in_bam) as out_bam:

        for chrom in in_bam.references:
            print("processing chromosome", chrom, file=sys.stderr)

            umi_idx = defaultdict(set)
            read_counts = Counter()

            for read in in_bam.fetch(chrom):
                if read.is_unmapped:
                    continue

                # get the iupac umi sequence
                umi = umi_from_name(read.qname)

                # get actual read start
                # read.pos accounts for 5' soft clipping
                if read.is_reverse:
                    # read.alen alignment length accounting for 3' soft clipping
                    # UMIs are then compared to reads with the same start
                    read_start = read.pos + read.alen
                else:
                    read_start = read.pos

                # add count for this start; counts all reads
                read_counts[read_start] += 1

                # check if UMI seen
                if umi in umi_idx[read_start]:
                    continue
                elif args.mismatches > 0 and passing_distances(umi, umi_idx[read_start], args.mismatches):
                    # add UMI as unique hit
                    umi_idx[read_start].add(umi)
                    continue

                # keep track of unique UMIs - set eliminates duplicates
                umi_idx[read_start].add(umi)

                out_bam.write(read)

            # process before and after counts over chrom
            for start, before_count in sorted(read_counts.items()):
                print(chrom, start, start + 1, before_count, len(umi_idx[start]), sep="\t")

def readfq(fq):
    with nopen(fq) as fh:
        fqclean = (x.strip("\r\n") for x in fh if x.strip())
        while True:
            rd = [x for x in islice(fqclean, 4)]
            if not rd:
                raise StopIteration
            assert all(rd) and len(rd) == 4
            yield Fastq(rd)

def valid_umi(iupac, umi):
    """
    parse UMI sequence to validate against IUPAC sequence.

    >>> valid_umi("NNNV", "ACGT")
    False
    >>> valid_umi("NNNV", "ACGG")
    True
    """
    for code, base in izip(iupac, umi):
        try:
            if base not in IUPAC[code]:
                return False
        except KeyError:
            return False
    return True

def clip_umi(record, iupac_umi, n, end):
    """
    >>> fq = Fastq(["@cluster_455 2",\
                        "GGGGGAGCCACGAGGTGTGTTTTATTTTCATTATTC",\
                        "+",\
                        "C===>=B=@:<;4A;8=9?6EEC0?DDA72B@3EB4"])
    >>> clip_umi(fq, "NNNNNV", 6, "5")
    Fastq(cluster_455:UMI_GGGGGA 2)
    >>> fq = Fastq(["@cluster_455 2",\
                        "GGXXGAGCCACGAGGTGTGTTTTATTTTCATTATTC",\
                        "+",\
                        "C===>=B=@:<;4A;8=9?6EEC0?DDA72B@3EB4"])
    >>> clip_umi(fq, "NNNNNV", 6, "5")
    'GGXXGA'
    """
    if end == "5":
        umi = record.seq[:n]
        record.seq = record.seq[n:]
        record.qual = record.qual[n:]
    else:
        umi = record.seq[-n:]
        record.seq = record.seq[:-n]
        record.qual = record.qual[:-n]
    if not valid_umi(iupac_umi, umi):
        return umi
    try:
        name, pair = record.name.split(" ", 1)
        record.name = "{name}:UMI_{umi} {pair}".format(name=name,
                                                       umi=umi,
                                                       pair=pair)
    except ValueError:
        record.name = "{name}:UMI_{umi}".format(name=record.name, umi=umi)
    return record

def process_fastq(args):
    """
    for every valid umi, trim while incorporating into read name.

    args:
    fastq       reads to process
    umi         IUPAC sequence
    end         5 or 3 as a string
    top         int number of invalid sequences to output
    verbose
    """
    umi_stats = Counter()
    iupac = args.umi.upper()
    u_leng = len(iupac)
    end = args.end
    for r in readfq(args.fastq):
        r = clip_umi(r, iupac, u_leng, end)
        if type(r) is Fastq:
            print(r)
        else:
            umi_stats.update([r])
    if args.verbose:
        print("Invalid UMI Total:", sum(umi_stats.values()), file=sys.stderr)
        print("Unique UMIs Removed:", len(list(umi_stats)), file=sys.stderr)
        print("Top %d Invalid UMIs:" % args.top, file=sys.stderr)
        for umi, val in umi_stats.most_common(args.top):
            print(umi, val, sep="\t", file=sys.stderr)

def main():
    p = ArgumentParser(description=__doc__, version=__version__)
    subp = p.add_subparsers(help='commands')

    # fastq processing
    fastq = subp.add_parser('trim', description=("Trims the UMI sequence from the read, incorporating the unique "
                                                 "sequence in the read name facilitating filtering of the alignments."),
                            formatter_class=ArgumentDefaultsHelpFormatter,
                            help="trim UMI and incorporate sequence into read name")
    fastq.add_argument('fastq', metavar='FASTQ',
                       help='reads with untrimmed UMI')
    fastq.add_argument('umi', metavar='UMI',
                       help='IUPAC UMI sequence, e.g. NNNNNV')
    fastq.add_argument('--end', choices=['5', '3'], default="5",
                       help="UMI location on the read")
    fastq.add_argument('--verbose', action='store_true',
                       help="print UMI stats to stderr")
    fastq.add_argument('--top', type=int, default=10,
                       help="when verbose, print this many of the top filtered UMI sequences")
    fastq.set_defaults(func=process_fastq)

    # bam processing
    bam = subp.add_parser('rmdup', description=("Removes duplicate reads, that were previously characterized by "
                                                "their UMI, at any given start location. Coverage differences before "
                                                "and after are written to STDOUT as BED3+."),
                          formatter_class=ArgumentDefaultsHelpFormatter,
                          help="remove duplicate UMI entries from all start positions")
    bam.add_argument('abam', metavar='INPUT_BAM',
                     help='bam with UMI in read name')
    bam.add_argument('bbam', metavar='OUTPUT_BAM',
                     help='non-duplicate UMIs at any given start position')
    bam.add_argument('-m', '--mismatches', default=0, type=int,
                     help="allowable mismatches when comparing UMIs at any given start location")
    bam.set_defaults(func=process_bam)

    args = p.parse_args()
    args.func(args)

if __name__ == "__main__":
    if doctest.testmod(optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE).failed == 0:
        main()
<|endoftext|>


<commit_before>import getdns
import inspect
import StringIO
import sys, platform
import unittest

un = platform.uname()
d = 'lib.' + un[0].lower() + '-' + un[4] + '-' + '.'.join(platform.python_version().split('.')[:2])
sys.path.append(d)

class TestGetdnsMethods(unittest.TestCase):
    def test_context(self):
        c = getdns.Context()
        self.assertIsNotNone(c)
        del(c)

    def test_bogus_attribute(self):
        c = getdns.Context()
        with self.assertRaises(AttributeError, c.asdf) as e:

        except
        del(c)

    def test_append_name(self):
        c = getdns.Context()
        c.append_name = getdns.APPEND_NAME_NEVER
        self.assertEqual(c.append_name, getdns.APPEND_NAME_NEVER)
        del(c)

    def test_dns_root_servers(self):
        c = getdns.Context()
        addrs = [ { 'address_type': 'IPv4', 'address_data': '127.0.0.254' } ]
        c.dns_root_servers = addrs
        self.assertEqual(c.dns_root_servers, addrs)
        del(c)

    def test_dns_transport_list(self):
        c = getdns.Context()
        transports = [ getdns.TRANSPORT_TLS, getdns.TRANSPORT_UDP, getdns.TRANSPORT_TCP ]
        c.dns_transport_list = transports
        self.assertEqual(c.dns_transport_list, transports)
        del(c)

    def test_sync_address(self):
        c = getdns.Context()
        c.resolution_type=getdns.RESOLUTION_STUB
        r = c.address('www.getdnsapi.net')
        self.assertEqual(r.status, getdns.RESPSTATUS_GOOD)
        self.assertTrue('185.49.141.37' in [ x['address_data'] for x in r.just_address_answers ])
        del(c)
        del(r)

    def test_sync_service(self):
        c = getdns.Context()
        c.resolution_type=getdns.RESOLUTION_STUB
        r = c.service('_xmpp-server._tcp.jabber.org')
        self.assertEqual(r.status, getdns.RESPSTATUS_GOOD)
        del(c)
        del(r)

    def test_sync_hostname(self):
        c = getdns.Context()
        c.resolution_type=getdns.RESOLUTION_STUB
        r = c.hostname( { 'address_type': 'IPv4', 'address_data': '185.49.141.37' } )
        self.assertEqual(r.status, getdns.RESPSTATUS_GOOD)
        del(c)
        del(r)

    def test_sync_general(self):
        c = getdns.Context()
        c.resolution_type=getdns.RESOLUTION_STUB
        r = c.general('nlnetlabs.nl', request_type=getdns.RRTYPE_NS)
        self.assertEqual(r.status, getdns.RESPSTATUS_GOOD)
        del(c)
        del(r)

    def test_file_to_list(self):
        ns1 = {'class': 1,
               'name': 'example.com.',
               'rdata': {'nsdname': 'ns1.example.com.',
                         'rdata_raw':'ns1.example.com.'},
               'ttl': 3600,
               'type': 2
               }
        ns2 = {'class': 1,
               'name': 'example.com.',
               'rdata': {'nsdname': 'ns2.example.com.',
                         'rdata_raw': 'ns2.example.com.'},
               'ttl': 3600,
               'type': 2
               }
        f = open('example.com.zone')
        r = getdns.file_to_list(f, 'example.com', 3600 )
        self.assertIsInstance(r, list)
        self.assertEqual(r[1], ns1)
        self.assertEqual(r[2], ns2)
        del(f)
        del(r)

if __name__ == "__main__":
    unittest.main(verbosity=2)
<commit_msg>flake8 cleanup<commit_after>import getdns
import platform
import sys
import unittest

un = platform.uname()
d = "lib.{0}-{1}-{2}".format(
    un[0].lower(), un[4], '.'.join(platform.python_version().split('.')[:2])
)
sys.path.append(d)

class TestGetdnsMethods(unittest.TestCase):
    def test_context(self):
        c = getdns.Context()
        self.assertIsNotNone(c)
        del(c)

    def test_bogus_attribute(self):
        c = getdns.Context()
        with self.assertRaises(AttributeError, c.asdf) as e:

        except
        del(c)

    def test_append_name(self):
        c = getdns.Context()
        c.append_name = getdns.APPEND_NAME_NEVER
        self.assertEqual(c.append_name, getdns.APPEND_NAME_NEVER)
        del(c)

    def test_dns_root_servers(self):
        c = getdns.Context()
        addrs = [{'address_type': 'IPv4', 'address_data': '127.0.0.254'}]
        c.dns_root_servers = addrs
        self.assertEqual(c.dns_root_servers, addrs)
        del(c)

    def test_dns_transport_list(self):
        c = getdns.Context()
        transports = [getdns.TRANSPORT_TLS,
                      getdns.TRANSPORT_UDP,
                      getdns.TRANSPORT_TCP]
        c.dns_transport_list = transports
        self.assertEqual(c.dns_transport_list, transports)
        del(c)

    def test_sync_address(self):
        c = getdns.Context()
        c.resolution_type = getdns.RESOLUTION_STUB
        r = c.address('www.getdnsapi.net')
        self.assertEqual(r.status, getdns.RESPSTATUS_GOOD)
        self.assertTrue('185.49.141.37' in [x['address_data'] for x in
                                            r.just_address_answers])
        del(c)
        del(r)

    def test_sync_service(self):
        c = getdns.Context()
        c.resolution_type = getdns.RESOLUTION_STUB
        r = c.service('_xmpp-server._tcp.jabber.org')
        self.assertEqual(r.status, getdns.RESPSTATUS_GOOD)
        del(c)
        del(r)

    def test_sync_hostname(self):
        c = getdns.Context()
        c.resolution_type = getdns.RESOLUTION_STUB
        r = c.hostname({'address_type': 'IPv4',
                        'address_data': '185.49.141.37'})
        self.assertEqual(r.status, getdns.RESPSTATUS_GOOD)
        del(c)
        del(r)

    def test_sync_general(self):
        c = getdns.Context()
        c.resolution_type = getdns.RESOLUTION_STUB
        r = c.general('nlnetlabs.nl', request_type=getdns.RRTYPE_NS)
        self.assertEqual(r.status, getdns.RESPSTATUS_GOOD)
        del(c)
        del(r)

    def test_file_to_list(self):
        ns1 = {'class': 1,
               'name': 'example.com.',
               'rdata': {'nsdname': 'ns1.example.com.',
                         'rdata_raw': 'ns1.example.com.'},
               'ttl': 3600,
               'type': 2
               }
        ns2 = {'class': 1,
               'name': 'example.com.',
               'rdata': {'nsdname': 'ns2.example.com.',
                         'rdata_raw': 'ns2.example.com.'},
               'ttl': 3600,
               'type': 2
               }
        f = open('example.com.zone')
        r = getdns.file_to_list(f, 'example.com', 3600)
        self.assertIsInstance(r, list)
        self.assertEqual(r[1], ns1)
        self.assertEqual(r[2], ns2)
        del(f)
        del(r)

if __name__ == "__main__":
    unittest.main(verbosity=2)
<|endoftext|>