<commit_before>#!/usr/bin/python
#coding:utf-8

###########################################################
### Utilities for Plexon data collection
### Written by Huangxin
###########################################################

import numpy as np
import logging
logger = logging.getLogger('SpikeRecord.Plexon')
from SpikeRecord import Plexon

def reconstruct_word_in_python(WORD_BITS,bits_num,unstrobed_bits,words_buffer,timestamps_buffer):
    bits_indices = np.array([0]*WORD_BITS)
    oldest_timestamps = np.array([unstrobed_bits[bit][0] for bit in xrange(WORD_BITS)])
    # synonyms
    bits_num = bits_num
    unstrobed_bits = unstrobed_bits
    words_buffer = words_buffer
    timestamps_buffer = timestamps_buffer

    where = np.where
    left_shift = np.left_shift
    timestamps_min = oldest_timestamps.min

    words_count = 0
    indices_sum = 0
    while indices_sum < bits_num:
        timestamp = timestamps_min()
        word_bits = where(oldest_timestamps==timestamp)[0]
        # construct word from bits
        word = left_shift(1,word_bits).sum()
        # increment the indices of previous word bits
        bits_indices[word_bits] += 1
        # increment indices sum so that the loop will run until all bits are processed
        indices_sum += word_bits.size
        # update oldest timestamp of previous word bits
        oldest_timestamps[word_bits] = unstrobed_bits[word_bits,bits_indices[word_bits]]
        # fill word and timestamp in buffers
        words_buffer[words_count] = word
        timestamps_buffer[words_count] = timestamp
        words_count += 1
    return words_count

try:
    import _unstrobed_word
    reconstruct_word = _unstrobed_word.reconstruct_word_32
except ImportError,ValueError:
    reconstruct_word = reconstruct_word_in_python
    logger.info("Cannot import C version of reconstruct_word. Building a C version is highly recommended. We will use Python version this time.")

class PlexUtil(object):
    """
    Utilities for data collection
    """
    def __init__(self):
        self.last_word = None
        self.last_timestamp = None

    def GetSpikesInfo(self,data):
        """
        GetSpikesInfo(data) -> info

        Return spike units collected in this period of time.
        Parameters
        ----------
        data: dict
            {'type', 'channel', 'unit', 'timestamp'} dictionary from the return value of PlexClient.GetTimeStampArray().

        Returns
        info: list of units for every spikes occurring channels
            [(channel, units)]
        """
        sorted_spikes = (data['type'] == Plexon.PL_SingleWFType) & (data['unit'] > 0)
        info = []
        for channel in np.unique(data['channel'][sorted_spikes]):
            channel_units = map(chr, np.unique(data['unit'][sorted_spikes & (data['channel'] == channel)]) + (ord('a')-1))
            info.append((channel, channel_units))
        return info

    def GetSpikeTrains(self,data):
        spike_trains = {}
        sorted_spikes = (data['type'] == Plexon.PL_SingleWFType) & (data['unit'] > 0)
        for channel in np.unique(data['channel'][sorted_spikes]):
            spike_trains[str(channel)] = {}
            for unit in map(chr,np.unique(data['unit'][sorted_spikes & (data['channel'] == channel)]) + (ord('a')-1)):
                spike_trains[str(channel)][unit] = self.GetSpikeTrain(data, channel=channel, unit=unit)
        return spike_trains

    def GetSpikeTrain(self, data, channel, unit):
        """
        GetSpikeTrain(data) -> spike_train

        Return sorted spikes of the specific unit in one channel.
        Parameters
        ----------
        data: dict
            {'type', 'channel', 'unit', 'timestamp'} dictionary from the return value of PlexClient.GetTimeStampArray().
        channel: int
            currently 1-128
        unit: str
            a-z

        Returns
        -------
        spiketrain: array
            timestamp array of the specific unit
        """
        unit_spikes = (data['type'] == Plexon.PL_SingleWFType) & \
                      (data['channel'] == channel) & \
                      (data['unit'] == ord(unit)-ord('a')+1)

        return np.copy(data['timestamp'][unit_spikes])

    def GetEventsNum(self, data):
        return len(data['timestamp'])

    #@profile
    def GetExtEvents(self, data, event, bit=None, online=True):
        """
        GetExtEvents(data) -> extevents

        Return external events.
        Parameters
        ----------
        data: dict
            {'type', 'channel', 'unit', 'timestamp'} dictionary from the return value of PlexClient.GetTimeStampArray().
        event: string
            event types: 'first_strobe_word','second_strobe_word','start','stop','pause','resume','unstrobed_bit'
        bit: int
            currently 1-32, used only in unstrobed_bit event
        Returns
        -------
        extevents: timestamp array of 'first_strobe', 'second_strobe', 'start', 'stop', 'pause', 'resume' events. for strobe_word events
        the array is contained in a dictionary which take the key 'value' as the strobed word and the key 'timestamp' as event stamp.
        """
        ext_event_type = (data['type'] == Plexon.PL_ExtEventType)
        #extevents = data['type'][ext_event_type]
        channel = data['channel'][ext_event_type]
        unit = data['unit'][ext_event_type]
        timestamp = data['timestamp'][ext_event_type]
        # for strobed word event
        if event in ('first_strobe_word','second_strobe_word'):
            strobed_events = (channel == Plexon.PL_StrobedExtChannel)
            strobed_unit = unit[strobed_events]
            strobed_timestamp = timestamp[strobed_events]
            first_strobe  = (strobed_unit & 0x8000) == 0
            second_strobe = (strobed_unit & 0x8000) != 0
            if event == 'first_strobe_word':
                return {'value':strobed_unit[first_strobe] & 0x7FFF , 'timestamp':strobed_timestamp[first_strobe]}
            else:
                return {'value':strobed_unit[second_strobe]         , 'timestamp':strobed_timestamp[second_strobe]}
        # for start event
        if event == 'start':
            return timestamp[channel == Plexon.PL_StartExtChannel]
        # for stop event
        if event == 'stop':
            return timestamp[channel == Plexon.PL_StopExtChannel]
        # for pause event
        if event == 'pause':
            return timestamp[channel == Plexon.PL_Pause]
        # for resume event
        if event == 'resume':
            return timestamp[channel == Plexon.PL_Resume]
        # for unstrobed events
        if event == 'unstrobed_bit':
            return timestamp[channel == bit + 1 ]
        # reconstruct unstrobed word from unstrobed bits
        if event == 'unstrobed_word':
            infinity = float('inf')
            WORD_BITS = 32
            # add an additional infinity in array end so that index of unstrobed_bits will not get out of range
            unstrobed_bits_list = [timestamp[channel == bit+1] for bit in xrange(WORD_BITS)]
            bits_length = [len(unstrobed_bits_list[bit]) for bit in xrange(WORD_BITS)] # actural bits length
            max_length = max(bits_length)
            bits_num = sum(bits_length)
            # make 2d array of timestamp
            unstrobed_bits = np.array([np.append(unstrobed_bits_list[bit], [infinity]*(max_length-bits_length[bit]+1)) \
                                       for bit in xrange(WORD_BITS)],dtype=np.float32)
            # create numpy buffer to hold words and timestamps
            words_buffer = np.empty(bits_num,dtype=np.int32)
            timestamps_buffer = np.empty(bits_num,dtype=np.float32)

            words_count = reconstruct_word(WORD_BITS,bits_num,unstrobed_bits,words_buffer,timestamps_buffer)

            words = words_buffer[:words_count]
            timestamps = timestamps_buffer[:words_count]
            if len(timestamps) and self.last_timestamp == timestamps[0]:
                words[0] += self.last_word
            elif self.last_word is not None:
                words = np.append(self.last_word, words)
                timestamps = np.append(self.last_timestamp, timestamps)
                if len(timestamps)==1:
                    self.last_word = None
                    self.last_timestamp = None
                    return {'value': np.array(words), 'timestamp': np.array(timestamps)}
            if online:      # in offline mode all timestamps are read at once
                if len(timestamps):
                    self.last_word = words[-1]
                    self.last_timestamp = timestamps[-1]
                return {'value': np.array(words[:-1]), 'timestamp': np.array(timestamps[:-1])}
            else:
                return {'value': np.array(words), 'timestamp': np.array(timestamps)}
<commit_msg>use int key for spike channel<commit_after>#!/usr/bin/python
#coding:utf-8

###########################################################
### Utilities for Plexon data collection
### Written by Huangxin
###########################################################

import numpy as np
import logging
logger = logging.getLogger('SpikeRecord.Plexon')
from SpikeRecord import Plexon

def reconstruct_word_in_python(WORD_BITS,bits_num,unstrobed_bits,words_buffer,timestamps_buffer):
    bits_indices = np.array([0]*WORD_BITS)
    oldest_timestamps = np.array([unstrobed_bits[bit][0] for bit in xrange(WORD_BITS)])
    # synonyms
    bits_num = bits_num
    unstrobed_bits = unstrobed_bits
    words_buffer = words_buffer
    timestamps_buffer = timestamps_buffer

    where = np.where
    left_shift = np.left_shift
    timestamps_min = oldest_timestamps.min

    words_count = 0
    indices_sum = 0
    while indices_sum < bits_num:
        timestamp = timestamps_min()
        word_bits = where(oldest_timestamps==timestamp)[0]
        # construct word from bits
        word = left_shift(1,word_bits).sum()
        # increment the indices of previous word bits
        bits_indices[word_bits] += 1
        # increment indices sum so that the loop will run until all bits are processed
        indices_sum += word_bits.size
        # update oldest timestamp of previous word bits
        oldest_timestamps[word_bits] = unstrobed_bits[word_bits,bits_indices[word_bits]]
        # fill word and timestamp in buffers
        words_buffer[words_count] = word
        timestamps_buffer[words_count] = timestamp
        words_count += 1
    return words_count

try:
    import _unstrobed_word
    reconstruct_word = _unstrobed_word.reconstruct_word_32
except ImportError or ValueError:
    reconstruct_word = reconstruct_word_in_python
    logger.info("Cannot import C version of reconstruct_word. Building a C version is highly recommended. We will use Python version this time.")

class PlexUtil(object):
    """
    Utilities for data collection
    """
    def __init__(self):
        self.last_word = None
        self.last_timestamp = None

    def GetSpikesInfo(self,data):
        """
        GetSpikesInfo(data) -> info

        Return spike units collected in this period of time.
        Parameters
        ----------
        data: dict
            {'type', 'channel', 'unit', 'timestamp'} dictionary from the return value of PlexClient.GetTimeStampArray().

        Returns
        info: list of units for every spikes occurring channels
            [(channel, units)]
        """
        sorted_spikes = (data['type'] == Plexon.PL_SingleWFType) & (data['unit'] > 0)
        info = []
        for channel in np.unique(data['channel'][sorted_spikes]):
            channel_units = map(chr, np.unique(data['unit'][sorted_spikes & (data['channel'] == channel)]) + (ord('a')-1))
            info.append((channel, channel_units))
        return info

    def GetSpikeTrains(self,data):
        spike_trains = {}
        sorted_spikes = (data['type'] == Plexon.PL_SingleWFType) & (data['unit'] > 0)
        for channel in np.unique(data['channel'][sorted_spikes]):
            spike_trains[channel] = {}
            for unit in map(chr,np.unique(data['unit'][sorted_spikes & (data['channel'] == channel)]) + (ord('a')-1)):
                spike_trains[channel][unit] = self.GetSpikeTrain(data, channel=channel, unit=unit)
        return spike_trains

    def GetSpikeTrain(self, data, channel, unit):
        """
        GetSpikeTrain(data) -> spike_train

        Return sorted spikes of the specific unit in one channel.
        Parameters
        ----------
        data: dict
            {'type', 'channel', 'unit', 'timestamp'} dictionary from the return value of PlexClient.GetTimeStampArray().
        channel: int
            currently 1-128
        unit: str
            a-z

        Returns
        -------
        spiketrain: array
            timestamp array of the specific unit
        """
        unit_spikes = (data['type'] == Plexon.PL_SingleWFType) & \
                      (data['channel'] == channel) & \
                      (data['unit'] == ord(unit)-ord('a')+1)

        return np.copy(data['timestamp'][unit_spikes])

    def GetEventsNum(self, data):
        return len(data['timestamp'])

    #@profile
    def GetExtEvents(self, data, event, bit=None, online=True):
        """
        GetExtEvents(data) -> extevents

        Return external events.
        Parameters
        ----------
        data: dict
            {'type', 'channel', 'unit', 'timestamp'} dictionary from the return value of PlexClient.GetTimeStampArray().
        event: string
            event types: 'first_strobe_word','second_strobe_word','start','stop','pause','resume','unstrobed_bit'
        bit: int
            currently 1-32, used only in unstrobed_bit event
        Returns
        -------
        extevents: timestamp array of 'first_strobe', 'second_strobe', 'start', 'stop', 'pause', 'resume' events. for strobe_word events
        the array is contained in a dictionary which take the key 'value' as the strobed word and the key 'timestamp' as event stamp.
        """
        ext_event_type = (data['type'] == Plexon.PL_ExtEventType)
        #extevents = data['type'][ext_event_type]
        channel = data['channel'][ext_event_type]
        unit = data['unit'][ext_event_type]
        timestamp = data['timestamp'][ext_event_type]
        # for strobed word event
        if event in ('first_strobe_word','second_strobe_word'):
            strobed_events = (channel == Plexon.PL_StrobedExtChannel)
            strobed_unit = unit[strobed_events]
            strobed_timestamp = timestamp[strobed_events]
            first_strobe  = (strobed_unit & 0x8000) == 0
            second_strobe = (strobed_unit & 0x8000) != 0
            if event == 'first_strobe_word':
                return {'value':strobed_unit[first_strobe] & 0x7FFF , 'timestamp':strobed_timestamp[first_strobe]}
            else:
                return {'value':strobed_unit[second_strobe]         , 'timestamp':strobed_timestamp[second_strobe]}
        # for start event
        if event == 'start':
            return timestamp[channel == Plexon.PL_StartExtChannel]
        # for stop event
        if event == 'stop':
            return timestamp[channel == Plexon.PL_StopExtChannel]
        # for pause event
        if event == 'pause':
            return timestamp[channel == Plexon.PL_Pause]
        # for resume event
        if event == 'resume':
            return timestamp[channel == Plexon.PL_Resume]
        # for unstrobed events
        if event == 'unstrobed_bit':
            return timestamp[channel == bit + 1 ]
        # reconstruct unstrobed word from unstrobed bits
        if event == 'unstrobed_word':
            infinity = float('inf')
            WORD_BITS = 32
            # add an additional infinity in array end so that index of unstrobed_bits will not get out of range
            unstrobed_bits_list = [timestamp[channel == bit+1] for bit in xrange(WORD_BITS)]
            bits_length = [len(unstrobed_bits_list[bit]) for bit in xrange(WORD_BITS)] # actural bits length
            max_length = max(bits_length)
            bits_num = sum(bits_length)
            # make 2d array of timestamp
            unstrobed_bits = np.array([np.append(unstrobed_bits_list[bit], [infinity]*(max_length-bits_length[bit]+1)) \
                                       for bit in xrange(WORD_BITS)],dtype=np.float32)
            # create numpy buffer to hold words and timestamps
            words_buffer = np.empty(bits_num,dtype=np.int32)
            timestamps_buffer = np.empty(bits_num,dtype=np.float32)

            words_count = reconstruct_word(WORD_BITS,bits_num,unstrobed_bits,words_buffer,timestamps_buffer)

            words = words_buffer[:words_count]
            timestamps = timestamps_buffer[:words_count]
            if len(timestamps) and self.last_timestamp == timestamps[0]:
                words[0] += self.last_word
            elif self.last_word is not None:
                words = np.append(self.last_word, words)
                timestamps = np.append(self.last_timestamp, timestamps)
                if len(timestamps)==1:
                    self.last_word = None
                    self.last_timestamp = None
                    return {'value': np.array(words), 'timestamp': np.array(timestamps)}
            if online:      # in offline mode all timestamps are read at once
                if len(timestamps):
                    self.last_word = words[-1]
                    self.last_timestamp = timestamps[-1]
                return {'value': np.array(words[:-1]), 'timestamp': np.array(timestamps[:-1])}
            else:
                return {'value': np.array(words), 'timestamp': np.array(timestamps)}
<|endoftext|>


<commit_before>#!/usr/bin/python -OO

# This file is part of Archivematica.
#
# Copyright 2010-2013 Artefactual Systems Inc. <http://artefactual.com>
#
# Archivematica is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Archivematica is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Archivematica.  If not, see <http://www.gnu.org/licenses/>.

# @package Archivematica
# @subpackage DevCleanup
# @author Mike Cantelon <mike@artefactual.com>

import subprocess

def check_for_string_presence_in_file(string, file):
    fileContents = open(file, 'r').read()
    return string in fileContents

atomProjectConfigFile = '/var/www/ica-atom/config/ProjectConfiguration.class.php'
if not check_for_string_presence_in_file('qtSwordPlugin', atomProjectConfigFile):
    print "The qtSwordPlugin plugin hasn't been enabled in " \
      + atomProjectConfigFile + '.'
    exit(1)

processData = subprocess.check_output(['ps', 'aux'])
if not 'gearman:worker' in processData:
    print "The sword service doesn't seem to be running."
    exit(1)

print 'No issues detected in your AtoM/SWORD configuration!'
<commit_msg>Added a check to see if AtoM setup has been done.<commit_after>#!/usr/bin/python -OO

# This file is part of Archivematica.
#
# Copyright 2010-2013 Artefactual Systems Inc. <http://artefactual.com>
#
# Archivematica is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Archivematica is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Archivematica.  If not, see <http://www.gnu.org/licenses/>.

# @package Archivematica
# @subpackage DevCleanup
# @author Mike Cantelon <mike@artefactual.com>

import subprocess

def check_for_string_presence_in_file(string, file):
    fileContents = open(file, 'r').read()
    return string in fileContents

atomProjectConfigFile = '/var/www/ica-atom/config/ProjectConfiguration.class.php'
if not check_for_string_presence_in_file('qtSwordPlugin', atomProjectConfigFile):
    print "The qtSwordPlugin plugin hasn't been enabled in " \
      + atomProjectConfigFile + '.'
    exit(1)

installCheckData = subprocess.check_output([
 '/usr/bin/curl',
 '-s',
 'http://192.168.1.73/ica-atom/index.php/;sfInstallPlugin/checkSystem'
])
if '
Installation
' in installCheckData:
    print "Please visit http://127.0.0.1/ica-atom to complete installation."
    exit(1)

processData = subprocess.check_output(['ps', 'aux'])
if not 'gearman:worker' in processData:
    print "The sword service doesn't seem to be running."
    print
    print "Enter 'sudo service qubit-sword start' to start the service."
    print "If this results in 'no job found', try restarting."
    exit(1)

print 'No issues detected in your AtoM/SWORD configuration!'
<|endoftext|>


<commit_before>import logging
logger = logging.getLogger(__name__)
import json
import base64
import functools
from time import time
try:
    import tensorflow as tf
except ImportError:
    logger.warning('could not import tensorflow')
import numpy as np
from scipy import optimize

def autospace(x, factor=0.1, num=50, mode='lin'):
    """
    Get a vector with equal spacing whose range is determined by `x`.

    Parameters
    ----------
    x : array_like
        vector to use for range determination
    num : int
        number of sample points
    mode : str
        'lin' for linear spacing, 'log' for log spacing
    factor : float
        factor by which to extend the range
    """
    if mode == 'log':
        x = np.log(x)

    # Evaluate the range
    xmin = np.min(x)
    xmax = np.max(x)
    xrng = xmax - xmin

    # Create the samples
    y = np.linspace(xmin - factor * xrng, xmax + factor * xrng, num)

    # Transform samples
    if mode == 'log':
        return np.exp(y)
    elif mode == 'lin':
        return y
    else:
        raise KeyError(mode)

def haversine(pos1, pos2, degree=True, radius=6371):
    """
    Calculate the great circle distance between two points on the Earth.

    Parameters
    ----------
    pos1 : tuple
        tuple of latitude and longitude coordinates
    pos2 : tuple
        tuple of latitude and longitude coordinates
    degree : bool
        whether the coordinates are in degree or radians
    radius : float
        radius of the Earth
    """
    lat1, lon1 = pos1
    lat2, lon2 = pos2
    # Convert decimal degrees to radians
    if degree:
        lon1, lat1, lon2, lat2 = np.deg2rad([lon1, lat1, lon2, lat2])

    # Haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
    c = 2 * np.arcsin(np.sqrt(a))
    return c * radius

class NumpyEncoder(json.JSONEncoder):
    """
    JSON encoder that supports numpy arrays.

    References
    ----------
    http://stackoverflow.com/a/24375113/1150961
    """
    def default(self, obj):
        """
        If input object is an ndarray it will be converted into a dict holding
        dtype, shape and the data, base64 encoded.

        Parameters
        ----------
        obj : object
            value to encode
        """
        if isinstance(obj, np.ndarray):
            if obj.flags['C_CONTIGUOUS']:
                obj_data = obj.data
            else:
                cont_obj = np.ascontiguousarray(obj)
                assert(cont_obj.flags['C_CONTIGUOUS'])
                obj_data = cont_obj.data
            data_b64 = base64.b64encode(obj_data)
            return dict(__ndarray__=data_b64,
                        dtype=str(obj.dtype),
                        shape=obj.shape)
        # Let the base class default method raise the TypeError
        return json.JSONEncoder(self, obj)

def json_numpy_obj_hook(dictionary):
    """
    Decodes a previously encoded numpy ndarray with proper shape and dtype.

    Parameters
    ----------
    dictionary : dict
        json encoded ndarray

    References
    ----------
    http://stackoverflow.com/a/24375113/1150961
    """
    if isinstance(dictionary, dict) and '__ndarray__' in dictionary:
        data = base64.b64decode(dictionary['__ndarray__'])
        return np.frombuffer(data, dictionary['dtype']).reshape(dictionary['shape'])
    # Return the original data if we cannot decode
    return dictionary

class StringEncoder(json.JSONEncoder):
    """
    JSON encoder that encodes unsupported types as strings.
    """
    def default(self, obj):
        try:
            return super(StringEncoder, self).default(obj)
        except TypeError:
            return repr(obj)

def json_load(f, encoding=None, cls=None, object_hook=None, parse_float=None, parse_int=None, parse_constant=None,
              object_pairs_hook=None, **kwargs):
    """
    Convenience function to load JSON.

    Parameters
    ----------
    f : str or file
        file pointer or file name to read from

    Notes
    -----
    See json.load for details of the other parameters.
    """
    # Open a file if a string is given
    if isinstance(f, str):
        fp = open(f)
    elif hasattr(f, 'read'):
        fp = f
    else:
        raise ValueError('`f` must be a file pointer or file name')

    # Load the data
    result = json.load(fp, encoding, cls, object_hook or json_numpy_obj_hook, parse_float, parse_int, parse_constant,
                       object_pairs_hook, **kwargs)

    # Close the file if we just opened it
    if isinstance(f, str):
        fp.close()

    return result

def json_loads(s, encoding=None, cls=None, object_hook=None, parse_float=None, parse_int=None, parse_constant=None,
               object_pairs_hook=None, **kwargs):
    """
    Convenience function to load JSON from a string.

    Notes
    -----
    See json.loads for details of the parameters.
    """
    return json.loads(s, encoding, cls, object_hook or json_numpy_obj_hook, parse_float, parse_int, parse_constant,
                      object_pairs_hook, **kwargs)

def json_dump(obj, f, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, cls=None, indent=None,
              separators=None, encoding='utf-8', default=None, sort_keys=False, **kwargs):
    """
    Convenience function to dump JSON.

    Parameters
    ----------
    f : str or file
        file pointer or file name to write to

    Notes
    -----
    See json.dump for details of the other parameters.
    """
    if isinstance(f, str):
        fp = open(f, 'w')
    elif hasattr(f, 'write'):
        fp = f
    else:
        raise ValueError("`f` must be a file pointer or file name but got '{}'".format(f))

    # Dump the data to disk
    json.dump(obj, fp, skipkeys, ensure_ascii, check_circular, allow_nan, cls or NumpyEncoder, indent,
              separators, encoding, default, sort_keys, **kwargs)

    # Close the file if we just opened it
    if isinstance(f, str):
        fp.close()

def json_dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, cls=None, indent=None,
              separators=None, encoding='utf-8', default=None, sort_keys=False, **kwargs):
    """
    Convenience function to dump JSON to a string.

    Notes
    -----
    See json.dumps for details of the parameters.
    """
    return json.dumps(obj, skipkeys, ensure_ascii, check_circular, allow_nan, cls or NumpyEncoder, indent,
                       separators, encoding, default, sort_keys, **kwargs)

def read(filename):
    """
    Read all text from the specified file.

    Parameters
    ----------
    filename : str
        path to read from
    """
    with open(filename) as fp:
        return fp.read()

def recursive_update(self, other):
    """
    Update dictionary recursively.

    Parameters
    ----------
    self : dict
        dictionary to update
    other : dict
        dictionary to copy values from
    """
    for key, value in other.iteritems():
        # Simply copy the values if the value is not in the original dictionary or if we aren't dealing with dicts
        if key not in self or not isinstance(value, dict) or not isinstance(self[key], dict):
            self[key] = value
        # Perform a recursive update otherwise
        else:
            recursive_update(self[key], value)

    return self

class Timer(object):
    """
    Timer object that can be used as a `with` statement.
    """
    def __init__(self, message="duration : {0}", logger='print', level='info'):
        self.logger = logger
        self.level = level.lower()
        self.start = self.end = self.duration = None
        self.message = message

    def __enter__(self):
        self.start = time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time()
        self.duration = self.end - self.start

        # Format a message
        message = self.message.format(self.duration)

        if self.logger == 'print':
            print(message)
        elif self.logger:
            logger = logging.getLogger(self.logger)
            getattr(logger, self.level)(message)

class timeit:
    """
    Time the execution of a function and log the results.

    Parameters
    ----------
    logger : str
        name of the logger or 'print' to print to `stdout`
    level : str
        logging level for the reporting
    """
    def __init__(self, logger='print', level='info'):
        self.logger = logger
        self.level = level.lower()

    def __call__(self, func):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            # Call the function wrapped in a timer
            with Timer(func.func_name + " : {0}", self.logger, self.level):
                return func(*args, **kwargs)

        return inner

def negate(fun):
    """
    Negate a function.

    Parameters
    ----------
    fun : callable
        the function to negate
    """
    if fun:
        @functools.wraps(fun)
        def inner(*args, **kwargs):
            return -fun(*args, **kwargs)

        return inner
    else:
        return None

@functools.wraps(optimize.minimize)
def maximize(fun, x0, args=(), method=None, jac=None, hess=None, hessp=None, bounds=None, constraints=(), tol=None,
             callback=None, options=None):
    return optimize.minimize(negate(fun), x0, args, method, negate(jac), negate(hess), negate(hessp), bounds,
                             constraints, tol, callback, options)

def log_gaussian(x, mean=0.0, covariance=1.0):
    """
    Evaluate the log of the (unnormalised) Gaussian distribution.

    Parameters
    ----------
    x : float
        point at which to evaluate the distribution
    mean : float
        mean of the distribution
    covariance : float
        covariance of the distribution

    Returns
    -------
    value : float
        value of the log of the pdf
    jac : float
        derivative of the log of the pdf
    """
    # Convert a one-dimensional and two-dimensional arrays
    mean = np.atleast_1d(mean)
    covariance = np.atleast_1d(covariance)
    if covariance.ndim == 1:
        covariance = np.diag(covariance)

    # Compute the values
    precision = np.linalg.inv(covariance)
    residual = x - mean
    value = -0.5 * np.dot(np.dot(residual, precision), residual)
    jac = np.dot(residual, precision)
    return value, jac

def acorr(x, ax=None):
    """
    Compute the autocorrelation.

    References
    ----------
    http://stackoverflow.com/a/27558293/1150961
    """
    x = np.asarray(x)
    x = x - np.mean(x)

    autocorr = np.correlate(x, x, mode='full')
    autocorr = autocorr[x.size:]
    autocorr /= autocorr.max()

    return autocorr

def iterable(x):
    """
    Check whether the input is iterable but not a string.

    Parameters
    ----------
    x
        input to check

    Returns
    -------
    iterable : bool
        whether the input is iterable but not a string
    """
    return hasattr(x, '__iter__') and not isinstance(x, (str, bytes, tf.Variable, tf.Tensor))
<commit_msg>Add minmax function.<commit_after>import logging
logger = logging.getLogger(__name__)
import json
import base64
import functools
from time import time
try:
    import tensorflow as tf
except ImportError:
    logger.warning('could not import tensorflow')
import numpy as np
from scipy import optimize

def autospace(x, factor=0.1, num=50, mode='lin'):
    """
    Get a vector with equal spacing whose range is determined by `x`.

    Parameters
    ----------
    x : array_like
        vector to use for range determination
    num : int
        number of sample points
    mode : str
        'lin' for linear spacing, 'log' for log spacing
    factor : float
        factor by which to extend the range
    """
    if mode == 'log':
        x = np.log(x)

    # Evaluate the range
    xmin = np.min(x)
    xmax = np.max(x)
    xrng = xmax - xmin

    # Create the samples
    y = np.linspace(xmin - factor * xrng, xmax + factor * xrng, num)

    # Transform samples
    if mode == 'log':
        return np.exp(y)
    elif mode == 'lin':
        return y
    else:
        raise KeyError(mode)

def haversine(pos1, pos2, degree=True, radius=6371):
    """
    Calculate the great circle distance between two points on the Earth.

    Parameters
    ----------
    pos1 : tuple
        tuple of latitude and longitude coordinates
    pos2 : tuple
        tuple of latitude and longitude coordinates
    degree : bool
        whether the coordinates are in degree or radians
    radius : float
        radius of the Earth
    """
    lat1, lon1 = pos1
    lat2, lon2 = pos2
    # Convert decimal degrees to radians
    if degree:
        lon1, lat1, lon2, lat2 = np.deg2rad([lon1, lat1, lon2, lat2])

    # Haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
    c = 2 * np.arcsin(np.sqrt(a))
    return c * radius

class NumpyEncoder(json.JSONEncoder):
    """
    JSON encoder that supports numpy arrays.

    References
    ----------
    http://stackoverflow.com/a/24375113/1150961
    """
    def default(self, obj):
        """
        If input object is an ndarray it will be converted into a dict holding
        dtype, shape and the data, base64 encoded.

        Parameters
        ----------
        obj : object
            value to encode
        """
        if isinstance(obj, np.ndarray):
            if obj.flags['C_CONTIGUOUS']:
                obj_data = obj.data
            else:
                cont_obj = np.ascontiguousarray(obj)
                assert(cont_obj.flags['C_CONTIGUOUS'])
                obj_data = cont_obj.data
            data_b64 = base64.b64encode(obj_data)
            return dict(__ndarray__=data_b64,
                        dtype=str(obj.dtype),
                        shape=obj.shape)
        # Let the base class default method raise the TypeError
        return json.JSONEncoder(self, obj)

def json_numpy_obj_hook(dictionary):
    """
    Decodes a previously encoded numpy ndarray with proper shape and dtype.

    Parameters
    ----------
    dictionary : dict
        json encoded ndarray

    References
    ----------
    http://stackoverflow.com/a/24375113/1150961
    """
    if isinstance(dictionary, dict) and '__ndarray__' in dictionary:
        data = base64.b64decode(dictionary['__ndarray__'])
        return np.frombuffer(data, dictionary['dtype']).reshape(dictionary['shape'])
    # Return the original data if we cannot decode
    return dictionary

class StringEncoder(json.JSONEncoder):
    """
    JSON encoder that encodes unsupported types as strings.
    """
    def default(self, obj):
        try:
            return super(StringEncoder, self).default(obj)
        except TypeError:
            return repr(obj)

def json_load(f, encoding=None, cls=None, object_hook=None, parse_float=None, parse_int=None, parse_constant=None,
              object_pairs_hook=None, **kwargs):
    """
    Convenience function to load JSON.

    Parameters
    ----------
    f : str or file
        file pointer or file name to read from

    Notes
    -----
    See json.load for details of the other parameters.
    """
    # Open a file if a string is given
    if isinstance(f, str):
        fp = open(f)
    elif hasattr(f, 'read'):
        fp = f
    else:
        raise ValueError('`f` must be a file pointer or file name')

    # Load the data
    result = json.load(fp, encoding, cls, object_hook or json_numpy_obj_hook, parse_float, parse_int, parse_constant,
                       object_pairs_hook, **kwargs)

    # Close the file if we just opened it
    if isinstance(f, str):
        fp.close()

    return result

def json_loads(s, encoding=None, cls=None, object_hook=None, parse_float=None, parse_int=None, parse_constant=None,
               object_pairs_hook=None, **kwargs):
    """
    Convenience function to load JSON from a string.

    Notes
    -----
    See json.loads for details of the parameters.
    """
    return json.loads(s, encoding, cls, object_hook or json_numpy_obj_hook, parse_float, parse_int, parse_constant,
                      object_pairs_hook, **kwargs)

def json_dump(obj, f, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, cls=None, indent=None,
              separators=None, encoding='utf-8', default=None, sort_keys=False, **kwargs):
    """
    Convenience function to dump JSON.

    Parameters
    ----------
    f : str or file
        file pointer or file name to write to

    Notes
    -----
    See json.dump for details of the other parameters.
    """
    if isinstance(f, str):
        fp = open(f, 'w')
    elif hasattr(f, 'write'):
        fp = f
    else:
        raise ValueError("`f` must be a file pointer or file name but got '{}'".format(f))

    # Dump the data to disk
    json.dump(obj, fp, skipkeys, ensure_ascii, check_circular, allow_nan, cls or NumpyEncoder, indent,
              separators, encoding, default, sort_keys, **kwargs)

    # Close the file if we just opened it
    if isinstance(f, str):
        fp.close()

def json_dumps(obj, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, cls=None, indent=None,
              separators=None, encoding='utf-8', default=None, sort_keys=False, **kwargs):
    """
    Convenience function to dump JSON to a string.

    Notes
    -----
    See json.dumps for details of the parameters.
    """
    return json.dumps(obj, skipkeys, ensure_ascii, check_circular, allow_nan, cls or NumpyEncoder, indent,
                       separators, encoding, default, sort_keys, **kwargs)

def read(filename):
    """
    Read all text from the specified file.

    Parameters
    ----------
    filename : str
        path to read from
    """
    with open(filename) as fp:
        return fp.read()

def recursive_update(self, other):
    """
    Update dictionary recursively.

    Parameters
    ----------
    self : dict
        dictionary to update
    other : dict
        dictionary to copy values from
    """
    for key, value in other.iteritems():
        # Simply copy the values if the value is not in the original dictionary or if we aren't dealing with dicts
        if key not in self or not isinstance(value, dict) or not isinstance(self[key], dict):
            self[key] = value
        # Perform a recursive update otherwise
        else:
            recursive_update(self[key], value)

    return self

class Timer(object):
    """
    Timer object that can be used as a `with` statement.
    """
    def __init__(self, message="duration : {0}", logger='print', level='info'):
        self.logger = logger
        self.level = level.lower()
        self.start = self.end = self.duration = None
        self.message = message

    def __enter__(self):
        self.start = time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time()
        self.duration = self.end - self.start

        # Format a message
        message = self.message.format(self.duration)

        if self.logger == 'print':
            print(message)
        elif self.logger:
            logger = logging.getLogger(self.logger)
            getattr(logger, self.level)(message)

class timeit:
    """
    Time the execution of a function and log the results.

    Parameters
    ----------
    logger : str
        name of the logger or 'print' to print to `stdout`
    level : str
        logging level for the reporting
    """
    def __init__(self, logger='print', level='info'):
        self.logger = logger
        self.level = level.lower()

    def __call__(self, func):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            # Call the function wrapped in a timer
            with Timer(func.func_name + " : {0}", self.logger, self.level):
                return func(*args, **kwargs)

        return inner

def negate(fun):
    """
    Negate a function.

    Parameters
    ----------
    fun : callable
        the function to negate
    """
    if fun:
        @functools.wraps(fun)
        def inner(*args, **kwargs):
            return -fun(*args, **kwargs)

        return inner
    else:
        return None

@functools.wraps(optimize.minimize)
def maximize(fun, x0, args=(), method=None, jac=None, hess=None, hessp=None, bounds=None, constraints=(), tol=None,
             callback=None, options=None):
    return optimize.minimize(negate(fun), x0, args, method, negate(jac), negate(hess), negate(hessp), bounds,
                             constraints, tol, callback, options)

def log_gaussian(x, mean=0.0, covariance=1.0):
    """
    Evaluate the log of the (unnormalised) Gaussian distribution.

    Parameters
    ----------
    x : float
        point at which to evaluate the distribution
    mean : float
        mean of the distribution
    covariance : float
        covariance of the distribution

    Returns
    -------
    value : float
        value of the log of the pdf
    jac : float
        derivative of the log of the pdf
    """
    # Convert a one-dimensional and two-dimensional arrays
    mean = np.atleast_1d(mean)
    covariance = np.atleast_1d(covariance)
    if covariance.ndim == 1:
        covariance = np.diag(covariance)

    # Compute the values
    precision = np.linalg.inv(covariance)
    residual = x - mean
    value = -0.5 * np.dot(np.dot(residual, precision), residual)
    jac = np.dot(residual, precision)
    return value, jac

def acorr(x, ax=None):
    """
    Compute the autocorrelation.

    References
    ----------
    http://stackoverflow.com/a/27558293/1150961
    """
    x = np.asarray(x)
    x = x - np.mean(x)

    autocorr = np.correlate(x, x, mode='full')
    autocorr = autocorr[x.size:]
    autocorr /= autocorr.max()

    return autocorr

def iterable(x):
    """
    Check whether the input is iterable but not a string.

    Parameters
    ----------
    x
        input to check

    Returns
    -------
    iterable : bool
        whether the input is iterable but not a string
    """
    return hasattr(x, '__iter__') and not isinstance(x, (str, bytes, tf.Variable, tf.Tensor))

def minmax(x, axis=None):
    """
    Return the minimum and maximum of an array.
    """
    return np.min(x, axis=axis), np.max(x, axis=axis)
<|endoftext|>


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

import sys
import argparse
import configuration

from datetime import datetime
from gemini import GeminiQuery
from collections import defaultdict
from cassandra.cqlengine import connection

from variantstore import Variant
from cyvcf2 import VCF

def _var_is_rare(variant_data):
    """Determine if the MAF of the variant is < 1% in all populations"""

    if variant_data['in_esp'] != 0 or variant_data['in_1kg'] != 0 or variant_data['in_exac'] != 0:
        if variant_data['max_aaf_all'] > 0.01:
            return False
        else:
            return True
    else:
        return True

def _var_is_in_cosmic(variant_data):
    """Determine if the variant has a COSMIC identifier"""

    if variant_data['cosmic_ids'] is not None:
        return True
    else:
        return False

def _var_is_in_clinvar(variant_data):
    """Determine if there is ClinVar data for the variant"""

    if variant_data['clinvar_sig'] is not None:
        return True
    else:
        return False

def _var_is_pathogenic(variant_data):
    """Determine if a variant in clinvar is potentially pathogenic"""

    if variant_data['clinvar_sig'] is not None:
        if "pathogenic" in variant_data['clinvar_sig']:
            return True
        else:
            return False
    else:
        return False

def _var_is_protein_effecting(variant_data):
    if variant_data['impact_severity'] != "LOW":
        return True
    else:
        return False

def _var_in_gene(variant_data, genes):
    if variant_data['gene'] in genes:
        return True
    else:
        return False

def parse_vcf(vcf_file, caller, caller_vcf_records):
    sys.stdout.write("Reading {}\n".format(vcf_file))
    vcf = VCF(vcf_file)
    for record in vcf:
        if len(record.ALT) > 1:
            sys.stderr.write("ERROR: More than one alternative allele detected in file "
                             "{}\n Record: {}\n".format(vcf_file, record))
            sys.exit()
        key = (unicode("chr{}".format(record.CHROM)), int(record.start), int(record.end), unicode(record.REF),
               unicode(record.ALT[0]))
        caller_vcf_records[caller][key] = record

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--samples_file', help="Input configuration file for samples")
    parser.add_argument('-c', '--configuration', help="Configuration file for various settings")
    args = parser.parse_args()

    sys.stdout.write("Parsing configuration data\n")
    config = configuration.configure_runtime(args.configuration)

    sys.stdout.write("Parsing sample data\n")
    samples = configuration.configure_samples(args.samples_file, config)

    connection.setup(['127.0.0.1'], "variantstore")

    for sample in samples:
        query = "SELECT chrom, start, end, ref, alt, vcf_id, rs_ids, cosmic_ids, filter, qual, qual_depth, depth, " \
                "type, sub_type, " \
                "gene, transcript, exon, codon_change, aa_change, biotype, impact, impact_so, impact_severity, " \
                "aa_length, is_lof, is_conserved, pfam_domain, in_omim, clinvar_sig, clinvar_disease_name, " \
                "is_exonic, is_coding, is_splicing, " \
                "clinvar_origin, clinvar_causal_allele, clinvar_dbsource, clinvar_dbsource_id, " \
                "clinvar_on_diag_assay, rmsk, in_segdup, strand_bias, rms_map_qual, in_hom_run, num_mapq_zero, " \
                "num_reads_w_dels, grc, gms_illumina, in_cse, num_alleles, allele_count, haplotype_score, " \
                "is_somatic, somatic_score, aaf_esp_ea, aaf_esp_aa, aaf_esp_all, aaf_1kg_amr, " \
                "aaf_1kg_eas, aaf_1kg_sas, aaf_1kg_afr, aaf_1kg_eur, aaf_1kg_all, aaf_exac_all, aaf_adj_exac_all, " \
                "aaf_adj_exac_afr, aaf_adj_exac_amr, aaf_adj_exac_eas, aaf_adj_exac_fin, aaf_adj_exac_nfe, " \
                "aaf_adj_exac_oth, aaf_adj_exac_sas, max_aaf_all, in_esp, in_1kg, in_exac, info," \
                "(gts).(*), (gt_depths).(*), (gt_ref_depths).(*), (gt_alt_depths).(*) FROM variants"

        sys.stdout.write("Running GEMINI query\n")
        gq = GeminiQuery(samples[sample]['db'])
        gq.run(query)

        caller_vcf_records = defaultdict(lambda: dict())

        parse_vcf(samples[sample]['mutect'], "mutect", caller_vcf_records)
        parse_vcf(samples[sample]['vardict'], "vardict", caller_vcf_records)
        parse_vcf(samples[sample]['freebayes'], "freebayes", caller_vcf_records)
        parse_vcf(samples[sample]['scalpel'], "scalpel", caller_vcf_records)

        # Filter out variants with minor allele frequencies above the threshold but
        # retain any that are above the threshold but in COSMIC or in ClinVar and not listed as benign.
        for variant_data in gq:
            cassandra_variant = Variant(chr=variant_data['chrom'], start=variant_data['start'], end=variant_data['end'],
                                        ref=variant_data['ref'], alt=variant_data['alt'],
                                        sample=samples[sample]['sample_name'], extraction=samples[sample]['extraction'],
                                        library_name=sample, panel_name=samples[sample]['panel'],
                                        target_pool=samples[sample]['target_pool'],
                                        reference_genome=config['genome_version'], date_annotated=datetime.now(),
                                        type=variant_data['type'],
                                        subtype=variant_data['sub_type'], rs_id=variant_data['vcf_id'],
                                        gene=variant_data['gene'], max_aaf=variant_data['max_aaf_all'],
                                        transcript=variant_data['transcript'], exon=variant_data['exon'],
                                        codon_change=variant_data['codon_change'],
                                        aa_change=variant_data['aa_change'], biotype=variant_data['biotype'],
                                        impact=variant_data['impact'], impact_so=variant_data['impact_so'])

            cassandra_variant['in_clinvar'] = _var_is_in_clinvar(variant_data)
            cassandra_variant['in_cosmic'] = _var_is_in_cosmic(variant_data)
            cassandra_variant['is_pathogenic'] = _var_is_pathogenic(variant_data)

            if variant_data['is_lof']:
                cassandra_variant['is_lof'] = True
            else:
                cassandra_variant['is_lof'] = False

            if variant_data['is_coding']:
                cassandra_variant['is_coding'] = True
            else:
                cassandra_variant['is_coding'] = False

            if variant_data['is_splicing']:
                cassandra_variant['is_splicing'] = True
            else:
                cassandra_variant['is_splicing'] = False

            if variant_data['rs_ids'] is not None:
                cassandra_variant['rs_ids'] = variant_data['rs_ids'].split(',')

            if variant_data['cosmic_ids'] is not None:
                cassandra_variant['cosmic_ids'] = variant_data['cosmic_ids'].split(',')

            if variant_data['info']['CALLERS'] is not None:
                cassandra_variant['callers'] = variant_data['info']['CALLERS'].split(',')

            population_freqs = {'esp_ea': variant_data['aaf_esp_ea'],
                                'esp_aa': variant_data['aaf_esp_aa'],
                                'esp_all': variant_data['aaf_esp_all'],
                                '1kg_amr': variant_data['aaf_1kg_amr'],
                                '1kg_eas': variant_data['aaf_1kg_eas'],
                                '1kg_sas': variant_data['aaf_1kg_sas'],
                                '1kg_afr': variant_data['aaf_1kg_afr'],
                                '1kg_eur': variant_data['aaf_1kg_eur'],
                                '1kg_all': variant_data['aaf_1kg_all'],
                                'exac_all': variant_data['aaf_exac_all'],
                                'adj_exac_all': variant_data['aaf_adj_exac_all'],
                                'adj_exac_afr': variant_data['aaf_adj_exac_afr'],
                                'adj_exac_amr': variant_data['aaf_adj_exac_amr'],
                                'adj_exac_eas': variant_data['aaf_adj_exac_eas'],
                                'adj_exac_fin': variant_data['aaf_adj_exac_fin'],
                                'adj_exac_nfe': variant_data['aaf_adj_exac_nfe'],
                                'adj_exac_oth': variant_data['aaf_adj_exac_oth'],
                                'adj_exac_sas': variant_data['aaf_adj_exac_sas']}

            cassandra_variant['population_freqs'] = population_freqs

            key = (variant_data['chrom'], variant_data['start'], variant_data['end'], variant_data['ref'],
                   variant_data['alt'])

            if 'mutect' in variant_data['info']['CALLERS']:
                record = caller_vcf_records['mutect'][key]
                # filters = record.FILTER.split(',')

                mutect_info = {'GTF_DP': str(record.gt_depths[0]),
                               'GTF_AD': str(record.gt_alt_depths[0])}

                cassandra_variant['mutect'] = mutect_info

            if 'vardict' in variant_data['info']['CALLERS']:
                record = caller_vcf_records['vardict'][key]
                # filters = record.FILTER.split(',')

                vardict_info = {'DP': str(record.INFO.get('DP')),
                                'VD': str(record.INFO.get('VD')),
                                'AF': str(record.INFO.get('AF')),
                                'BIAS': str(record.INFO.get('BIAS')),
                                'REFBIAS': str(record.INFO.get('REFBIAS')),
                                'VARBIAS': str(record.INFO.get('VARBIAS')),
                                'QUAL': str(record.INFO.get('QUAL')),
                                'QSTD': str(record.INFO.get('QSTD')),
                                'SBF': str(record.INFO.get('SBF')),
                                'ODDRATIO': str(record.INFO.get('ODDRATIO')),
                                'MQ': str(record.INFO.get('MQ')),
                                'SN': str(record.INFO.get('SN')),
                                'HIAF': str(record.INFO.get('HIAF')),
                                'ADJAF': str(record.INFO.get('ADJAF')),
                                'MSI': str(record.INFO.get('MSI')),
                                'MSILEN': str(record.INFO.get('MSILEN')),
                                'SHIFT3': str(record.INFO.get('SHIFT3')),
                                'NM': str(record.INFO.get('NM')),
                                'GDAMP': str(record.INFO.get('GDAMP')),
                                'LSEQ': str(record.INFO.get('LSEQ')),
                                'RSEQ': str(record.INFO.get('RSEQ')),
                                'TLAMP': str(record.INFO.get('TLAMP')),
                                'NCAMP': str(record.INFO.get('NCAMP')),
                                'AMPFLAG': str(record.INFO.get('AMPFLAG')),
                                'HICNT': str(record.INFO.get('HICNT')),
                                'HICOV': str(record.INFO.get('HICOV')),
                                'GTF_DP': str(record.gt_depths[0]),
                                'GTF_AD': str(record.gt_alt_depths[0])}

                cassandra_variant['vardict'] = vardict_info

            if 'freebayes' in variant_data['info']['CALLERS']:
                record = caller_vcf_records['freebayes'][key]
                # filters = record.FILTER.split(',')

                freebayes_info = {'DP': str(record.INFO.get('DP')),
                                  'AF': str(record.INFO.get('AF')),
                                  'AC': str(record.INFO.get('AC')),
                                  'RO': str(record.INFO.get('RO')),
                                  'AO': str(record.INFO.get('AO')),
                                  'PRO': str(record.INFO.get('PRO')),
                                  'PAO': str(record.INFO.get('PAO')),
                                  'QR': str(record.INFO.get('QR')),
                                  'QA': str(record.INFO.get('QA')),
                                  'PQR': str(record.INFO.get('PQR')),
                                  'PQA': str(record.INFO.get('PQA')),
                                  'SRF': str(record.INFO.get('SRF')),
                                  'SRR': str(record.INFO.get('SRR')),
                                  'SAF': str(record.INFO.get('SAF')),
                                  'SAR': str(record.INFO.get('SAR')),
                                  'SRP': str(record.INFO.get('SRP')),
                                  'SAP': str(record.INFO.get('SAP')),
                                  'AB': str(record.INFO.get('AB')),
                                  'ABP': str(record.INFO.get('ABP')),
                                  'RUN': str(record.INFO.get('RUN')),
                                  'RPP': str(record.INFO.get('RPP')),
                                  'RPPR': str(record.INFO.get('RPPR')),
                                  'RPL': str(record.INFO.get('RPL')),
                                  'RPR': str(record.INFO.get('RPR')),
                                  'EPP': str(record.INFO.get('EPP')),
                                  'EPPR': str(record.INFO.get('EPPR')),
                                  'DRPA': str(record.INFO.get('DRPA')),
                                  'ODDS': str(record.INFO.get('ODDS')),
                                  'GTI': str(record.INFO.get('GTI')),
                                  'TYPE': str(record.INFO.get('TYPE')),
                                  'CIGAR': str(record.INFO.get('CIGAR')),
                                  'NUMALT': str(record.INFO.get('NUMALT')),
                                  'MEANALT': str(record.INFO.get('MEANALT')),
                                  'LEN': str(record.INFO.get('LEN')),
                                  'MQM': str(record.INFO.get('MQM')),
                                  'MQMR': str(record.INFO.get('MQMR')),
                                  'PAIRED': str(record.INFO.get('PAIRED')),
                                  'PAIREDR': str(record.INFO.get('PAIREDR')),
                                  'GTF_DP': str(record.gt_depths[0])}

                cassandra_variant['freebayes'] = freebayes_info

            if 'scalpel' in variant_data['info']['CALLERS']:
                record = caller_vcf_records['scalpel'][key]
                # filters = record.FILTER.split(',')

                scalpel_info = {'AVGCOV': str(record.INFO.get('AVGCOV')),
                                'MINCOV': str(record.INFO.get('MINCOV')),
                                'ALTCOV': str(record.INFO.get('ALTCOV')),
                                'COVRATIO': str(record.INFO.get('COVRATIO')),
                                'ZYG': str(record.INFO.get('ZYG')),
                                'CHI2': str(record.INFO.get('CHI2')),
                                'FISHERPHREDSCORE': str(record.INFO.get('FISHERPHREDSCORE')),
                                'INH': str(record.INFO.get('INH')),
                                'BESTSTATE': str(record.INFO.get('BESTSTATE')),
                                'COVSTATE': str(record.INFO.get('COVSTATE')),
                                'SOMATIC': str(record.INFO.get('SOMATIC')),
                                'DENOVO': str(record.INFO.get('DENOVO')),
                                'GTF_DP': str(record.gt_depths[0]),
                                'GTF_AD': str(record.gt_alt_depths[0])}

                cassandra_variant['scalpel'] = scalpel_info

            cassandra_variant.save()
<commit_msg>Import from ddb core module for configuration and GEMINI interface<commit_after>#!/usr/bin/env python

import sys
import argparse

from datetime import datetime
from gemini import GeminiQuery
from collections import defaultdict
from cassandra.cqlengine import connection

from variantstore import Variant
from cyvcf2 import VCF
from ddb import gemini_interface
from ddb import configuration

def parse_vcf(vcf_file, caller, caller_vcf_records):
    sys.stdout.write("Reading {}\n".format(vcf_file))
    vcf = VCF(vcf_file)
    for record in vcf:
        if len(record.ALT) > 1:
            sys.stderr.write("ERROR: More than one alternative allele detected in file "
                             "{}\n Record: {}\n".format(vcf_file, record))
            sys.exit()
        key = (unicode("chr{}".format(record.CHROM)), int(record.start), int(record.end), unicode(record.REF),
               unicode(record.ALT[0]))
        caller_vcf_records[caller][key] = record

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--samples_file', help="Input configuration file for samples")
    parser.add_argument('-c', '--configuration', help="Configuration file for various settings")
    args = parser.parse_args()

    sys.stdout.write("Parsing configuration data\n")
    config = configuration.configure_runtime(args.configuration)

    sys.stdout.write("Parsing sample data\n")
    samples = configuration.configure_samples(args.samples_file, config)

    connection.setup(['127.0.0.1'], "variantstore")

    for sample in samples:
        query = "SELECT chrom, start, end, ref, alt, vcf_id, rs_ids, cosmic_ids, filter, qual, qual_depth, depth, " \
                "type, sub_type, " \
                "gene, transcript, exon, codon_change, aa_change, biotype, impact, impact_so, impact_severity, " \
                "aa_length, is_lof, is_conserved, pfam_domain, in_omim, clinvar_sig, clinvar_disease_name, " \
                "is_exonic, is_coding, is_splicing, " \
                "clinvar_origin, clinvar_causal_allele, clinvar_dbsource, clinvar_dbsource_id, " \
                "clinvar_on_diag_assay, rmsk, in_segdup, strand_bias, rms_map_qual, in_hom_run, num_mapq_zero, " \
                "num_reads_w_dels, grc, gms_illumina, in_cse, num_alleles, allele_count, haplotype_score, " \
                "is_somatic, somatic_score, aaf_esp_ea, aaf_esp_aa, aaf_esp_all, aaf_1kg_amr, " \
                "aaf_1kg_eas, aaf_1kg_sas, aaf_1kg_afr, aaf_1kg_eur, aaf_1kg_all, aaf_exac_all, aaf_adj_exac_all, " \
                "aaf_adj_exac_afr, aaf_adj_exac_amr, aaf_adj_exac_eas, aaf_adj_exac_fin, aaf_adj_exac_nfe, " \
                "aaf_adj_exac_oth, aaf_adj_exac_sas, max_aaf_all, in_esp, in_1kg, in_exac, info," \
                "(gts).(*), (gt_depths).(*), (gt_ref_depths).(*), (gt_alt_depths).(*) FROM variants"

        sys.stdout.write("Running GEMINI query\n")
        gq = GeminiQuery(samples[sample]['db'])
        gq.run(query)

        caller_vcf_records = defaultdict(lambda: dict())

        parse_vcf(samples[sample]['mutect'], "mutect", caller_vcf_records)
        parse_vcf(samples[sample]['vardict'], "vardict", caller_vcf_records)
        parse_vcf(samples[sample]['freebayes'], "freebayes", caller_vcf_records)
        parse_vcf(samples[sample]['scalpel'], "scalpel", caller_vcf_records)

        # Filter out variants with minor allele frequencies above the threshold but
        # retain any that are above the threshold but in COSMIC or in ClinVar and not listed as benign.
        for variant_data in gq:
            cassandra_variant = Variant(chr=variant_data['chrom'], start=variant_data['start'], end=variant_data['end'],
                                        ref=variant_data['ref'], alt=variant_data['alt'],
                                        sample=samples[sample]['sample_name'], extraction=samples[sample]['extraction'],
                                        library_name=sample, panel_name=samples[sample]['panel'],
                                        target_pool=samples[sample]['target_pool'],
                                        reference_genome=config['genome_version'], date_annotated=datetime.now(),
                                        type=variant_data['type'],
                                        subtype=variant_data['sub_type'], rs_id=variant_data['vcf_id'],
                                        gene=variant_data['gene'], max_aaf=variant_data['max_aaf_all'],
                                        transcript=variant_data['transcript'], exon=variant_data['exon'],
                                        codon_change=variant_data['codon_change'],
                                        aa_change=variant_data['aa_change'], biotype=variant_data['biotype'],
                                        impact=variant_data['impact'], impact_so=variant_data['impact_so'])

            cassandra_variant['in_clinvar'] = gemini_interface.var_is_in_clinvar(variant_data)
            cassandra_variant['in_cosmic'] = gemini_interface.var_is_in_cosmic(variant_data)
            cassandra_variant['is_pathogenic'] = gemini_interface.var_is_pathogenic(variant_data)

            if variant_data['is_lof']:
                cassandra_variant['is_lof'] = True
            else:
                cassandra_variant['is_lof'] = False

            if variant_data['is_coding']:
                cassandra_variant['is_coding'] = True
            else:
                cassandra_variant['is_coding'] = False

            if variant_data['is_splicing']:
                cassandra_variant['is_splicing'] = True
            else:
                cassandra_variant['is_splicing'] = False

            if variant_data['rs_ids'] is not None:
                cassandra_variant['rs_ids'] = variant_data['rs_ids'].split(',')

            if variant_data['cosmic_ids'] is not None:
                cassandra_variant['cosmic_ids'] = variant_data['cosmic_ids'].split(',')

            if variant_data['info']['CALLERS'] is not None:
                cassandra_variant['callers'] = variant_data['info']['CALLERS'].split(',')

            population_freqs = {'esp_ea': variant_data['aaf_esp_ea'],
                                'esp_aa': variant_data['aaf_esp_aa'],
                                'esp_all': variant_data['aaf_esp_all'],
                                '1kg_amr': variant_data['aaf_1kg_amr'],
                                '1kg_eas': variant_data['aaf_1kg_eas'],
                                '1kg_sas': variant_data['aaf_1kg_sas'],
                                '1kg_afr': variant_data['aaf_1kg_afr'],
                                '1kg_eur': variant_data['aaf_1kg_eur'],
                                '1kg_all': variant_data['aaf_1kg_all'],
                                'exac_all': variant_data['aaf_exac_all'],
                                'adj_exac_all': variant_data['aaf_adj_exac_all'],
                                'adj_exac_afr': variant_data['aaf_adj_exac_afr'],
                                'adj_exac_amr': variant_data['aaf_adj_exac_amr'],
                                'adj_exac_eas': variant_data['aaf_adj_exac_eas'],
                                'adj_exac_fin': variant_data['aaf_adj_exac_fin'],
                                'adj_exac_nfe': variant_data['aaf_adj_exac_nfe'],
                                'adj_exac_oth': variant_data['aaf_adj_exac_oth'],
                                'adj_exac_sas': variant_data['aaf_adj_exac_sas']}

            cassandra_variant['population_freqs'] = population_freqs

            key = (variant_data['chrom'], variant_data['start'], variant_data['end'], variant_data['ref'],
                   variant_data['alt'])

            if 'mutect' in variant_data['info']['CALLERS']:
                record = caller_vcf_records['mutect'][key]
                # filters = record.FILTER.split(',')

                mutect_info = {'GTF_DP': str(record.gt_depths[0]),
                               'GTF_AD': str(record.gt_alt_depths[0])}

                cassandra_variant['mutect'] = mutect_info

            if 'vardict' in variant_data['info']['CALLERS']:
                record = caller_vcf_records['vardict'][key]
                # filters = record.FILTER.split(',')

                vardict_info = {'DP': str(record.INFO.get('DP')),
                                'VD': str(record.INFO.get('VD')),
                                'AF': str(record.INFO.get('AF')),
                                'BIAS': str(record.INFO.get('BIAS')),
                                'REFBIAS': str(record.INFO.get('REFBIAS')),
                                'VARBIAS': str(record.INFO.get('VARBIAS')),
                                'QUAL': str(record.INFO.get('QUAL')),
                                'QSTD': str(record.INFO.get('QSTD')),
                                'SBF': str(record.INFO.get('SBF')),
                                'ODDRATIO': str(record.INFO.get('ODDRATIO')),
                                'MQ': str(record.INFO.get('MQ')),
                                'SN': str(record.INFO.get('SN')),
                                'HIAF': str(record.INFO.get('HIAF')),
                                'ADJAF': str(record.INFO.get('ADJAF')),
                                'MSI': str(record.INFO.get('MSI')),
                                'MSILEN': str(record.INFO.get('MSILEN')),
                                'SHIFT3': str(record.INFO.get('SHIFT3')),
                                'NM': str(record.INFO.get('NM')),
                                'GDAMP': str(record.INFO.get('GDAMP')),
                                'LSEQ': str(record.INFO.get('LSEQ')),
                                'RSEQ': str(record.INFO.get('RSEQ')),
                                'TLAMP': str(record.INFO.get('TLAMP')),
                                'NCAMP': str(record.INFO.get('NCAMP')),
                                'AMPFLAG': str(record.INFO.get('AMPFLAG')),
                                'HICNT': str(record.INFO.get('HICNT')),
                                'HICOV': str(record.INFO.get('HICOV')),
                                'GTF_DP': str(record.gt_depths[0]),
                                'GTF_AD': str(record.gt_alt_depths[0])}

                cassandra_variant['vardict'] = vardict_info

            if 'freebayes' in variant_data['info']['CALLERS']:
                record = caller_vcf_records['freebayes'][key]
                # filters = record.FILTER.split(',')

                freebayes_info = {'DP': str(record.INFO.get('DP')),
                                  'AF': str(record.INFO.get('AF')),
                                  'AC': str(record.INFO.get('AC')),
                                  'RO': str(record.INFO.get('RO')),
                                  'AO': str(record.INFO.get('AO')),
                                  'PRO': str(record.INFO.get('PRO')),
                                  'PAO': str(record.INFO.get('PAO')),
                                  'QR': str(record.INFO.get('QR')),
                                  'QA': str(record.INFO.get('QA')),
                                  'PQR': str(record.INFO.get('PQR')),
                                  'PQA': str(record.INFO.get('PQA')),
                                  'SRF': str(record.INFO.get('SRF')),
                                  'SRR': str(record.INFO.get('SRR')),
                                  'SAF': str(record.INFO.get('SAF')),
                                  'SAR': str(record.INFO.get('SAR')),
                                  'SRP': str(record.INFO.get('SRP')),
                                  'SAP': str(record.INFO.get('SAP')),
                                  'AB': str(record.INFO.get('AB')),
                                  'ABP': str(record.INFO.get('ABP')),
                                  'RUN': str(record.INFO.get('RUN')),
                                  'RPP': str(record.INFO.get('RPP')),
                                  'RPPR': str(record.INFO.get('RPPR')),
                                  'RPL': str(record.INFO.get('RPL')),
                                  'RPR': str(record.INFO.get('RPR')),
                                  'EPP': str(record.INFO.get('EPP')),
                                  'EPPR': str(record.INFO.get('EPPR')),
                                  'DRPA': str(record.INFO.get('DRPA')),
                                  'ODDS': str(record.INFO.get('ODDS')),
                                  'GTI': str(record.INFO.get('GTI')),
                                  'TYPE': str(record.INFO.get('TYPE')),
                                  'CIGAR': str(record.INFO.get('CIGAR')),
                                  'NUMALT': str(record.INFO.get('NUMALT')),
                                  'MEANALT': str(record.INFO.get('MEANALT')),
                                  'LEN': str(record.INFO.get('LEN')),
                                  'MQM': str(record.INFO.get('MQM')),
                                  'MQMR': str(record.INFO.get('MQMR')),
                                  'PAIRED': str(record.INFO.get('PAIRED')),
                                  'PAIREDR': str(record.INFO.get('PAIREDR')),
                                  'GTF_DP': str(record.gt_depths[0])}

                cassandra_variant['freebayes'] = freebayes_info

            if 'scalpel' in variant_data['info']['CALLERS']:
                record = caller_vcf_records['scalpel'][key]
                # filters = record.FILTER.split(',')

                scalpel_info = {'AVGCOV': str(record.INFO.get('AVGCOV')),
                                'MINCOV': str(record.INFO.get('MINCOV')),
                                'ALTCOV': str(record.INFO.get('ALTCOV')),
                                'COVRATIO': str(record.INFO.get('COVRATIO')),
                                'ZYG': str(record.INFO.get('ZYG')),
                                'CHI2': str(record.INFO.get('CHI2')),
                                'FISHERPHREDSCORE': str(record.INFO.get('FISHERPHREDSCORE')),
                                'INH': str(record.INFO.get('INH')),
                                'BESTSTATE': str(record.INFO.get('BESTSTATE')),
                                'COVSTATE': str(record.INFO.get('COVSTATE')),
                                'SOMATIC': str(record.INFO.get('SOMATIC')),
                                'DENOVO': str(record.INFO.get('DENOVO')),
                                'GTF_DP': str(record.gt_depths[0]),
                                'GTF_AD': str(record.gt_alt_depths[0])}

                cassandra_variant['scalpel'] = scalpel_info

            cassandra_variant.save()
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
Turn the gitenv into RPM spec files which can be built at a later stage.

"""
from __future__ import print_function

import datetime
from glob import glob
import os
import shutil
import time

import conda.api
import conda.fetch
from conda.resolve import Resolve, MatchSpec
from conda_gitenv import manifest_branch_prefix
from conda_gitenv.deploy import tags_by_label, tags_by_env
from conda_gitenv.lock import Locked
from conda_gitenv.resolve import tempdir, create_tracking_branches
from git import Repo, Commit
import yaml

import logging
import conda_rpms.generate as generate
import conda_rpms.install as conda_install

class Config(dict):
    def __init__(self, fname, store=None, key=None):
        self.fname = os.path.abspath(os.path.expanduser(fname))
        self._store = store
        if store is None:
            self._load()
        self._key = []
        if key is not None:
            self._key = key

    def _load(self):
        if not os.path.exists(self.fname):
            emsg = 'The configuration file {!r} does not exist.'
            raise ValueError(emsg.format(os.path.basename(self.fname)))
        with open(self.fname, 'r') as fh:
            try:
                self._store = yaml.safe_load(fh)
            except yaml.YAMLError as e:
                emsg = 'YAML error in configuration file {!r}: {}'
                line, column = e.problem_mark.line, e.problem_mark.column
                ymsg = '{}, line {}, column {}.'.format(e.context,
                                                        line+1,
                                                        column+1)
                raise ValueError(emsg.format(os.path.basename(self.fname),
                                             ymsg))

    def __getitem__(self, key):
        try:
            result = self._store[key]
        except KeyError:
            emsg = 'The YAML file {!r} does not contain key [{}].'
            full_key = ']['.join(self._key + [key])
            raise ValueError(emsg.format(os.path.basename(self.fname),
                                         full_key))
        if isinstance(result, dict):
            result = Config(self.fname, result, self._key + [key])
        return result

    def __iter__(self):
        return iter(self._store)

    def __len__(self):
        return len(self._store)

    def __repr__(self):
        return repr(self._store)

def create_rpmbuild_for_env(pkgs, target, config):
    rpm_prefix = config['rpm']['prefix']
    pkg_cache = os.path.join(target, 'SOURCES')
    pkg_names = set(pkg for _, pkg in pkgs)
    if os.path.exists(target):
        # The environment we want to deploy already exists. We should
        # just double check that there aren't already packages in there which
        # we need to remove before we install anything new.
        linked = conda_install.linked(target)
        for pkg in linked:
            if pkg not in pkg_names:
                conda_install.unlink(target, pkg)
    else:
        linked = []

    if set(linked) == pkg_names:
        # We don't need to re-link everything - it is already as expected.
        # The downside is that we are not verifying that each package is
        # installed correctly.
        return

    spec_dir = os.path.join(target, 'SPECS')
    if not os.path.exists(spec_dir):
        os.makedirs(spec_dir)
    for source, pkg in pkgs:
        index = conda.fetch.fetch_index([source], use_cache=False)
        pkg_index = {pkg_info['fn']: pkg_info for pkg_info in index.values()}
        tar_name = pkg + '.tar.bz2'
        pkg_info = pkg_index.get(tar_name, None)
        if pkg_info is None:
            raise ValueError('Distribution {} is no longer available '
                             'in the channel {}.'.format(tar_name, source))
        dist_name = pkg
        if not conda_install.is_fetched(pkg_cache, dist_name):
            print('Fetching {}'.format(dist_name))
            conda.fetch.fetch_pkg(pkg_info, pkg_cache)
        spec_path = os.path.join(spec_dir, '{}-pkg-{}.spec'.format(rpm_prefix,
                                                                   pkg))
        if not os.path.exists(spec_path):
            spec = generate.render_dist_spec(os.path.join(pkg_cache,
                                                          tar_name), config)
            with open(spec_path, 'w') as fh:
                fh.write(spec)

def create_rpmbuild_for_tag(repo, tag_name, target, config):
    rpm_prefix = config['rpm']['prefix']
    print("CREATE FOR {}".format(tag_name))
    tag = repo.tags[tag_name]
    # Checkout the tag in a detached head form.
    repo.head.reference = tag.commit
    repo.head.reset(working_tree=True)

    manifest_fname = os.path.join(repo.working_dir, 'env.manifest')
    if not os.path.exists(manifest_fname):
        raise ValueError("The tag '{}' doesn't have a manifested "
                         "environment.".format(tag_name))
    with open(manifest_fname, 'r') as fh:
        manifest = sorted(line.strip().split('\t') for line in fh)

    spec_fname = os.path.join(repo.working_dir, 'env.spec')
    if not os.path.exists(spec_fname):
        raise ValueError("The tag '{}' doesn't have an environment specification.".format(tag_name))
    with open(spec_fname, 'r') as fh:
        env_spec = yaml.safe_load(fh).get('env', [])
    create_rpmbuild_for_env(manifest, target, config)
    pkgs = [pkg for _, pkg in manifest]
    env_name, tag = tag_name.split('-')[1:]
    fname = '{}-env-{}-tag-{}.spec'.format(rpm_prefix, env_name, tag)
    with open(os.path.join(target, 'SPECS', fname), 'w') as fh:
        fh.write(generate.render_taggedenv(env_name, tag, pkgs, config, env_spec))

def create_rpmbuild_content(repo, target, config):
    rpm_prefix = config['rpm']['prefix']
    for branch in repo.branches:
        # We only want environment branches, not manifest branches.
        if not branch.name.startswith(manifest_branch_prefix):
            manifest_branch_name = manifest_branch_prefix + branch.name
            # If there is no equivalent manifest branch, we need to
            # skip this environment.
            if manifest_branch_name not in repo.branches:
                continue
            branch.checkout()
            labelled_tags = tags_by_label(os.path.join(repo.working_dir,
                                                       'labels'))

            # Get number of commits to determine the version of the env rpm.
            commit_num = len(list(Commit.iter_items(repo, branch.commit)))

            # Keep track of the labels which have tags - its those we want.
            for label, tag in labelled_tags.items():
                create_rpmbuild_for_tag(repo, tag, target, config)
                fname = '{}-env-{}-label-{}.spec'.format(rpm_prefix, branch.name, label)
                with open(os.path.join(target, 'SPECS', fname), 'w') as fh:
                    fh.write(generate.render_env(branch.name, label,
                                                 repo, config, tag, commit_num))

def create_rpm_installer(target, config, python_spec='python'):
    rpm_prefix = config['rpm']['prefix']
    index = conda.api.get_index()
    matches = Resolve(index).get_pkgs(MatchSpec(python_spec))
    if not matches:
        raise RuntimeError('No python found in the channels.')
    # Pick the latest Python match.
    pkg_info = sorted(matches)[-1].info
    dist_name = '{}-{}-{}'.format(pkg_info['name'], pkg_info['version'],
                                  pkg_info['build'])
    pkg_cache = os.path.join(target, 'SOURCES')
    if not conda_install.is_fetched(pkg_cache, dist_name):
        print('Fetching {}'.format(dist_name))
        conda.fetch.fetch_pkg(pkg_info, pkg_cache)

    installer_source = os.path.join(os.path.dirname(__file__), 'install.py')
    installer_target = os.path.join(pkg_cache, 'install.py')

    shutil.copyfile(installer_source, installer_target)

    spec_dir = os.path.join(target, 'SPECS')
    if not os.path.exists(spec_dir):
        os.makedirs(spec_dir)

    specfile = os.path.join(spec_dir, '{}-installer.spec'.format(rpm_prefix))
    with open(specfile, 'w') as fh:
        fh.write(generate.render_installer(pkg_info, config))

def configure_parser(parser):
    parser.add_argument('repo_uri', help='Repo to deploy.')
    parser.add_argument('target', help='Location to put the RPMBUILD content.')
    parser.add_argument('--config', '-c', type=str, default='config.yaml',
                        help='YAML configuration filename.')
    parser.set_defaults(function=handle_args)
    return parser

def handle_args(args):
    # To reduce the noise coming from conda/conda-build we set
    # all loggers to WARN level.
    logging.getLogger('').setLevel(logging.WARNING)
    for logger_name in logging.Logger.manager.loggerDict.keys():
        logger = logging.getLogger(logger_name)
        logger.setLevel(logging.WARNING)

    config = Config(args.config)
    with tempdir() as repo_directory:
        repo = Repo.clone_from(args.repo_uri, repo_directory)
        create_tracking_branches(repo)
        create_rpmbuild_content(repo, args.target, config)
        create_rpm_installer(args.target, config)

def main():
    import argparse
    parser = argparse.ArgumentParser(description='Deploy the tracked '
                                                 'environments.')
    configure_parser(parser)
    args = parser.parse_args()
    return args.function(args)

if __name__ == '__main__':
    main()
<commit_msg>Fix bad tag resolution. (#39)<commit_after>#!/usr/bin/env python
"""
Turn the gitenv into RPM spec files which can be built at a later stage.

"""
from __future__ import print_function

import datetime
from glob import glob
import os
import shutil
import time

import conda.api
import conda.fetch
from conda.resolve import Resolve, MatchSpec
from conda_gitenv import manifest_branch_prefix
from conda_gitenv.deploy import tags_by_label, tags_by_env
from conda_gitenv.lock import Locked
from conda_gitenv.resolve import tempdir, create_tracking_branches
from git import Repo, Commit
import yaml

import logging
import conda_rpms.generate as generate
import conda_rpms.install as conda_install

class Config(dict):
    def __init__(self, fname, store=None, key=None):
        self.fname = os.path.abspath(os.path.expanduser(fname))
        self._store = store
        if store is None:
            self._load()
        self._key = []
        if key is not None:
            self._key = key

    def _load(self):
        if not os.path.exists(self.fname):
            emsg = 'The configuration file {!r} does not exist.'
            raise ValueError(emsg.format(os.path.basename(self.fname)))
        with open(self.fname, 'r') as fh:
            try:
                self._store = yaml.safe_load(fh)
            except yaml.YAMLError as e:
                emsg = 'YAML error in configuration file {!r}: {}'
                line, column = e.problem_mark.line, e.problem_mark.column
                ymsg = '{}, line {}, column {}.'.format(e.context,
                                                        line+1,
                                                        column+1)
                raise ValueError(emsg.format(os.path.basename(self.fname),
                                             ymsg))

    def __getitem__(self, key):
        try:
            result = self._store[key]
        except KeyError:
            emsg = 'The YAML file {!r} does not contain key [{}].'
            full_key = ']['.join(self._key + [key])
            raise ValueError(emsg.format(os.path.basename(self.fname),
                                         full_key))
        if isinstance(result, dict):
            result = Config(self.fname, result, self._key + [key])
        return result

    def __iter__(self):
        return iter(self._store)

    def __len__(self):
        return len(self._store)

    def __repr__(self):
        return repr(self._store)

def create_rpmbuild_for_env(pkgs, target, config):
    rpm_prefix = config['rpm']['prefix']
    pkg_cache = os.path.join(target, 'SOURCES')
    pkg_names = set(pkg for _, pkg in pkgs)
    if os.path.exists(target):
        # The environment we want to deploy already exists. We should
        # just double check that there aren't already packages in there which
        # we need to remove before we install anything new.
        linked = conda_install.linked(target)
        for pkg in linked:
            if pkg not in pkg_names:
                conda_install.unlink(target, pkg)
    else:
        linked = []

    if set(linked) == pkg_names:
        # We don't need to re-link everything - it is already as expected.
        # The downside is that we are not verifying that each package is
        # installed correctly.
        return

    spec_dir = os.path.join(target, 'SPECS')
    if not os.path.exists(spec_dir):
        os.makedirs(spec_dir)
    for source, pkg in pkgs:
        index = conda.fetch.fetch_index([source], use_cache=False)
        pkg_index = {pkg_info['fn']: pkg_info for pkg_info in index.values()}
        tar_name = pkg + '.tar.bz2'
        pkg_info = pkg_index.get(tar_name, None)
        if pkg_info is None:
            raise ValueError('Distribution {} is no longer available '
                             'in the channel {}.'.format(tar_name, source))
        dist_name = pkg
        if not conda_install.is_fetched(pkg_cache, dist_name):
            print('Fetching {}'.format(dist_name))
            conda.fetch.fetch_pkg(pkg_info, pkg_cache)
        spec_path = os.path.join(spec_dir, '{}-pkg-{}.spec'.format(rpm_prefix,
                                                                   pkg))
        if not os.path.exists(spec_path):
            spec = generate.render_dist_spec(os.path.join(pkg_cache,
                                                          tar_name), config)
            with open(spec_path, 'w') as fh:
                fh.write(spec)

def create_rpmbuild_for_tag(repo, tag_name, target, config):
    rpm_prefix = config['rpm']['prefix']
    print("CREATE FOR {}".format(tag_name))
    tag = repo.tags[tag_name]
    # Checkout the tag in a detached head form.
    repo.head.reference = tag.commit
    repo.head.reset(working_tree=True)

    manifest_fname = os.path.join(repo.working_dir, 'env.manifest')
    if not os.path.exists(manifest_fname):
        raise ValueError("The tag '{}' doesn't have a manifested "
                         "environment.".format(tag_name))
    with open(manifest_fname, 'r') as fh:
        manifest = sorted(line.strip().split('\t') for line in fh)

    spec_fname = os.path.join(repo.working_dir, 'env.spec')
    if not os.path.exists(spec_fname):
        raise ValueError("The tag '{}' doesn't have an environment specification.".format(tag_name))
    with open(spec_fname, 'r') as fh:
        env_spec = yaml.safe_load(fh).get('env', [])
    create_rpmbuild_for_env(manifest, target, config)
    pkgs = [pkg for _, pkg in manifest]
    env_name, tag = tag_name.split('-', 2)[1:]
    fname = '{}-env-{}-tag-{}.spec'.format(rpm_prefix, env_name, tag)
    with open(os.path.join(target, 'SPECS', fname), 'w') as fh:
        fh.write(generate.render_taggedenv(env_name, tag, pkgs, config, env_spec))

def create_rpmbuild_content(repo, target, config):
    rpm_prefix = config['rpm']['prefix']
    for branch in repo.branches:
        # We only want environment branches, not manifest branches.
        if not branch.name.startswith(manifest_branch_prefix):
            manifest_branch_name = manifest_branch_prefix + branch.name
            # If there is no equivalent manifest branch, we need to
            # skip this environment.
            if manifest_branch_name not in repo.branches:
                continue
            branch.checkout()
            labelled_tags = tags_by_label(os.path.join(repo.working_dir,
                                                       'labels'))

            # Get number of commits to determine the version of the env rpm.
            commit_num = len(list(Commit.iter_items(repo, branch.commit)))

            # Keep track of the labels which have tags - its those we want.
            for label, tag in labelled_tags.items():
                create_rpmbuild_for_tag(repo, tag, target, config)
                fname = '{}-env-{}-label-{}.spec'.format(rpm_prefix, branch.name, label)
                with open(os.path.join(target, 'SPECS', fname), 'w') as fh:
                    fh.write(generate.render_env(branch.name, label,
                                                 repo, config, tag, commit_num))

def create_rpm_installer(target, config, python_spec='python'):
    rpm_prefix = config['rpm']['prefix']
    index = conda.api.get_index()
    matches = Resolve(index).get_pkgs(MatchSpec(python_spec))
    if not matches:
        raise RuntimeError('No python found in the channels.')
    # Pick the latest Python match.
    pkg_info = sorted(matches)[-1].info
    dist_name = '{}-{}-{}'.format(pkg_info['name'], pkg_info['version'],
                                  pkg_info['build'])
    pkg_cache = os.path.join(target, 'SOURCES')
    if not conda_install.is_fetched(pkg_cache, dist_name):
        print('Fetching {}'.format(dist_name))
        conda.fetch.fetch_pkg(pkg_info, pkg_cache)

    installer_source = os.path.join(os.path.dirname(__file__), 'install.py')
    installer_target = os.path.join(pkg_cache, 'install.py')

    shutil.copyfile(installer_source, installer_target)

    spec_dir = os.path.join(target, 'SPECS')
    if not os.path.exists(spec_dir):
        os.makedirs(spec_dir)

    specfile = os.path.join(spec_dir, '{}-installer.spec'.format(rpm_prefix))
    with open(specfile, 'w') as fh:
        fh.write(generate.render_installer(pkg_info, config))

def configure_parser(parser):
    parser.add_argument('repo_uri', help='Repo to deploy.')
    parser.add_argument('target', help='Location to put the RPMBUILD content.')
    parser.add_argument('--config', '-c', type=str, default='config.yaml',
                        help='YAML configuration filename.')
    parser.set_defaults(function=handle_args)
    return parser

def handle_args(args):
    # To reduce the noise coming from conda/conda-build we set
    # all loggers to WARN level.
    logging.getLogger('').setLevel(logging.WARNING)
    for logger_name in logging.Logger.manager.loggerDict.keys():
        logger = logging.getLogger(logger_name)
        logger.setLevel(logging.WARNING)

    config = Config(args.config)
    with tempdir() as repo_directory:
        repo = Repo.clone_from(args.repo_uri, repo_directory)
        create_tracking_branches(repo)
        create_rpmbuild_content(repo, args.target, config)
        create_rpm_installer(args.target, config)

def main():
    import argparse
    parser = argparse.ArgumentParser(description='Deploy the tracked '
                                                 'environments.')
    configure_parser(parser)
    args = parser.parse_args()
    return args.function(args)

if __name__ == '__main__':
    main()
<|endoftext|>