<commit_before>import pytest
import warnings
import logging
from .errors import AnsibleHostUnreachable

try:
    from logging import NullHandler
except ImportError:
    from logging import Handler

    class NullHandler(Handler):

        def emit(self, record):
            pass

log = logging.getLogger(__name__)
log.addHandler(NullHandler())

@pytest.fixture(scope='class')
def ansible_module_cls(request):
    '''
    Return AnsibleV1Module instance with class scope.
    '''
    warnings.warn("Use of ansible_facts_cls is deprecated and will be removed in a future release", DeprecationWarning)
    ansible_helper = request.config.pluginmanager.getplugin("ansible")
    return ansible_helper.initialize(request)

@pytest.fixture(scope='function')
def ansible_module(request):
    '''
    Return AnsibleV1Module instance with function scope.
    '''
    ansible_helper = request.config.pluginmanager.getplugin("ansible")
    return ansible_helper.initialize(request)

@pytest.fixture(scope='function')
def ansible_facts_cls(ansible_module_cls):
    '''
    Return ansible_facts dictionary
    '''
    warnings.warn("Use of ansible_facts_cls is deprecated and will be removed in a future release", DeprecationWarning)
    try:
        return ansible_module_cls.setup()
    except AnsibleHostUnreachable, e:
        log.warning("Hosts unreachable: %s" % e.dark.keys())
        return e.contacted

@pytest.fixture(scope='function')
def ansible_facts(ansible_module):
    '''
    Return ansible_facts dictionary
    '''
    try:
        return ansible_module.setup()
    except AnsibleHostUnreachable, e:
        log.warning("Hosts unreachable: %s" % e.dark.keys())
        return e.contacted
<commit_msg>should reference ansible_module_cls<commit_after>import pytest
import warnings
import logging
from .errors import AnsibleHostUnreachable

try:
    from logging import NullHandler
except ImportError:
    from logging import Handler

    class NullHandler(Handler):

        def emit(self, record):
            pass

log = logging.getLogger(__name__)
log.addHandler(NullHandler())

@pytest.fixture(scope='class')
def ansible_module_cls(request):
    '''
    Return AnsibleV1Module instance with class scope.
    '''
    warnings.warn("Use of ansible_module_cls is deprecated and will be removed in a future release", DeprecationWarning)
    ansible_helper = request.config.pluginmanager.getplugin("ansible")
    return ansible_helper.initialize(request)

@pytest.fixture(scope='function')
def ansible_module(request):
    '''
    Return AnsibleV1Module instance with function scope.
    '''
    ansible_helper = request.config.pluginmanager.getplugin("ansible")
    return ansible_helper.initialize(request)

@pytest.fixture(scope='function')
def ansible_facts_cls(ansible_module_cls):
    '''
    Return ansible_facts dictionary
    '''
    warnings.warn("Use of ansible_facts_cls is deprecated and will be removed in a future release", DeprecationWarning)
    try:
        return ansible_module_cls.setup()
    except AnsibleHostUnreachable, e:
        log.warning("Hosts unreachable: %s" % e.dark.keys())
        return e.contacted

@pytest.fixture(scope='function')
def ansible_facts(ansible_module):
    '''
    Return ansible_facts dictionary
    '''
    try:
        return ansible_module.setup()
    except AnsibleHostUnreachable, e:
        log.warning("Hosts unreachable: %s" % e.dark.keys())
        return e.contacted
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
#
# bin.py
#
# Created by José Sánchez-Gallego on 6 Nov 2016.

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

import warnings

from marvin.core.exceptions import MarvinError, MarvinUserWarning
from marvin.tools.maps import Maps, _is_MPL4
from marvin.tools.modelcube import ModelCube
from marvin.tools.spaxel import Spaxel

class Bin(object):

    def __init__(self, binid, **kwargs):

        self.plateifu = None
        self.mangaid = None

        self._maps, self._modelcube = self._get_dap_objects(**kwargs)
        self.binid = binid

        # Drops some keyword that could make spaxel load fail.
        kwargs.pop('bintype', None)
        kwargs.pop('mode', None)

        self._load_spaxels(**kwargs)
        self._load_data(**kwargs)

    def __repr__(self):
        return '<Marvin Bin (binid={0}, n_spaxels={1})>'.format(self.binid, len(self.spaxels))

    def _get_dap_objects(self, **kwargs):
        """Gets the Maps and ModelCube object."""

        try:
            kwargs_maps = kwargs.copy()
            kwargs_maps.pop('modelcube_filename', None)
            kwargs_maps['filename'] = kwargs_maps.pop('maps_filename', None)
            maps = Maps(**kwargs_maps)
        except MarvinError as ee:
            raise MarvinError('failed to open a Maps: {0}'.format(str(ee)))

        self.plateifu = maps.plateifu
        self.mangaid = maps.mangaid

        if _is_MPL4(maps._dapver):
            return maps, None

        try:
            kwargs_modelcube = kwargs.copy()
            kwargs_modelcube.pop('maps_filename', None)
            kwargs_modelcube['filename'] = kwargs_modelcube.pop('modelcube_filename', None)
            modelcube = ModelCube(**kwargs_modelcube)
        except Exception:
            warnings.warn('cannot open a ModelCube for this combination of '
                          'parameters. Some fetures will not be available.', MarvinUserWarning)
            modelcube = False

        return maps, modelcube

    def _load_spaxels(self, **kwargs):
        """Creates a list of unloaded spaxels for this binid."""

        load_spaxels = kwargs.pop('load_spaxels', False)

        self.spaxel_coords = self._maps.get_bin_spaxels(self.binid, only_list=True)

        if len(self.spaxel_coords) == 0:
            raise MarvinError('there are no spaxels associated with binid={0}.'.format(self.binid))
        else:

            if 'plateifu' not in kwargs:
                kwargs['plateifu'] = self._maps.plateifu

            modelcube_for_spaxel = False if not self._modelcube else self._modelcube.get_unbinned()
            self.spaxels = [Spaxel(x=cc[0], y=cc[1], cube=True, maps=self._maps.get_unbinned(),
                                   modelcube=modelcube_for_spaxel,
                                   load=load_spaxels,
                                   **kwargs)
                            for cc in self.spaxel_coords]

    def _load_data(self, **kwargs):
        """Loads one of the spaxels to get the DAP properties for the binid."""

        assert len(self.spaxel_coords) > 0

        sample_coords = self.spaxel_coords[0]

        if 'plateifu' not in kwargs:
            kwargs['plateifu'] = self._maps.plateifu

        sample_spaxel = Spaxel(x=sample_coords[0], y=sample_coords[1],
                               cube=True, maps=self._maps, modelcube=self._modelcube,
                               load=True, allow_binned=True, **kwargs)

        self.specres = sample_spaxel.specres
        self.specresd = sample_spaxel.specresd
        self.spectrum = sample_spaxel.spectrum
        self.properties = sample_spaxel.properties

        self.model_flux = sample_spaxel.model_flux
        self.redcorr = sample_spaxel.redcorr
        self.model = sample_spaxel.model
        self.emline = sample_spaxel.emline
        self.emline_base = sample_spaxel.emline_base
        self.stellar_continuum = sample_spaxel.stellar_continuum

    def load_all(self):
        """Loads all the spaxels that for this bin."""

        for spaxel in self.spaxels:
            spaxel.load()

        return
<commit_msg>added docstring for Bin<commit_after>#!/usr/bin/env python
# encoding: utf-8
#
# bin.py
#
# Created by José Sánchez-Gallego on 6 Nov 2016.

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

import warnings

from marvin.core.exceptions import MarvinError, MarvinUserWarning
from marvin.tools.maps import Maps, _is_MPL4
from marvin.tools.modelcube import ModelCube
from marvin.tools.spaxel import Spaxel

class Bin(object):
    """A class to interface with a bin.

    This class represents a bin of spaxels in a MaNGA DAP
    `:class:~marvin.tools.maps.Maps` or
    `:class:~marvin.tools.modelcube.ModelCube`.
    By definition, a bin is identified by a ``binid`` and can correspond to a
    binned or unbinned MAPS, although in the latter case there is no formal
    difference between a bin and a spaxel.

    When a bin is instantiated, a list of its contituent
    `:class:~marvin.tools.spaxel.Spaxel` is created. ``Spaxel`` objects are
    created unloaded for efficiency, but they can be loaded by doing, for
    example for the first spaxel, ``bin.spaxels[0].load()``.

    MAPS properties and (if available) Model Cube spectra are loaded for the
    bin.

    Parameters:
        binid (int):
            The binid of the requested bin.
        extra_kwargs:
            Any other keyword parameter necessary to instantiate the
            `:class:~marvin.tools.maps.Maps`, the
            `:class:~marvin.tools.modelcube.ModelCube`, or the
            `:class:~marvin.tools.spaxel.Spaxel` related to this bin. Refer to
            the documentation of those classes for details on the available
            keywords.

    Attributes:
        spaxels (list):
            A list of `:class:~marvin.tools.spaxel.Spaxel` objects that form
            the bin. The spaxels are created unloaded.
        properties (dict):
            Same as the `:class:~marvin.tools.maps.Maps` properties for a
            given spaxel, but in this case for the bin.
        modelcube_attributes:
            All the attributes derived for a given spaxel from the
            `:class:~marvin.tools.modelcube.ModelCube`, but in this case for
            the bin. If a modelcube is not available or it is not instantiated,
            these attributes will be ``None``.

    """

    def __init__(self, binid, **kwargs):

        self.plateifu = None
        self.mangaid = None

        self._maps, self._modelcube = self._get_dap_objects(**kwargs)
        self.binid = binid

        # Drops some keyword that could make spaxel load fail.
        kwargs.pop('bintype', None)
        kwargs.pop('mode', None)

        self._load_spaxels(**kwargs)
        self._load_data(**kwargs)

    def __repr__(self):
        return '<Marvin Bin (binid={0}, n_spaxels={1})>'.format(self.binid, len(self.spaxels))

    def _get_dap_objects(self, **kwargs):
        """Gets the Maps and ModelCube object."""

        try:
            kwargs_maps = kwargs.copy()
            kwargs_maps.pop('modelcube_filename', None)
            kwargs_maps['filename'] = kwargs_maps.pop('maps_filename', None)
            maps = Maps(**kwargs_maps)
        except MarvinError as ee:
            raise MarvinError('failed to open a Maps: {0}'.format(str(ee)))

        self.plateifu = maps.plateifu
        self.mangaid = maps.mangaid

        if _is_MPL4(maps._dapver):
            return maps, None

        try:
            kwargs_modelcube = kwargs.copy()
            kwargs_modelcube.pop('maps_filename', None)
            kwargs_modelcube['filename'] = kwargs_modelcube.pop('modelcube_filename', None)
            modelcube = ModelCube(**kwargs_modelcube)
        except Exception:
            warnings.warn('cannot open a ModelCube for this combination of '
                          'parameters. Some fetures will not be available.', MarvinUserWarning)
            modelcube = False

        return maps, modelcube

    def _load_spaxels(self, **kwargs):
        """Creates a list of unloaded spaxels for this binid."""

        load_spaxels = kwargs.pop('load_spaxels', False)

        self.spaxel_coords = self._maps.get_bin_spaxels(self.binid, only_list=True)

        if len(self.spaxel_coords) == 0:
            raise MarvinError('there are no spaxels associated with binid={0}.'.format(self.binid))
        else:

            if 'plateifu' not in kwargs:
                kwargs['plateifu'] = self._maps.plateifu

            modelcube_for_spaxel = False if not self._modelcube else self._modelcube.get_unbinned()
            self.spaxels = [Spaxel(x=cc[0], y=cc[1], cube=True, maps=self._maps.get_unbinned(),
                                   modelcube=modelcube_for_spaxel,
                                   load=load_spaxels,
                                   **kwargs)
                            for cc in self.spaxel_coords]

    def _load_data(self, **kwargs):
        """Loads one of the spaxels to get the DAP properties for the binid."""

        assert len(self.spaxel_coords) > 0

        sample_coords = self.spaxel_coords[0]

        if 'plateifu' not in kwargs:
            kwargs['plateifu'] = self._maps.plateifu

        sample_spaxel = Spaxel(x=sample_coords[0], y=sample_coords[1],
                               cube=True, maps=self._maps, modelcube=self._modelcube,
                               load=True, allow_binned=True, **kwargs)

        self.specres = sample_spaxel.specres
        self.specresd = sample_spaxel.specresd
        self.spectrum = sample_spaxel.spectrum
        self.properties = sample_spaxel.properties

        self.model_flux = sample_spaxel.model_flux
        self.redcorr = sample_spaxel.redcorr
        self.model = sample_spaxel.model
        self.emline = sample_spaxel.emline
        self.emline_base = sample_spaxel.emline_base
        self.stellar_continuum = sample_spaxel.stellar_continuum

    def load_all(self):
        """Loads all the spaxels that for this bin."""

        for spaxel in self.spaxels:
            spaxel.load()

        return
<|endoftext|>


<commit_before>#!/usr/bin/python
"""
@package nilib
@file redisflush.py
@brief To flush the entire Redis database used by niserver.
@version $Revision: 1.01 $ $Author: elwynd $
@version Copyright (C) 2012 Trinity College Dublin and Folly Consulting Ltd
      This is an adjunct to the NI URI library developed as
      part of the SAIL project. (http://sail-project.eu)

      Specification(s) - note, versions may change
          - http://tools.ietf.org/html/draft-farrell-decade-ni-10
          - http://tools.ietf.org/html/draft-hallambaker-decade-ni-params-03
          - http://tools.ietf.org/html/draft-kutscher-icnrg-netinf-proto-00

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

       - http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

================================================================================

@details
Clear out the Redis database and the corresponding object storage cache.

@code
Revision History
================
Version   Date       Author         Notes
1.0       13/12/2012 Elwyn Davies   Create.

@endcode
"""

#==============================================================================#

import redis
import os

STORAGE_ROOT_KEY = "NISERVER_STORAGE_ROOT"
try:
    ans = raw_input("This will flush the Redis database.  Are you sure you want to continue? [y/n] " )
except:
    print "Interrupted"
    os._exit(1)
if ans == "y":
    c = redis.StrictRedis()
    try:
        storage_root = c.get(STORAGE_ROOT_KEY)
        if c.flushdb():
            print "Database flushed"
            # Now go delete the object cache
        else:
            print "Database flushing failed"
            os._exit(1)
    except Exception, e:
        print "Connection to database failed: %s" % str(e)
        os._exit(1)
else:
    print "Abandoned flushing"
    os._exit(1)

if storage_root is None:
    print "Key %s was not set in database - need to manually remove object cache" % \
          STORAGE_ROOT_KEY
    os._exit(1)

if os.path.isdir(storage_root):
    ans = raw_input("OK to delete storage cache from %s? [y/n] " % storage_root)
    if ans == "y" :
        os.system("rm -rf %s/*" % storage_root)
        os._exit(0)
    else:
        print "Abandoned deleteing object cache"
        os._exit(1)
else:
    print "Storage root directory %s does not exist" % storage_root
    os._exit(1)

<commit_msg>Added database number parameter to redisflush.py.<commit_after>#!/usr/bin/python
"""
@package nilib
@file redisflush.py
@brief To flush the entire Redis database used by niserver.
@version $Revision: 1.01 $ $Author: elwynd $
@version Copyright (C) 2012 Trinity College Dublin and Folly Consulting Ltd
      This is an adjunct to the NI URI library developed as
      part of the SAIL project. (http://sail-project.eu)

      Specification(s) - note, versions may change
          - http://tools.ietf.org/html/draft-farrell-decade-ni-10
          - http://tools.ietf.org/html/draft-hallambaker-decade-ni-params-03
          - http://tools.ietf.org/html/draft-kutscher-icnrg-netinf-proto-00

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

       - http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

================================================================================

@details
Clear out the Redis database and the corresponding object storage cache.

@code
Revision History
================
Version   Date       Author         Notes
1.1       25/01/2013 Elwyn Davies   Add parameter to select database to flush.
1.0       13/12/2012 Elwyn Davies   Create.

@endcode
"""

#==============================================================================#

import redis
import os
import sys

STORAGE_ROOT_KEY = "NISERVER_STORAGE_ROOT"

# Get database number from command line (if given)
if len(sys.argv) > 1:
    try:
        redis_db = int(sys.argv[1])
    except Exception, e:
        print "Usage: %s [<Redis database - integer>]" % sys.argv[0]
        os._exit(1)
else:
    redis_db = 0

try:
    ans = raw_input("This will flush the Redis database #%d.  Are you sure you want to continue? [y/n] " %
                    redis_db)
except:
    print "Interrupted"
    os._exit(1)
if ans == "y":
    c = redis.StrictRedis(db=redis_db)
    try:
        storage_root = c.get(STORAGE_ROOT_KEY)
        if c.flushdb():
            print "Database flushed"
            # Now go delete the object cache
        else:
            print "Database flushing failed"
            os._exit(1)
    except Exception, e:
        print "Connection to database failed: %s" % str(e)
        os._exit(1)
else:
    print "Abandoned flushing"
    os._exit(1)

if storage_root is None:
    print "Key %s was not set in database - need to manually remove object cache" % \
          STORAGE_ROOT_KEY
    os._exit(1)

if os.path.isdir(storage_root):
    ans = raw_input("OK to delete storage cache from %s? [y/n] " % storage_root)
    if ans == "y" :
        os.system("rm -rf %s/*" % storage_root)
        os._exit(0)
    else:
        print "Abandoned deleting object cache"
        os._exit(1)
else:
    print "Storage root directory %s does not exist" % storage_root
    os._exit(1)

<|endoftext|>


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

from __future__ import print_function
import vcf
import argparse
import numpy
import banyan
import gzip

#Functions
def overlapValid((s1, e1), (s2, e2), reciprocalOverlap=0.5, maxOffset=500):
    if (e1 < s2) or (s1 > e2):
        return False
    overlapLen = float(min(e1, e2) - max(s1, s2))
    # Check reciprocal overlap
    if (overlapLen < 1) or (float(e1-s1)/overlapLen < reciprocalOverlap) or (float(e2-s2)/overlapLen < reciprocalOverlap):
        return False
    # Check offset
    if (abs(s2-s1) > maxOffset) or (abs(e2-e1) > maxOffset):
        return False
    return True

# Parse command line
parser = argparse.ArgumentParser(description='Filter for reliable SV sites.')
parser.add_argument('-v', '--vcf', metavar='variants.vcf', required=True, dest='vcfFile', help='input vcf file (required)')
parser.add_argument('-o', '--out', metavar='out.vcf', required=True, dest='outFile', help='output vcf file (required)')
parser.add_argument('-g', '--gqAlt', metavar='15', required=False, dest='gqAlt', help='min. GQ for carriers (optional)')
parser.add_argument('-q', '--gqRef', metavar='15', required=False, dest='gqRef', help='min. GQ for non-carriers (optional)')
parser.add_argument('-m', '--minsize', metavar='500', required=False, dest='minSize', help='min. size (optional)')
parser.add_argument('-n', '--maxsize', metavar='5000000', required=False, dest='maxSize', help='max. size (optional)')
parser.add_argument('-a', '--altaf', metavar='0.4', required=False, dest='altAF', help='min. alt. AF (optional)')
parser.add_argument('-r', '--ratioGeno', metavar='0.4', required=False, dest='ratioGeno', help='min. fraction of genotyped samples (optional)')
parser.add_argument('-s', '--sample', metavar='NA12878', required=False, dest='sampleID', help='required carrier sample (optional)')
parser.add_argument('-f', '--filter', dest='siteFilter', action='store_true', help='Filter sites for PASS')
parser.add_argument('-p', '--paired', dest='pairedFilter', action='store_true', help='Require 3to3 and 5to5 inversion support')
args = parser.parse_args()

# Command-line args
sampleID = ""
if args.sampleID:
    sampleID = args.sampleID
gqAltCut = 15
if args.gqAlt:
    gqAltCut = int(args.gqAlt)
gqRefCut = 15
if args.gqRef:
    gqRefCut = int(args.gqRef)
minSize = 500
if args.minSize:
    minSize = int(args.minSize)
maxSize = 5000000
if args.maxSize:
    maxSize = int(args.maxSize)
altAF = 0.4
if args.altAF:
    altAF = float(args.altAF)
ratioGeno = 0.4
if args.ratioGeno:
    ratioGeno = float(args.ratioGeno)

# Collect high-quality SVs
sv = dict()
if args.vcfFile:
    vcfIn = gzip.open(args.vcfFile) if args.vcfFile.endswith('.gz') else open(args.vcfFile)
    vcf_reader = vcf.Reader(vcfIn)
    for record in vcf_reader:
        if (record.INFO['SVLEN'] >= minSize) and (record.INFO['SVLEN'] <= maxSize) and ((not args.siteFilter) or (len(record.FILTER) == 0)):
            precise = False
            if 'PRECISE' in record.INFO.keys():
                precise = record.INFO['PRECISE']
            gqRef = []
            gqAlt = []
            ratioRef = [0]
            ratioAlt = []
            carrierSample = False
            if sampleID == "":
                carrierSample = True
            for call in record.samples:
                if call.called:
                    if call.gt_type == 0:
                        if precise:
                            ratioRef.append(float(call['RV'])/float(call['RR'] + call['RV']))
                        else:
                            ratioRef.append(float(call['DV'])/float(call['DR'] + call['DV']))
                        if ((precise) and (call['RV'] == 0)) or ((not precise) and (call['DV'] == 0)):
                            gqRef.append(call['GQ'])
                    if call.gt_type != 0:
                        if precise:
                            ratioAlt.append(float(call['RV'])/float(call['RR'] + call['RV']))
                        else:
                            ratioAlt.append(float(call['DV'])/float(call['DR'] + call['DV']))
                        if (not carrierSample) and (call.sample == sampleID):
                            carrierSample = True
                        if ((precise) and (call['RV'] >= 2)) or ((not precise) and (call['DV'] >= 2)):
                            gqAlt.append(call['GQ'])
            genotypeRatio = float(len(gqAlt)+len(gqRef)) /  float(len(record.samples))
            if (carrierSample) and (genotypeRatio > ratioGeno):
                if (len(gqRef)) and (len(gqAlt)) and (numpy.median(gqRef) >= gqRefCut) and (numpy.median(gqAlt) >= gqAltCut):
                    if (numpy.percentile(ratioRef, 99) == 0) and (numpy.median(ratioAlt) >= altAF):
                        #print(record.INFO['END']-record.POS, len(gqRef), len(gqAlt), numpy.median(gqRef), numpy.median(gqAlt), numpy.percentile(ratioRef, 95), numpy.median(ratioAlt), genotypeRatio, sep="\t")
                        if not sv.has_key(record.CHROM):
                            sv[record.CHROM] = banyan.SortedDict(key_type=(int, int), alg=banyan.RED_BLACK_TREE, updator=banyan.OverlappingIntervalsUpdator)
                        sv[record.CHROM][(record.POS, record.INFO['END'])] = (record.ID, record.INFO['PE'], record.INFO['CT'])

# Kick-out the unpaired SVs
if args.pairedFilter:
    filteredSVs = dict()
    for chrName in sv.keys():
        for start, end in sv[chrName].keys():
            overlapList = sv[chrName].overlap((start, end))
            svID, score, ct = sv[chrName][(start, end)]
            for cStart, cEnd in overlapList:
                cSvID, cScore, cCt = sv[chrName][(cStart, cEnd)]
                if (svID != cSvID) and (ct != cCt) and (overlapValid((start, end), (cStart, cEnd), 0.9, 100)):
                    if not filteredSVs.has_key(chrName):
                        filteredSVs[chrName] = banyan.SortedDict(key_type=(int, int), alg=banyan.RED_BLACK_TREE, updator=banyan.OverlappingIntervalsUpdator)
                    filteredSVs[chrName][(start, end)] = sv[chrName][(start, end)]
                    break
    sv = filteredSVs

# Output vcf records
if args.vcfFile:
    vcf_reader = vcf.Reader(open(args.vcfFile), 'r')
    vcf_writer = vcf.Writer(open(args.outFile, 'w'), vcf_reader, lineterminator='\n')
    for record in vcf_reader:
        if (record.CHROM not in sv.keys()) or ((record.POS, record.INFO['END']) not in sv[record.CHROM].keys()):
            continue
        overlapList = sv[record.CHROM].overlap((record.POS, record.INFO['END']))
        svID, score, ct = sv[record.CHROM][(record.POS, record.INFO['END'])]
        foundBetterHit = False
        for cStart, cEnd in overlapList:
            cSvID, cScore, cCt = sv[record.CHROM][(cStart, cEnd)]
            if svID == cSvID:
                continue
            # There should be at least 10% overlap otherwise ignore it
            if not overlapValid((record.POS, record.INFO['END']), (cStart, cEnd), 0.1, 10000000):
                continue
            if (cScore > score) or ((cScore == score) and (cSvID < svID)):
                foundBetterHit = True
                break
        if not foundBetterHit:
            vcf_writer.write_record(record)
<commit_msg>Fixed vcf reader problem.<commit_after>#! /usr/bin/env python

from __future__ import print_function
import vcf
import argparse
import numpy
import banyan
import gzip

#Functions
def overlapValid((s1, e1), (s2, e2), reciprocalOverlap=0.5, maxOffset=500):
    if (e1 < s2) or (s1 > e2):
        return False
    overlapLen = float(min(e1, e2) - max(s1, s2))
    # Check reciprocal overlap
    if (overlapLen < 1) or (float(e1-s1)/overlapLen < reciprocalOverlap) or (float(e2-s2)/overlapLen < reciprocalOverlap):
        return False
    # Check offset
    if (abs(s2-s1) > maxOffset) or (abs(e2-e1) > maxOffset):
        return False
    return True

# Parse command line
parser = argparse.ArgumentParser(description='Filter for reliable SV sites.')
parser.add_argument('-v', '--vcf', metavar='variants.vcf', required=True, dest='vcfFile', help='input vcf file (required)')
parser.add_argument('-o', '--out', metavar='out.vcf', required=True, dest='outFile', help='output vcf file (required)')
parser.add_argument('-g', '--gqAlt', metavar='15', required=False, dest='gqAlt', help='min. GQ for carriers (optional)')
parser.add_argument('-q', '--gqRef', metavar='15', required=False, dest='gqRef', help='min. GQ for non-carriers (optional)')
parser.add_argument('-m', '--minsize', metavar='500', required=False, dest='minSize', help='min. size (optional)')
parser.add_argument('-n', '--maxsize', metavar='5000000', required=False, dest='maxSize', help='max. size (optional)')
parser.add_argument('-a', '--altaf', metavar='0.4', required=False, dest='altAF', help='min. alt. AF (optional)')
parser.add_argument('-r', '--ratioGeno', metavar='0.4', required=False, dest='ratioGeno', help='min. fraction of genotyped samples (optional)')
parser.add_argument('-s', '--sample', metavar='NA12878', required=False, dest='sampleID', help='required carrier sample (optional)')
parser.add_argument('-f', '--filter', dest='siteFilter', action='store_true', help='Filter sites for PASS')
parser.add_argument('-p', '--paired', dest='pairedFilter', action='store_true', help='Require 3to3 and 5to5 inversion support')
args = parser.parse_args()

# Command-line args
sampleID = ""
if args.sampleID:
    sampleID = args.sampleID
gqAltCut = 15
if args.gqAlt:
    gqAltCut = int(args.gqAlt)
gqRefCut = 15
if args.gqRef:
    gqRefCut = int(args.gqRef)
minSize = 500
if args.minSize:
    minSize = int(args.minSize)
maxSize = 5000000
if args.maxSize:
    maxSize = int(args.maxSize)
altAF = 0.4
if args.altAF:
    altAF = float(args.altAF)
ratioGeno = 0.4
if args.ratioGeno:
    ratioGeno = float(args.ratioGeno)

# Collect high-quality SVs
sv = dict()
if args.vcfFile:
    vcf_reader = vcf.Reader(open(args.vcfFile), 'r', compressed=True) if args.vcfFile.endswith('.gz') else vcf.Reader(open(args.vcfFile), 'r', compressed=False)
    for record in vcf_reader:
        if (record.INFO['SVLEN'] >= minSize) and (record.INFO['SVLEN'] <= maxSize) and ((not args.siteFilter) or (len(record.FILTER) == 0)):
            precise = False
            if 'PRECISE' in record.INFO.keys():
                precise = record.INFO['PRECISE']
            gqRef = []
            gqAlt = []
            ratioRef = [0]
            ratioAlt = []
            carrierSample = False
            if sampleID == "":
                carrierSample = True
            for call in record.samples:
                if call.called:
                    if call.gt_type == 0:
                        if precise:
                            ratioRef.append(float(call['RV'])/float(call['RR'] + call['RV']))
                        else:
                            ratioRef.append(float(call['DV'])/float(call['DR'] + call['DV']))
                        if ((precise) and (call['RV'] == 0)) or ((not precise) and (call['DV'] == 0)):
                            gqRef.append(call['GQ'])
                    if call.gt_type != 0:
                        if precise:
                            ratioAlt.append(float(call['RV'])/float(call['RR'] + call['RV']))
                        else:
                            ratioAlt.append(float(call['DV'])/float(call['DR'] + call['DV']))
                        if (not carrierSample) and (call.sample == sampleID):
                            carrierSample = True
                        if ((precise) and (call['RV'] >= 2)) or ((not precise) and (call['DV'] >= 2)):
                            gqAlt.append(call['GQ'])
            genotypeRatio = float(len(gqAlt)+len(gqRef)) /  float(len(record.samples))
            if (carrierSample) and (genotypeRatio > ratioGeno):
                if (len(gqRef)) and (len(gqAlt)) and (numpy.median(gqRef) >= gqRefCut) and (numpy.median(gqAlt) >= gqAltCut):
                    if (numpy.percentile(ratioRef, 99) == 0) and (numpy.median(ratioAlt) >= altAF):
                        #print(record.INFO['END']-record.POS, len(gqRef), len(gqAlt), numpy.median(gqRef), numpy.median(gqAlt), numpy.percentile(ratioRef, 95), numpy.median(ratioAlt), genotypeRatio, sep="\t")
                        if not sv.has_key(record.CHROM):
                            sv[record.CHROM] = banyan.SortedDict(key_type=(int, int), alg=banyan.RED_BLACK_TREE, updator=banyan.OverlappingIntervalsUpdator)
                        sv[record.CHROM][(record.POS, record.INFO['END'])] = (record.ID, record.INFO['PE'], record.INFO['CT'])

# Kick-out the unpaired SVs
if args.pairedFilter:
    filteredSVs = dict()
    for chrName in sv.keys():
        for start, end in sv[chrName].keys():
            overlapList = sv[chrName].overlap((start, end))
            svID, score, ct = sv[chrName][(start, end)]
            for cStart, cEnd in overlapList:
                cSvID, cScore, cCt = sv[chrName][(cStart, cEnd)]
                if (svID != cSvID) and (ct != cCt) and (overlapValid((start, end), (cStart, cEnd), 0.9, 100)):
                    if not filteredSVs.has_key(chrName):
                        filteredSVs[chrName] = banyan.SortedDict(key_type=(int, int), alg=banyan.RED_BLACK_TREE, updator=banyan.OverlappingIntervalsUpdator)
                    filteredSVs[chrName][(start, end)] = sv[chrName][(start, end)]
                    break
    sv = filteredSVs

# Output vcf records
if args.vcfFile:
    vcf_reader = vcf.Reader(open(args.vcfFile), 'r', compressed=True) if args.vcfFile.endswith('.gz') else vcf.Reader(open(args.vcfFile), 'r', compressed=False)
    vcf_writer = vcf.Writer(open(args.outFile, 'w'), vcf_reader, lineterminator='\n')
    for record in vcf_reader:
        if (record.CHROM not in sv.keys()) or ((record.POS, record.INFO['END']) not in sv[record.CHROM].keys()):
            continue
        overlapList = sv[record.CHROM].overlap((record.POS, record.INFO['END']))
        svID, score, ct = sv[record.CHROM][(record.POS, record.INFO['END'])]
        foundBetterHit = False
        for cStart, cEnd in overlapList:
            cSvID, cScore, cCt = sv[record.CHROM][(cStart, cEnd)]
            if svID == cSvID:
                continue
            # There should be at least 10% overlap otherwise ignore it
            if not overlapValid((record.POS, record.INFO['END']), (cStart, cEnd), 0.1, 10000000):
                continue
            if (cScore > score) or ((cScore == score) and (cSvID < svID)):
                foundBetterHit = True
                break
        if not foundBetterHit:
            vcf_writer.write_record(record)
<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function, unicode_literals
import struct
import math
from subprocess import Popen, PIPE
import tempfile
import os

from scannerpy.stdlib import parsers
from scannerpy.common import *
from scannerpy.job import Job
from scannerpy.bulk_job import BulkJob

class Column(object):
    """
    A column of a Table.
    """

    def __init__(self, table, name):
        self._table = table
        self._name = name
        self._db = table._db
        self._storage = table._db.config.storage
        self._db_path = table._db.config.db_path

        self._loaded = False
        self._descriptor = None
        self._video_descriptor = None

    def _load_meta(self):
        if not self._loaded:
            self._loaded = True
            descriptor, video_descriptor = self._table._load_column(self._name)
            self._descriptor = descriptor
            self._video_descriptor = video_descriptor

    def name(self):
        return self._name

    def type(self):
        self._load_meta()
        return self._descriptor.type

    def id(self):
        self._load_meta()
        return self._descriptor.id

    def keyframes(self):
        self._load_meta()
        if (self._descriptor.type == self._db.protobufs.Video and
            self._video_descriptor.codec_type ==
            self._db.protobufs.VideoDescriptor.H264):
            # For each encoded video, add start frame offset
            frame_offset = 0
            kf_offset = 0
            keyframes = []
            for frames_per_video, kfs_per_video in zip(
                    self._video_descriptor.frames_per_video,
                    self._video_descriptor.keyframes_per_video):
                keyframes += [
                    frame_offset + kfi
                    for kfi in self._video_descriptor.keyframe_indices[
                            kf_offset:kf_offset + kfs_per_video]]
                frame_offset += frames_per_video
                kf_offset += kfs_per_video
            return keyframes
        else:
            return range(self._table.num_rows())

    def _load_output_file(self, item_id, rows, fn=None):
        assert len(rows) > 0

        metadata_path = '{}/tables/{}/{}_{}_metadata.bin'.format(
            self._db_path, self._table._descriptor.id,
            self._descriptor.id, item_id)
        try:
            metadata_contents = self._storage.read(
                metadata_path.encode('ascii'))
        except UserWarning:
            raise ScannerException('Path {} does not exist'.format(
                metadata_path))

        path = '{}/tables/{}/{}_{}.bin'.format(
            self._db_path, self._table._descriptor.id,
            self._descriptor.id, item_id)
        try:
            contents = self._storage.read(path.encode('ascii'))
        except UserWarning:
            raise ScannerException('Path {} does not exist'.format(path))

        lens = []
        total_rows = 0
        i = 0
        while i < len(metadata_contents):
            (num_rows,) = struct.unpack("=Q", metadata_contents[i:i+8])
            total_rows += num_rows
            i += 8
            for fi in range(num_rows):
                (buf_len,) = struct.unpack("=Q", metadata_contents[i:i+8])
                lens.append(buf_len)
                i += 8

        start_pos = None
        pos = 0
        rows = rows if len(rows) > 0 else range(total_rows)
        for fi in range(total_rows):
            old_pos = pos
            pos += lens[fi]
            if start_pos is None:
                start_pos = old_pos

        rows_idx = 0
        i = start_pos
        for j, buf_len in enumerate(lens):
            if rows_idx < len(rows) and j == rows[rows_idx]:
                buf = contents[i:i+buf_len]
                # len(buf) == 0 when element is null
                if len(buf) == 0:
                    yield None
                elif fn is not None:
                    yield fn(buf, self._db.protobufs)
                else:
                    yield buf
                rows_idx += 1
            i += buf_len

    def _load(self, fn=None, rows=None):
        table_descriptor = self._table._descriptor
        total_rows = table_descriptor.end_rows[-1]

        # Integer divide, round up
        num_items = len(table_descriptor.end_rows)
        bufs = []
        input_rows = list(range(self._table.num_rows()))
        assert len(input_rows) == total_rows
        i = 0
        rows_so_far = 0
        rows_idx = 0
        rows = range(total_rows) if rows is None else rows
        prev = 0
        for item_id in range(num_items):
            start_row = prev
            end_row = table_descriptor.end_rows[item_id]
            item_rows = end_row - start_row
            prev = end_row
            select_rows = []
            while rows_idx < len(rows):
                r = rows[rows_idx]
                if r >= start_row and r < end_row:
                    select_rows.append(r - start_row)
                    rows_idx += 1
                else:
                    break
            if select_rows:
                for output in self._load_output_file(item_id, select_rows, fn):
                    yield (input_rows[i], output)
                    i += 1
            rows_so_far += item_rows

    # TODO(wcrichto): don't show progress bar when running decode png
    def load(self, fn=None, rows=None):
        """
        Loads the results of a Scanner computation into Python.

        Kwargs:
            fn: Optional function to apply to the binary blobs as they are read
                in.

        Returns:
            Generator that yields either a numpy array for frame columns or
            a binary blob for non-frame columns (optionally processed by the
            `fn`).
        """

        self._load_meta()
        # If the column is a video, then dump the requested frames to disk as
        # PNGs and return the decoded PNGs
        if (self._descriptor.type == self._db.protobufs.Video and
            self._video_descriptor.codec_type ==
            self._db.protobufs.VideoDescriptor.H264):
            png_table_name = self._db._png_dump_prefix.format(self._table.name())
            if self._db.has_table(png_table_name):
                png_table = self._db.table(png_table_name)
                if rows is None and \
                   png_table.num_rows() == self._table.num_rows() and \
                   png_table._descriptor.timestamp > \
                   self._table._descriptor.timestamp:
                    return png_table.load(['img'], parsers.image)
            pair = [(self._table.name(), png_table_name)]
            op_args = {}
            frame = self._db.ops.FrameInput()
            op_args[frame] = self
            enc_input = frame
            if rows is not None:
                sampled_frame = frame.sample()
                op_args[sampled_frame] = self._db.sampler.gather(rows)
                enc_input = sampled_frame
            img = self._db.ops.ImageEncoder(frame = enc_input)
            output_op = self._db.ops.Output(columns=[img])
            op_args[output_op] = png_table_name
            job = Job(op_args=op_args)
            bulk_job = BulkJob(output=output_op, jobs=[job])
            [out_tbl] = self._db.run(bulk_job, force=True, show_progress=False)
            return out_tbl.load(['img'], parsers.image)
        elif self._descriptor.type == self._db.protobufs.Video:
            frame_type = self._video_descriptor.frame_type
            if frame_type == self._db.protobufs.U8:
                dtype = np.uint8
            elif frame_type == self._db.protobufs.F32:
                dtype = np.float32
            elif frame_type == self._db.protobufs.F64:
                dtype = np.float64
            parser_fn = parsers.raw_frame_gen(self._video_descriptor.height,
                                              self._video_descriptor.width,
                                              self._video_descriptor.channels,
                                              dtype)
            return self._load(fn=parser_fn, rows=rows)
        else:
            return self._load(fn, rows=rows)

    def save_mp4(self, output_name, fps=None, scale=None):
        self._load_meta()
        if not (self._descriptor.type == self._db.protobufs.Video and
                self._video_descriptor.codec_type ==
                self._db.protobufs.VideoDescriptor.H264):
            raise ScannerException('Attempted to save a non-h264-compressed '
                                   'column as an mp4. Try compressing the '
                                   'column first by saving the output as '
                                   'an RGB24 frame')
        num_items = len(self._table._descriptor.end_rows)

        paths = ['{}/tables/{:d}/{:d}_{:d}.bin'.format(
            self._db._db_path,
            self._table._descriptor.id, self._descriptor.id, item_id)
                          for item_id in range(num_items)]
        temp_paths = []
        for _ in range(len(paths)):
            fd, p = tempfile.mkstemp()
            os.close(fd)
            temp_paths.append(p)
        # Copy all files locally before calling ffmpeg
        for in_path, temp_path in zip(paths, temp_paths):
            with open(temp_path, 'w') as f:
                f.write(self._storage.read(in_path.encode('ascii')))

        files = '|'.join(temp_paths)

        vid_fps = (fps or
                   (1.0/(self._video_descriptor.time_base_num /
                         float(self._video_descriptor.time_base_denom))))

        args = ''
        if scale:
            args += '-filter:v "scale={:d}x{:d}" '.format(scale[0], scale[1])

        cmd = (
            'ffmpeg -y '
            '-r {fps:f} ' # set the input fps
            '-i "concat:{input_files:s}" ' # concatenate the h264 files
            '-c:v libx264 '
            '-filter:v "setpts=N" ' # h264 does not have pts' in it
            '{extra_args:s}'
            '{output_name:s}.mp4'.format(
                input_files = files,
                fps = vid_fps,
                extra_args = args,
                output_name=output_name))
        rc = Popen(cmd, shell=True).wait()
        if rc != 0:
            raise ScannerException('ffmpeg failed during mp4 export!')
<commit_msg>Optimize sparse loading via Python API<commit_after>from __future__ import absolute_import, division, print_function, unicode_literals
import struct
import math
from subprocess import Popen, PIPE
import tempfile
import os

from storehouse import RandomReadFile
from scannerpy.stdlib import parsers
from scannerpy.common import *
from scannerpy.job import Job
from scannerpy.bulk_job import BulkJob

LOAD_SPARSITY_THRESHOLD = 10

class Column(object):
    """
    A column of a Table.
    """

    def __init__(self, table, name):
        self._table = table
        self._name = name
        self._db = table._db
        self._storage = table._db.config.storage
        self._db_path = table._db.config.db_path

        self._loaded = False
        self._descriptor = None
        self._video_descriptor = None

    def _load_meta(self):
        if not self._loaded:
            self._loaded = True
            descriptor, video_descriptor = self._table._load_column(self._name)
            self._descriptor = descriptor
            self._video_descriptor = video_descriptor

    def name(self):
        return self._name

    def type(self):
        self._load_meta()
        return self._descriptor.type

    def id(self):
        self._load_meta()
        return self._descriptor.id

    def keyframes(self):
        self._load_meta()
        if (self._descriptor.type == self._db.protobufs.Video and
            self._video_descriptor.codec_type ==
            self._db.protobufs.VideoDescriptor.H264):
            # For each encoded video, add start frame offset
            frame_offset = 0
            kf_offset = 0
            keyframes = []
            for frames_per_video, kfs_per_video in zip(
                    self._video_descriptor.frames_per_video,
                    self._video_descriptor.keyframes_per_video):
                keyframes += [
                    frame_offset + kfi
                    for kfi in self._video_descriptor.keyframe_indices[
                            kf_offset:kf_offset + kfs_per_video]]
                frame_offset += frames_per_video
                kf_offset += kfs_per_video
            return keyframes
        else:
            return range(self._table.num_rows())

    def _load_output_file(self, item_id, rows, fn=None):
        assert len(rows) > 0

        metadata_path = '{}/tables/{}/{}_{}_metadata.bin'.format(
            self._db_path, self._table._descriptor.id,
            self._descriptor.id, item_id)
        try:
            metadata_contents = self._storage.read(
                metadata_path.encode('ascii'))
        except UserWarning:
            raise ScannerException('Path {} does not exist'.format(
                metadata_path))

        lens = []
        total_rows = 0
        i = 0
        while i < len(metadata_contents):
            (num_rows,) = struct.unpack("=Q", metadata_contents[i:i+8])
            total_rows += num_rows
            i += 8
            for fi in range(num_rows):
                (buf_len,) = struct.unpack("=Q", metadata_contents[i:i+8])
                lens.append(buf_len)
                i += 8

        data_path = '{}/tables/{}/{}_{}.bin'.format(
            self._db_path, self._table._descriptor.id,
            self._descriptor.id, item_id)
        sparse_load = len(rows) < LOAD_SPARSITY_THRESHOLD
        try:
            if not sparse_load:
                contents = self._storage.read(data_path.encode('ascii'))
            else:
                fil = RandomReadFile(self._storage, data_path.encode('ascii'))
        except UserWarning:
            raise ScannerException('Path {} does not exist'.format(path))

        start_pos = None
        pos = 0
        rows = rows if len(rows) > 0 else range(total_rows)
        for fi in range(total_rows):
            old_pos = pos
            pos += lens[fi]
            if start_pos is None:
                start_pos = old_pos

        rows_idx = 0
        i = start_pos
        for j, buf_len in enumerate(lens):
            if rows_idx < len(rows) and j == rows[rows_idx]:
                if sparse_load:
                    fil.seek(i)
                    buf = fil.read(buf_len)
                else:
                    buf = contents[i:i+buf_len]

                # len(buf) == 0 when element is null
                if len(buf) == 0:
                    yield None
                elif fn is not None:
                    yield fn(buf, self._db.protobufs)
                else:
                    yield buf
                rows_idx += 1
            i += buf_len

    def _load(self, fn=None, rows=None):
        table_descriptor = self._table._descriptor
        total_rows = table_descriptor.end_rows[-1]

        # Integer divide, round up
        num_items = len(table_descriptor.end_rows)
        bufs = []
        input_rows = list(range(self._table.num_rows()))
        assert len(input_rows) == total_rows
        i = 0
        rows_so_far = 0
        rows_idx = 0
        rows = range(total_rows) if rows is None else rows
        prev = 0
        for item_id in range(num_items):
            start_row = prev
            end_row = table_descriptor.end_rows[item_id]
            item_rows = end_row - start_row
            prev = end_row
            select_rows = []
            while rows_idx < len(rows):
                r = rows[rows_idx]
                if r >= start_row and r < end_row:
                    select_rows.append(r - start_row)
                    rows_idx += 1
                else:
                    break
            if select_rows:
                for output in self._load_output_file(item_id, select_rows, fn):
                    yield (input_rows[i], output)
                    i += 1
            rows_so_far += item_rows

    # TODO(wcrichto): don't show progress bar when running decode png
    def load(self, fn=None, rows=None):
        """
        Loads the results of a Scanner computation into Python.

        Kwargs:
            fn: Optional function to apply to the binary blobs as they are read
                in.

        Returns:
            Generator that yields either a numpy array for frame columns or
            a binary blob for non-frame columns (optionally processed by the
            `fn`).
        """

        self._load_meta()
        # If the column is a video, then dump the requested frames to disk as
        # PNGs and return the decoded PNGs
        if (self._descriptor.type == self._db.protobufs.Video and
            self._video_descriptor.codec_type ==
            self._db.protobufs.VideoDescriptor.H264):
            png_table_name = self._db._png_dump_prefix.format(self._table.name())
            if self._db.has_table(png_table_name):
                png_table = self._db.table(png_table_name)
                if rows is None and \
                   png_table.num_rows() == self._table.num_rows() and \
                   png_table._descriptor.timestamp > \
                   self._table._descriptor.timestamp:
                    return png_table.load(['img'], parsers.image)
            pair = [(self._table.name(), png_table_name)]
            op_args = {}
            frame = self._db.ops.FrameInput()
            op_args[frame] = self
            enc_input = frame
            if rows is not None:
                sampled_frame = frame.sample()
                op_args[sampled_frame] = self._db.sampler.gather(rows)
                enc_input = sampled_frame
            img = self._db.ops.ImageEncoder(frame = enc_input)
            output_op = self._db.ops.Output(columns=[img])
            op_args[output_op] = png_table_name
            job = Job(op_args=op_args)
            bulk_job = BulkJob(output=output_op, jobs=[job])
            [out_tbl] = self._db.run(bulk_job, force=True, show_progress=False)
            return out_tbl.load(['img'], parsers.image)
        elif self._descriptor.type == self._db.protobufs.Video:
            frame_type = self._video_descriptor.frame_type
            if frame_type == self._db.protobufs.U8:
                dtype = np.uint8
            elif frame_type == self._db.protobufs.F32:
                dtype = np.float32
            elif frame_type == self._db.protobufs.F64:
                dtype = np.float64
            parser_fn = parsers.raw_frame_gen(self._video_descriptor.height,
                                              self._video_descriptor.width,
                                              self._video_descriptor.channels,
                                              dtype)
            return self._load(fn=parser_fn, rows=rows)
        else:
            return self._load(fn, rows=rows)

    def save_mp4(self, output_name, fps=None, scale=None):
        self._load_meta()
        if not (self._descriptor.type == self._db.protobufs.Video and
                self._video_descriptor.codec_type ==
                self._db.protobufs.VideoDescriptor.H264):
            raise ScannerException('Attempted to save a non-h264-compressed '
                                   'column as an mp4. Try compressing the '
                                   'column first by saving the output as '
                                   'an RGB24 frame')
        num_items = len(self._table._descriptor.end_rows)

        paths = ['{}/tables/{:d}/{:d}_{:d}.bin'.format(
            self._db._db_path,
            self._table._descriptor.id, self._descriptor.id, item_id)
                          for item_id in range(num_items)]
        temp_paths = []
        for _ in range(len(paths)):
            fd, p = tempfile.mkstemp()
            os.close(fd)
            temp_paths.append(p)
        # Copy all files locally before calling ffmpeg
        for in_path, temp_path in zip(paths, temp_paths):
            with open(temp_path, 'w') as f:
                f.write(self._storage.read(in_path.encode('ascii')))

        files = '|'.join(temp_paths)

        vid_fps = (fps or
                   (1.0/(self._video_descriptor.time_base_num /
                         float(self._video_descriptor.time_base_denom))))

        args = ''
        if scale:
            args += '-filter:v "scale={:d}x{:d}" '.format(scale[0], scale[1])

        cmd = (
            'ffmpeg -y '
            '-r {fps:f} ' # set the input fps
            '-i "concat:{input_files:s}" ' # concatenate the h264 files
            '-c:v libx264 '
            '-filter:v "setpts=N" ' # h264 does not have pts' in it
            '{extra_args:s}'
            '{output_name:s}.mp4'.format(
                input_files = files,
                fps = vid_fps,
                extra_args = args,
                output_name=output_name))
        rc = Popen(cmd, shell=True).wait()
        if rc != 0:
            raise ScannerException('ffmpeg failed during mp4 export!')
<|endoftext|>


<commit_before>#!/usr/bin/env python
import StringIO
import cv2
import angus
import numpy as np

if __name__ == '__main__':
    ### Web cam index might be different from 0 on your setup.
    stream_index = 0
    cap = cv2.VideoCapture(stream_index)

    if not cap.isOpened():
        print "Cannot open stream of index " + str(stream_index)
        exit(1)

    print "Video stream is of resolution " + str(cap.get(3)) + " x " + str(cap.get(4))

    conn = angus.connect()
    service = conn.services.get_service("age_and_gender_estimation", version=1)
    service.enable_session()

    while(cap.isOpened()):
        ret, frame = cap.read()
        if not frame == None:

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            ret, buff = cv2.imencode(".png", gray)
            buff = StringIO.StringIO(np.array(buff).tostring())

            job = service.process({"image": buff})
            res = job.result

            if res['nb_faces'] > 0:
                for i in range(0,res['nb_faces']):
                    roi = res['faces'][i]['roi']
                    cv2.rectangle(frame, (int(roi[0]), int(roi[1])),
                                         (int(roi[0] + roi[2]), int(roi[1] + roi[3])),
                                         (0,255,0))
                    age = res['faces'][i]['age']
                    gender = res['faces'][i]['gender']
                    cv2.putText(frame, "(age, gender) = (" + '%.1f'%age + ", " + str(gender) + ")",
                                (int(roi[0]), int(roi[1])), cv2.FONT_HERSHEY_SIMPLEX,
                                0.8, (255, 255, 255))

            cv2.imshow('original', frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    service.disable_session()

    cap.release()
    cv2.destroyAllWindows()
<commit_msg>fixing agegender example<commit_after>#!/usr/bin/env python
import StringIO
import cv2
import angus
import numpy as np

if __name__ == '__main__':
    ### Web cam index might be different from 0 on your setup.
    stream_index = 0
    cap = cv2.VideoCapture(stream_index)
    cap.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, 480)

    if not cap.isOpened():
        print "Cannot open stream of index " + str(stream_index)
        exit(1)

    print "Video stream is of resolution " + str(cap.get(3)) + " x " + str(cap.get(4))

    conn = angus.connect()
    service = conn.services.get_service("age_and_gender_estimation", version=1)
    service.enable_session()

    while cap.isOpened():
        ret, frame = cap.read()
        if not frame == None:

            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            ret, buff = cv2.imencode(".png", gray)
            buff = StringIO.StringIO(np.array(buff).tostring())

            job = service.process({"image": buff})
            res = job.result

            if res['nb_faces'] > 0:
                for i in range(0, res['nb_faces']):
                    roi = res['faces'][i]['roi']
                    cv2.rectangle(frame, (int(roi[0]), int(roi[1])),
                                         (int(roi[0] + roi[2]), int(roi[1] + roi[3])),
                                         (0,255,0))
                    age = res['faces'][i]['age']
                    gender = res['faces'][i]['gender']
                    cv2.putText(frame, "(age, gender) = (" + '%.1f'%age + ", " + str(gender) + ")",
                                (int(roi[0]), int(roi[1])), cv2.FONT_HERSHEY_SIMPLEX,
                                0.8, (255, 255, 255))

            cv2.imshow('original', frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    service.disable_session()

    cap.release()
    cv2.destroyAllWindows()
<|endoftext|>


<commit_before>from __future__ import print_function
import atexit
import os
import shutil
import sys
import ctypes
import subprocess

if sys.version_info[0] < 3 or sys.version_info[1] < 5:
  print("\nPlease restart with python3. \n(Taichi supports Python 3.5+)\n")
  print("Current version:", sys.version_info)
  exit(-1)

try:
  import pip
except Exception as e:
  print(e)
  print('  Please install pip3.')
  print('    [Ubuntu] sudo apt-get install python-pip3')
  print('    [Arch Linux] sudo pacman -S python-pip')
  exit(-1)

required_packages = [
    'numpy', ('Pillow', 'PIL'), 'scipy', 'pybind11', 'flask', 'flask_cors',
    ('GitPython', 'git'), 'yapf', 'colorama', 'pyglet', 'PyQt5'
]

def install_package(pkg):
  pip.main(['install', '--user', pkg])

def check_for_packages():
  for pkg in required_packages:
    if isinstance(pkg, tuple):
      import_name = pkg[1]
      pkg = pkg[0]
    else:
      import_name = pkg

    try:
      exec('import {}'.format(import_name))
    except ImportError as e:
      print("Installing package: ", pkg)
      install_package(pkg)

check_for_packages()
from colorama import Fore, Back, Style

def get_repo():
  from git import Repo
  import taichi as tc
  repo = Repo(tc.get_repo_directory())
  return repo

def update(include_projects=False):
  import git
  import taichi as tc

  g = git.cmd.Git(tc.get_repo_directory())
  print(Fore.GREEN + "Updating [taichi]..." + Style.RESET_ALL)
  g.pull('--rebase')
  print(Fore.GREEN + "   ...Done" + Style.RESET_ALL)

  for proj in os.listdir(tc.get_project_directory()):
    if proj in ['examples', 'toys'] or proj.startswith('_') or not os.path.isdir(
        tc.get_project_directory(proj)):
      continue
    print(
        Fore.GREEN + "Updating project [{}]...".format(proj) + Style.RESET_ALL)
    g = git.cmd.Git(os.path.join(tc.get_project_directory(proj)))
    g.pull('--rebase')
    print(Fore.GREEN + "   ...Done" + Style.RESET_ALL)

def format():
  import os
  import sys
  import taichi as tc
  from yapf.yapflib.yapf_api import FormatFile
  repo = get_repo()

  print('* Formatting code', end='')
  for item in repo.index.diff('HEAD'):
    fn = os.path.join(tc.get_repo_directory(), item.a_path)
    print(end='.')
    if fn.endswith('.py'):
      FormatFile(
          fn,
          in_place=True,
          style_config=os.path.join(tc.get_repo_directory(), '.style.yapf'))
    if fn.endswith('.cpp'):
      os.system('clang-format -i -style=file {}'.format(fn))
    repo.git.add(item.a_path)

  print('* Done!')

from taichi.misc.settings import get_output_directory, get_build_directory, get_bin_directory, get_root_directory
from taichi.misc.util import get_os_name, get_unique_task_id

CREATE_SAND_BOX_ON_WINDOWS = True

def build():
  tmp_cwd = os.getcwd()
  bin_dir = get_build_directory()

  try:
    os.mkdir(bin_dir)
  except:
    pass
  os.chdir(bin_dir)

  flags = ' -DPYTHON_EXECUTABLE:FILEPATH="{}"'.format(sys.executable)

  print('Running cmake...')
  if os.environ.get('TC_CI', '') == '1':
    print('  Note: building for CI. SIMD disabled.')
    flags += ' -DTC_DISABLE_SIMD:BOOL=1'
  if get_os_name() == 'win':
    flags += ' -G "Visual Studio 15 Win64"'
  cmake_ret = os.system('cmake .. ' + flags)
  if cmake_ret != 0:
    print('  Error: CMake failed.')
    exit(-1)

  import multiprocessing
  print('Building taichi...')
  num_make_threads = min(20, multiprocessing.cpu_count())
  if get_os_name() == 'win':
    make_ret = os.system("msbuild /p:Configuration=Release /p:Platform=x64 /m taichi.sln")
  else:
    make_ret = os.system('make -j {}'.format(num_make_threads))
  if make_ret != 0:
    print('  Error: Build failed.')
    exit(-1)

  os.chdir(tmp_cwd)

if get_os_name() == 'osx':
  bin_dir = get_bin_directory()
  if not os.path.exists(os.path.join(bin_dir, 'libtaichi_core.dylib')):
    build()
  tmp_cwd = os.getcwd()
  os.chdir(bin_dir)
  shutil.copy('libtaichi_core.dylib', 'taichi_core.so')
  sys.path.append(bin_dir)
  import taichi_core as tc_core
  os.chdir(tmp_cwd)
elif get_os_name() == 'linux':
  bin_dir = get_bin_directory()
  os.environ['LD_LIBRARY_PATH'] = '/usr/lib64/'
  if not os.path.exists(os.path.join(bin_dir, 'libtaichi_core.so')):
    build()
  tmp_cwd = os.getcwd()
  os.chdir(bin_dir)
  sys.path.append(bin_dir)
  # https://stackoverflow.com/questions/3855004/overwriting-library-file-causes-segmentation-fault
  if os.path.exists('taichi_core.so'):
    os.unlink('taichi_core.so')
  shutil.copy('libtaichi_core.so', 'taichi_core.so')
  try:
    import taichi_core as tc_core
  except Exception as e:
    from colorama import Fore, Back, Style
    print(e)
    print()
    print(Fore.RED + "Please make sure you are using python3 "
          "instead of python2." + Style.RESET_ALL)
    print()
    exit(-1)

  os.chdir(tmp_cwd)
elif get_os_name() == 'win':
  bin_dir = get_bin_directory()
  dll_path = os.path.join(bin_dir, 'Release', 'taichi_core.dll')
  if not os.path.exists(dll_path):
    build()

  # The problem here is, on windows, when an dll/pyd is loaded, we can not write to it any more

  old_wd = os.getcwd()
  os.chdir(bin_dir)

  if CREATE_SAND_BOX_ON_WINDOWS:
    # Create a sandbox for separated core lib development and loading
    dir = os.path.join(get_output_directory(), 'tmp', get_unique_task_id())

    os.environ['PATH'] += ';' + (os.path.join(get_root_directory(), 'taichi', 'external', 'lib'))
    os.makedirs(dir)
    shutil.copy(dll_path, os.path.join(dir, 'taichi_core.pyd'))
    sys.path.append(dir)
  else:
    shutil.copy(dll_path, os.path.join(bin_dir, 'taichi_core.pyd'))
    sys.path.append(bin_dir)
  import taichi_core as tc_core

  os.chdir(old_wd)

def get_dll_name(name):
  if get_os_name() == 'linux':
    return 'libtaichi_%s.so' % name
  else:
    assert False

def load_module(name, verbose=True):
  if verbose:
    print('Loading module', name)
  try:
    if '.so' in name:
      ctypes.PyDLL(name)
    else:
      ctypes.PyDLL(
        os.path.join(get_root_directory(), 'taichi', 'build',
                     get_dll_name(name)))
  except Exception as e:
    print(Fore.YELLOW + "Warning: module [{}] loading failed: {}".format(
      name, e) + Style.RESET_ALL)

def at_startup():
  assert os.path.exists(get_root_directory(
  )), 'Please make sure $TAICHI_ROOT_DIR [' + get_root_directory() + '] exists.'
  output_dir = get_output_directory()
  if not os.path.exists(output_dir):
    print('Making output directory')
    os.mkdir(output_dir)

  # Load modules
  f = open(os.path.join(get_root_directory(), 'taichi', 'modules.txt'), 'r')
  modules = f.readline().strip().split(';')
  for module in modules:
    if module != '':
      load_module(module)

  tc_core.set_core_state_python_imported(True)
  f.close()

at_startup()

@atexit.register
def clean_libs():
  pass
<commit_msg>Check package import after installation<commit_after>from __future__ import print_function
import atexit
import os
import shutil
import sys
import ctypes
import subprocess

if sys.version_info[0] < 3 or sys.version_info[1] < 5:
  print("\nPlease restart with python3. \n(Taichi supports Python 3.5+)\n")
  print("Current version:", sys.version_info)
  exit(-1)

try:
  import pip
except Exception as e:
  print(e)
  print('  Please install pip3.')
  print('    [Ubuntu] sudo apt-get install python-pip3')
  print('    [Arch Linux] sudo pacman -S python-pip')
  exit(-1)

required_packages = [
    'numpy', ('Pillow', 'PIL'), 'scipy', 'pybind11', 'flask', 'flask_cors',
    ('GitPython', 'git'), 'yapf', 'colorama', 'pyglet', 'PyQt5'
]

def install_package(pkg):
  pip.main(['install', '--user', pkg])

def check_for_packages():
  for pkg in required_packages:
    if isinstance(pkg, tuple):
      import_name = pkg[1]
      pkg = pkg[0]
    else:
      import_name = pkg

    try:
      exec('import {}'.format(import_name))
    except ImportError as e:
      print("Installing package:", pkg)
      install_package(pkg)
      print('Checking installation of "{}"'.format(import_name))
      exec('import {}'.format(import_name))

check_for_packages()
from colorama import Fore, Back, Style

def get_repo():
  from git import Repo
  import taichi as tc
  repo = Repo(tc.get_repo_directory())
  return repo

def update(include_projects=False):
  import git
  import taichi as tc

  g = git.cmd.Git(tc.get_repo_directory())
  print(Fore.GREEN + "Updating [taichi]..." + Style.RESET_ALL)
  g.pull('--rebase')
  print(Fore.GREEN + "   ...Done" + Style.RESET_ALL)

  for proj in os.listdir(tc.get_project_directory()):
    if proj in ['examples', 'toys'] or proj.startswith('_') or not os.path.isdir(
        tc.get_project_directory(proj)):
      continue
    print(
        Fore.GREEN + "Updating project [{}]...".format(proj) + Style.RESET_ALL)
    g = git.cmd.Git(os.path.join(tc.get_project_directory(proj)))
    g.pull('--rebase')
    print(Fore.GREEN + "   ...Done" + Style.RESET_ALL)

def format():
  import os
  import sys
  import taichi as tc
  from yapf.yapflib.yapf_api import FormatFile
  repo = get_repo()

  print('* Formatting code', end='')
  for item in repo.index.diff('HEAD'):
    fn = os.path.join(tc.get_repo_directory(), item.a_path)
    print(end='.')
    if fn.endswith('.py'):
      FormatFile(
          fn,
          in_place=True,
          style_config=os.path.join(tc.get_repo_directory(), '.style.yapf'))
    if fn.endswith('.cpp'):
      os.system('clang-format -i -style=file {}'.format(fn))
    repo.git.add(item.a_path)

  print('* Done!')

from taichi.misc.settings import get_output_directory, get_build_directory, get_bin_directory, get_root_directory
from taichi.misc.util import get_os_name, get_unique_task_id

CREATE_SAND_BOX_ON_WINDOWS = True

def build():
  tmp_cwd = os.getcwd()
  bin_dir = get_build_directory()

  try:
    os.mkdir(bin_dir)
  except:
    pass
  os.chdir(bin_dir)

  flags = ' -DPYTHON_EXECUTABLE:FILEPATH="{}"'.format(sys.executable)

  print('Running cmake...')
  if os.environ.get('TC_CI', '') == '1':
    print('  Note: building for CI. SIMD disabled.')
    flags += ' -DTC_DISABLE_SIMD:BOOL=1'
  if get_os_name() == 'win':
    flags += ' -G "Visual Studio 15 Win64"'
  cmake_ret = os.system('cmake .. ' + flags)
  if cmake_ret != 0:
    print('  Error: CMake failed.')
    exit(-1)

  import multiprocessing
  print('Building taichi...')
  num_make_threads = min(20, multiprocessing.cpu_count())
  if get_os_name() == 'win':
    make_ret = os.system("msbuild /p:Configuration=Release /p:Platform=x64 /m taichi.sln")
  else:
    make_ret = os.system('make -j {}'.format(num_make_threads))
  if make_ret != 0:
    print('  Error: Build failed.')
    exit(-1)

  os.chdir(tmp_cwd)

if get_os_name() == 'osx':
  bin_dir = get_bin_directory()
  if not os.path.exists(os.path.join(bin_dir, 'libtaichi_core.dylib')):
    build()
  tmp_cwd = os.getcwd()
  os.chdir(bin_dir)
  shutil.copy('libtaichi_core.dylib', 'taichi_core.so')
  sys.path.append(bin_dir)
  import taichi_core as tc_core
  os.chdir(tmp_cwd)
elif get_os_name() == 'linux':
  bin_dir = get_bin_directory()
  os.environ['LD_LIBRARY_PATH'] = '/usr/lib64/'
  if not os.path.exists(os.path.join(bin_dir, 'libtaichi_core.so')):
    build()
  tmp_cwd = os.getcwd()
  os.chdir(bin_dir)
  sys.path.append(bin_dir)
  # https://stackoverflow.com/questions/3855004/overwriting-library-file-causes-segmentation-fault
  if os.path.exists('taichi_core.so'):
    os.unlink('taichi_core.so')
  shutil.copy('libtaichi_core.so', 'taichi_core.so')
  try:
    import taichi_core as tc_core
  except Exception as e:
    from colorama import Fore, Back, Style
    print(e)
    print()
    print(Fore.RED + "Please make sure you are using python3 "
          "instead of python2." + Style.RESET_ALL)
    print()
    exit(-1)

  os.chdir(tmp_cwd)
elif get_os_name() == 'win':
  bin_dir = get_bin_directory()
  dll_path = os.path.join(bin_dir, 'Release', 'taichi_core.dll')
  if not os.path.exists(dll_path):
    build()

  # The problem here is, on windows, when an dll/pyd is loaded, we can not write to it any more

  old_wd = os.getcwd()
  os.chdir(bin_dir)

  if CREATE_SAND_BOX_ON_WINDOWS:
    # Create a sandbox for separated core lib development and loading
    dir = os.path.join(get_output_directory(), 'tmp', get_unique_task_id())

    os.environ['PATH'] += ';' + (os.path.join(get_root_directory(), 'taichi', 'external', 'lib'))
    os.makedirs(dir)
    shutil.copy(dll_path, os.path.join(dir, 'taichi_core.pyd'))
    sys.path.append(dir)
  else:
    shutil.copy(dll_path, os.path.join(bin_dir, 'taichi_core.pyd'))
    sys.path.append(bin_dir)
  import taichi_core as tc_core

  os.chdir(old_wd)

def get_dll_name(name):
  if get_os_name() == 'linux':
    return 'libtaichi_%s.so' % name
  else:
    assert False

def load_module(name, verbose=True):
  if verbose:
    print('Loading module', name)
  try:
    if '.so' in name:
      ctypes.PyDLL(name)
    else:
      ctypes.PyDLL(
        os.path.join(get_root_directory(), 'taichi', 'build',
                     get_dll_name(name)))
  except Exception as e:
    print(Fore.YELLOW + "Warning: module [{}] loading failed: {}".format(
      name, e) + Style.RESET_ALL)

def at_startup():
  assert os.path.exists(get_root_directory(
  )), 'Please make sure $TAICHI_ROOT_DIR [' + get_root_directory() + '] exists.'
  output_dir = get_output_directory()
  if not os.path.exists(output_dir):
    print('Making output directory')
    os.mkdir(output_dir)

  # Load modules
  f = open(os.path.join(get_root_directory(), 'taichi', 'modules.txt'), 'r')
  modules = f.readline().strip().split(';')
  for module in modules:
    if module != '':
      load_module(module)

  tc_core.set_core_state_python_imported(True)
  f.close()

at_startup()

@atexit.register
def clean_libs():
  pass
<|endoftext|>