<commit_before># Copyright (c) 2011-2012 OpenStack, LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gettext

#: Version information (major, minor, revision[, 'dev']).
version_info = (1, 1, 11,)
#: Version string 'major.minor.revision'.
version = __version__ = ".".join(map(str, version_info))
gettext.install('sos')
<commit_msg>upping vrs<commit_after># Copyright (c) 2011-2012 OpenStack, LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gettext

#: Version information (major, minor, revision[, 'dev']).
version_info = (1, 1, 12,)
#: Version string 'major.minor.revision'.
version = __version__ = ".".join(map(str, version_info))
gettext.install('sos')
<|endoftext|>


<commit_before>import time
from dtest import Tester, debug
from tools import no_vnodes
from threading import Event

class NotificationWaiter(object):
    """
    A helper class for waiting for pushed notifications from
    Cassandra over the native protocol.
    """

    def __init__(self, tester, node, notification_type, num_notifications=1):
        """
        `address` should be a ccmlib.node.Node instance
        `notification_type` should either be "TOPOLOGY_CHANGE", "STATUS_CHANGE",
        or "SCHEMA_CHANGE".
        """
        self.address = node.network_interfaces['binary'][0]
        self.notification_type = notification_type

        # get a single, new connection
        session = tester.patient_cql_connection(node)
        connection = session.cluster.connection_factory(self.address, is_control_connection=True)

        # coordinate with an Event
        self.event = Event()

        # the pushed notification
        self.notifications = []
        self.num_notifications = num_notifications

        # register a callback for the notification type
        connection.register_watcher(
            notification_type, self.handle_notification, register_timeout=5.0)

    def handle_notification(self, notification):
        """
        Called when a notification is pushed from Cassandra.
        """

        debug("Source %s sent %s for %s" % (self.address, notification["change_type"], notification["address"][0],))

        self.notifications.append(notification)
        self.event.set()

    def wait_for_notifications(self, timeout):
        """
        Waits up to `timeout` seconds for a notifications from Cassandra.
        """

        deadline = time.time() + timeout
        while time.time() < deadline:
            self.event.wait(deadline - time.time())
            self.event.clear()
            if len(self.notifications) >= self.num_notifications:
                break

        if len(self.notifications) < self.num_notifications:
            raise Exception("Timed out waiting for %s notifications from %s"
                            % (self.notification_type, self.address))
        else:
            return self.notifications

class TestPushedNotifications(Tester):
    """
    Tests for pushed native protocol notification from Cassandra.
    """

    @no_vnodes()
    def move_single_node_test(self):
        """
        Moving a token should result in NODE_MOVED notifications.
        """
        self.cluster.populate(3).start()

        waiters = [NotificationWaiter(self, node, "TOPOLOGY_CHANGE")
                   for node in self.cluster.nodes.values()]

        node1 = self.cluster.nodes.values()[0]
        node1.move("123")

        for waiter in waiters:
            debug("Waiting for notification from %s" % (waiter.address,))
            notifications = waiter.wait_for_notifications(60.0)
            self.assertEquals(1, len(notifications))
            notification = notifications[0]
            change_type = notification["change_type"]
            address, port = notification["address"]
            self.assertEquals("MOVED_NODE", change_type)
            self.assertEquals(self.get_ip_from_node(node1), address)

    def restart_node_test(self):
        """
        Restarting a node should generate exactly one DOWN and one UP notification
        """

        self.cluster.populate(2).start()

        node1, node2 = self.cluster.nodelist()

        waiter = NotificationWaiter(self, node1, "STATUS_CHANGE", num_notifications=2)

        for i in range(3):
            debug("Restarting second node...")
            node2.stop()
            node2.start()
            debug("Waiting for notifications from %s" % (waiter.address,))
            notifications = waiter.wait_for_notifications(60.0)
            self.assertEquals(2, len(notifications))
            self.assertEquals(self.get_ip_from_node(node2), notifications[0]["address"][0])
            self.assertEquals("DOWN", notifications[0]["change_type"])
            self.assertEquals(self.get_ip_from_node(node2), notifications[1]["address"][0])
            self.assertEquals("UP", notifications[1]["change_type"])

<commit_msg>Fixed bug for test of CASSANDRA-7816<commit_after>import time
from dtest import Tester, debug
from tools import no_vnodes
from threading import Event

class NotificationWaiter(object):
    """
    A helper class for waiting for pushed notifications from
    Cassandra over the native protocol.
    """

    def __init__(self, tester, node, notification_type, num_notifications=1):
        """
        `address` should be a ccmlib.node.Node instance
        `notification_type` should either be "TOPOLOGY_CHANGE", "STATUS_CHANGE",
        or "SCHEMA_CHANGE".
        """
        self.address = node.network_interfaces['binary'][0]
        self.notification_type = notification_type

        # get a single, new connection
        session = tester.patient_cql_connection(node)
        connection = session.cluster.connection_factory(self.address, is_control_connection=True)

        # coordinate with an Event
        self.event = Event()

        # the pushed notification
        self.notifications = []
        self.num_notifications = num_notifications

        # register a callback for the notification type
        connection.register_watcher(
            notification_type, self.handle_notification, register_timeout=5.0)

    def handle_notification(self, notification):
        """
        Called when a notification is pushed from Cassandra.
        """

        debug("Source %s sent %s for %s" % (self.address, notification["change_type"], notification["address"][0],))

        self.notifications.append(notification)
        self.event.set()

    def wait_for_notifications(self, timeout):
        """
        Waits up to `timeout` seconds for a notifications from Cassandra.
        """

        deadline = time.time() + timeout
        while time.time() < deadline:
            self.event.wait(deadline - time.time())
            self.event.clear()
            if len(self.notifications) >= self.num_notifications:
                break

        if len(self.notifications) < self.num_notifications:
            raise Exception("Timed out waiting for %s notifications from %s"
                            % (self.notification_type, self.address))
        else:
            return self.notifications

    def clear_notifications(self):
        self.notifications = []
        self.event.clear()

class TestPushedNotifications(Tester):
    """
    Tests for pushed native protocol notification from Cassandra.
    """

    @no_vnodes()
    def move_single_node_test(self):
        """
        Moving a token should result in NODE_MOVED notifications.
        """
        self.cluster.populate(3).start()

        waiters = [NotificationWaiter(self, node, "TOPOLOGY_CHANGE")
                   for node in self.cluster.nodes.values()]

        node1 = self.cluster.nodes.values()[0]
        node1.move("123")

        for waiter in waiters:
            debug("Waiting for notification from %s" % (waiter.address,))
            notifications = waiter.wait_for_notifications(60.0)
            self.assertEquals(1, len(notifications))
            notification = notifications[0]
            change_type = notification["change_type"]
            address, port = notification["address"]
            self.assertEquals("MOVED_NODE", change_type)
            self.assertEquals(self.get_ip_from_node(node1), address)

    def restart_node_test(self):
        """
        Restarting a node should generate exactly one DOWN and one UP notification
        """

        self.cluster.populate(2).start()
        node1, node2 = self.cluster.nodelist()

        waiter = NotificationWaiter(self, node1, "STATUS_CHANGE", num_notifications=2)

        for i in range(5):
            debug("Restarting second node...")
            node2.stop()
            node2.start()
            debug("Waiting for notifications from %s" % (waiter.address,))
            notifications = waiter.wait_for_notifications(60.0)
            self.assertEquals(2, len(notifications))
            self.assertEquals(self.get_ip_from_node(node2), notifications[0]["address"][0])
            self.assertEquals("DOWN", notifications[0]["change_type"])
            self.assertEquals(self.get_ip_from_node(node2), notifications[1]["address"][0])
            self.assertEquals("UP", notifications[1]["change_type"])
            waiter.clear_notifications()

<|endoftext|>


<commit_before>"""
Plot running time vs generation.
"""

import sys
import matplotlib
import os

def main():
    args = sys.argv[1:]
    if len(args) == 0:
        print 'autotune_plot tune_X/summary.txt [out.png]'

    filename = args[0]
    out_filename = args[1] if len(args) > 1 else None
    tunedir = os.path.split(os.path.split(os.path.abspath(filename))[0])[1]
    # Parse summary file (a bit of a hack -- should really dump a plot xy points file)
    L = open(filename, 'rt').read().strip().split('\n')
    gen = 1
    time_d = {}
    for i in range(len(L)):
        if len(L[i].split()) == 0:
            continue
        val = L[i].split()[0]
        try:
            val = float(val)
            if not gen in time_d:
                time_d[gen] = val
        except ValueError:
            if L[i].startswith('Generation'):
                gen = int(L[i].split()[1])
    x = []
    y = []
    for (xv, yv) in time_d.items():
        x.append(xv)
        y.append(yv)

    if out_filename is not None:
        matplotlib.use('Agg')     # headless

    import pylab

    pylab.plot(x, y)
    pylab.title('Best time vs Generation\n(%s)'%tunedir)
    pylab.ylabel('Best time [sec]')
    pylab.xlabel('Generation')
    if out_filename is not None:
        pylab.savefig(out_filename, dpi=300)
    else:
        pylab.show()

if __name__ == '__main__':
    main()
<commit_msg>py bindings -- minor plot tweak<commit_after>"""
Plot running time vs generation.
"""

import sys
import matplotlib
import os

def main():
    args = sys.argv[1:]
    if len(args) == 0:
        print 'autotune_plot tune_X/summary.txt [out.png]'

    filename = args[0]
    out_filename = args[1] if len(args) > 1 else None
    tunedir = os.path.split(os.path.split(os.path.abspath(filename))[0])[1]
    # Parse summary file (a bit of a hack -- should really dump a plot xy points file)
    L = open(filename, 'rt').read().strip().split('\n')
    gen = 1
    time_d = {}
    for i in range(len(L)):
        if len(L[i].split()) == 0:
            continue
        val = L[i].split()[0]
        try:
            val = float(val)
            if not gen in time_d:
                time_d[gen] = val
        except ValueError:
            if L[i].startswith('Generation'):
                gen = int(L[i].split()[1])
    x = []
    y = []
    for (xv, yv) in time_d.items():
        x.append(xv)
        y.append(yv*1000)

    if out_filename is not None:
        matplotlib.use('Agg')     # headless

    import pylab

    pylab.plot(x, y)
    pylab.title('Best time vs Generation\n(%s)'%tunedir)
    pylab.ylabel('Best time [ms]')
    pylab.xlabel('Generation')
    if out_filename is not None:
        pylab.savefig(out_filename, dpi=300)
    else:
        pylab.show()

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


<commit_before>import collections
import os
import warnings

import astropy.io.fits as pyfits
import numpy as np
import scipy.integrate
from math import log

from threeML.minimizer import minimization
from threeML.plugin_prototype import PluginPrototype
from threeML.plugins.gammaln import logfactorial
from threeML.plugins.ogip import OGIPPHA

from astromodels.parameter import Parameter

__instrument_name = "Fermi GBM (all detectors)"

class FermiGBMLike(PluginPrototype):

    def __init__(self, name, phafile, bkgfile, rspfile):
        '''
        If the input files are PHA2 files, remember to specify the spectrum number, for example:
        FermiGBMLike("GBM","spectrum.pha{2}","bkgfile.bkg{2}","rspfile.rsp{2}")
        to load the second spectrum, second background spectrum and second response.
        '''

        self.name = name

        # Check that all file exists
        notExistant = []

        if (not os.path.exists(phafile.split("{")[0])):
            notExistant.append(phafile.split("{")[0])

        if (not os.path.exists(bkgfile.split("{")[0])):
            notExistant.append(bkgfile.split("{")[0])

        if (not os.path.exists(rspfile.split("{")[0])):
            notExistant.append(rspfile.split("{")[0])

        if (len(notExistant) > 0):

            for nt in notExistant:
                print("File %s does not exists!" % (nt))

            raise IOError("One or more input file do not exist!")

        self.phafile = OGIPPHA(phafile, filetype='observed')
        self.exposure = self.phafile.getExposure()
        self.bkgfile = OGIPPHA(bkgfile, filetype="background")
        self.response = Response(rspfile)

        # Start with an empty mask (the user will overwrite it using the
        # setActiveMeasurement method)
        self.mask = np.asarray(
            np.ones(self.phafile.getRates().shape),
            np.bool)

        # Get the counts for this spectrum
        self.counts = (self.phafile.getRates()[self.mask]
                       * self.exposure)

        # Check that counts is positive
        idx = (self.counts < 0)

        if (np.sum(idx) > 0):
            warnings.warn("The observed spectrum for %s " % self.name +
                          "has negative channels! Fixing those to zero.",
                          RuntimeWarning)
            self.counts[idx] = 0

        # Get the background counts for this spectrum
        self.bkgCounts = (self.bkgfile.getRates()[self.mask]
                          * self.exposure)

        # Get the error on the background counts
        self.bkgErr = self.bkgfile.getRatesErrors()[self.mask] * self.exposure

        # Check that bkgCounts is positive
        idx = (self.bkgCounts < 0)

        if (np.sum(idx) > 0):
            warnings.warn("The background spectrum for %s " % self.name +
                          "has negative channels! Fixing those to zero.",
                          RuntimeWarning)
            self.bkgCounts[idx] = 0

        # Check that bkgErr is positive
        idx = (self.bkgCounts < 0)

        if (np.sum(idx) > 0):

            raise RuntimeError("The background spectrum for %s " % self.name +
                                "has negative errors! Fixing those to zero.")

        # Check that errors are zeros when the bkg is zero
        if np.any((self.bkgErr == 0) & (self.bkgCounts > 0)):

            raise RuntimeError("The background error is zero but the background counts are not zero for some bins. "
                               "Data are corrupted.")

        # Check that the observed counts are positive

        idx = self.counts < 0

        if np.sum(idx) > 0:
            raise RuntimeError("Negative counts in observed spectrum %s. Data are corrupted." % (phafile))

        # Keep a copy which will never be modified
        self.counts_backup = np.array(self.counts, copy=True)
        self.bkgCounts_backup = np.array(self.bkgCounts, copy=True)
        self.bkgErr_backup = np.array(self.bkgErr, copy=True)

        # Effective area correction is disabled by default, i.e.,
        # the nuisance parameter is fixed to 1
        self.nuisanceParameters = {}
        self.nuisanceParameters['InterCalib'] = Parameter("InterCalib", 1, min_value=0.9, max_value=1.1, delta=0.01)
        self.nuisanceParameters['InterCalib'].fix = True

    pass

    def use_intercalibration_constant(self, factorLowBound=0.9, factorHiBound=1.1):

        self.nuisanceParameters['InterCalib'].free()
        self.nuisanceParameters['InterCalib'].set_bounds(factorLowBound, factorHiBound)

        # Check that the parameter is within the provided bounds
        value = self.nuisanceParameters['InterCalib'].value

        if (value < factorLowBound):
            warnings.warn(
                "The intercalibration constant was %s, lower than the provided lower bound %s." % (value, factorLowBound) +
                " Setting it equal to the lower bound")

            self.nuisanceParameters['InterCalib'].setValue(float(factorLowBound))

        if (value > factorHiBound):
            warnings.warn(
                "The intercalibration constant was %s, larger than the provided hi bound %s." % (value, factorHiBound) +
                " Setting it equal to the hi bound")

            self.nuisanceParameters['InterCalib'].setValue(float(factorHiBound))

    def fix_intercalibration_constant(self, value=None):

        if (value is not None):
            # Fixing the constant to the provided value
            self.nuisanceParameters['InterCalib'].value = float(value)

        else:

            # Do nothing, i.e., leave the constant to the value
            # it currently has
            pass

        self.nuisanceParameters['InterCalib'].fix()

    def set_active_measurements(self, *args):
        '''Set the measurements to be used during the analysis.
        Use as many ranges as you need,
        specified as 'emin-emax'. Energies are in keV. Example:

        setActiveMeasurements('10-12.5','56.0-100.0')

        which will set the energy range 10-12.5 keV and 56-100 keV to be
        used in the analysis'''

        # To implelemnt this we will use an array of boolean index,
        # which will filter
        # out the non-used channels during the logLike

        # Now build the mask: values for which the mask is 0 will be masked
        mask = np.zeros(self.phafile.getRates().shape)

        for arg in args:
            ee = map(float, arg.replace(" ", "").split("-"))
            emin, emax = sorted(ee)
            idx1 = self.response.energy_to_channel(emin)
            idx2 = self.response.energy_to_channel(emax)
            mask[idx1:idx2 + 1] = True
        pass
        self.mask = np.array(mask, np.bool)

        self.counts = self.counts_backup[self.mask]
        self.bkgCounts = self.bkgCounts_backup[self.mask]
        self.bkgErr = self.bkgErr_backup[self.mask]

        print("Now using %s channels out of %s" % (np.sum(self.mask),
                                                   self.phafile.getRates().shape[0]
                                                   ))

    pass

    def get_name(self):
        '''
        Return a name for this dataset (likely set during the constructor)
        '''
        return self.name

    pass

    def set_model(self, likelihoodModel):
        '''
        Set the model to be used in the joint minimization.
        '''
        self.likelihoodModel = likelihoodModel

        nPointSources = self.likelihoodModel.get_number_of_point_sources()

        # This is a wrapper which iterates over all the point sources and get
        # the fluxes
        # We assume there are no extended sources, since the GBM cannot handle them

        def diffFlux(energies):
            fluxes = self.likelihoodModel.get_point_source_fluxes(0, energies)

            # If we have only one point source, this will never be executed
            for i in range(1, nPointSources):
                fluxes += self.likelihoodModel.get_point_source_fluxes(i, energies)

            return fluxes

        self.diffFlux = diffFlux

        # The following integrates the diffFlux function using Simpson's rule
        # This assume that the intervals e1,e2 are all small, which is guaranteed
        # for any reasonable response matrix, given that e1 and e2 are Monte-Carlo
        # energies. It also assumes that the function is smooth in the interval
        # e1 - e2 and twice-differentiable, again reasonable on small intervals for
        # decent models. It might fail for models with too sharp features, smaller
        # than the size of the monte carlo interval.

        def integral(e1, e2):
            # Simpson's rule

            return (e2 - e1) / 6.0 * (self.diffFlux(e1)
                                      + 4 * self.diffFlux((e1 + e2) / 2.0)
                                      + self.diffFlux(e2))

        self.response.set_function(diffFlux,
                                   integral)

    pass

    def inner_fit(self):

        return self.get_log_like()

    def get_folded_model(self):

        # Get the folded model for this spectrum
        # (this is the rate predicted, in cts/s)

        return self.response.convolve()[self.mask]

    def get_model_and_data(self):

        e1, e2 = (self.response.ebounds[:, 0],
                  self.response.ebounds[:, 1])

        return (self.response.convolve()[self.mask] * self.exposure
                + self.bkgCounts,
                e1[self.mask],
                e2[self.mask],
                self.counts)

    def _get_model_counts(self):

        # Get the folded model for this spectrum (this is the rate predicted,
        # in cts/s)

        folded = self.get_folded_model()

        # Model is folded+background (i.e., we assume negligible errors on the
        # background)
        modelCounts = self.nuisanceParameters['InterCalib'].value * folded * self.exposure

        # Put all negative predicted counts to zero
        clippedModelCounts = np.clip(modelCounts, 0, 1e10)

        return clippedModelCounts

    def _compute_log_like(self, modelCounts):

        # This loglike assume Gaussian errors on the background and Poisson uncertainties on the
        # observed counts. It is a profile likelihood.

        MB = self.bkgCounts + modelCounts
        s2 = self.bkgErr**2

        b = 0.5 * (np.sqrt( MB**2 - 2 * s2 * (MB - 2 * self.counts) + self.bkgErr**4)
                   + self.bkgCounts - modelCounts - s2)

        # Now there are two branches: when the background is 0 we are in the normal situation of a pure
        # Poisson likelihood, while when the background is not zero we use the profile likelihood

        # NOTE: In the constructor we enforced that bkgErr can be 0 only when also bkgCounts = 0
        # Also it is evident from the expression above that when bkgCounts = 0 and bkgErr=0 also b=0

        # Let's do the branch with background > 0 first

        idx = self.bkgCounts > 0

        log_likes = np.empty_like(modelCounts)

        log_likes[idx] = ( -(b[idx] - self.bkgCounts[idx])**2 / (2 * s2[idx])
                           + self.counts[idx] * np.log(b[idx] + modelCounts[idx])
                           - b[idx] - modelCounts[idx] - logfactorial(self.counts[idx])
                           - 0.5 * log(2*np.pi) - np.log(self.bkgErr[idx]))

        # Let's do the other branch

        nidx = ~idx

        # the 1e-100 in the log is to avoid zero divisions
        # This is the Poisson likelihood with no background
        log_likes[nidx] = self.counts[nidx] * np.log(modelCounts[nidx] + 1e-100) - modelCounts[nidx] \
                          - logfactorial(self.counts[nidx])

        return np.sum(log_likes)

    def get_log_like(self):
        '''
        Return the value of the log-likelihood with the current values for the
        parameters
        '''

        modelCounts = self._get_model_counts()

        logLike = self._compute_log_like(modelCounts)

        return logLike

    def get_nuisance_parameters(self):
        '''
        Return a list of nuisance parameter names. Return an empty list if there
        are no nuisance parameters
        '''
        return self.nuisanceParameters.keys()

class Response(object):

    def __init__(self, rspfile):

        rspNumber = 1

        # Check if we are dealing with a .rsp2 file (containing more than
        # one response). This is checked by looking for the syntax
        # [responseFile]{[responseNumber]}

        if ('{' in rspfile):
            tokens = rspfile.split("{")
            rspfile = tokens[0]
            rspNumber = int(tokens[-1].split('}')[0].replace(" ", ""))

        # Read the response
        with pyfits.open(rspfile) as f:

            try:

                # GBM typical response
                data = f['MATRIX', rspNumber].data

            except:
                # Other detectors might use the SPECRESP MATRIX name instead
                # Note that here we are not catching any exception, because
                # we have to fail if we cannot read the matrix

                data = f['SPECRESP MATRIX', rspNumber].data

            # Sometimes .rsp files contains a weird format featuring variable-length
            # arrays. Historically those confuse pyfits quite a lot, so we ensure
            # to transform them into standard numpy matrices to avoid issues

            self.matrix = variableToMatrix(data.field('MATRIX'))

            self.ebounds = np.vstack([f['EBOUNDS'].data.field("E_MIN"),
                                         f['EBOUNDS'].data.field("E_MAX")]).T

            self.mc_channels = np.vstack([data.field("ENERG_LO"),
                                             data.field("ENERG_HI")]).T

    def set_function(self, differentialFunction, integralFunction=None):
        '''
        Set the function to be used for the convolution
        '''

        self.differentialFunction = differentialFunction

        if (integralFunction == None):

            # This should never happen in 3ML, but we keep this functionality for
            # other uses. Print a warning anyway

            warnings.warn("Using the slow numerical integration in the GBM plugin!",
                          RuntimeWarning)

            def integral(x, y):
                return scipy.integrate.quad(self.differentialFunction, x, y)[0]

            # NB: vectorize is super slow!
            self.integralFunction = np.vectorize(integral, otypes=[np.float])

        else:

            self.integralFunction = integralFunction

    def convolve(self):

        trueFluxes = self.integralFunction(self.mc_channels[:, 0],
                                           self.mc_channels[:, 1])

        foldedCounts = np.dot(trueFluxes, self.matrix.T)

        return foldedCounts

    def get_counts_vector(self, e1, e2):

        trueFluxes = self.integralFunction(self.mc_channels[:, 0], self.mc_channels[:, 1])

    def energy_to_channel(self, energy):

        '''Finds the channel containing the provided energy.
        NOTE: returns the channel index (starting at zero),
        not the channel number (likely starting from 1)'''

        # Get the index of the first ebounds upper bound larger than energy

        try:

            idx = next(idx for idx,
                               value in enumerate(self.ebounds[:, 1])
                       if value >= energy)

        except StopIteration:

            # No values above the given energy, return the last channel
            return self.ebounds[:, 1].shape[0]

        return idx

def variableToMatrix(variableLengthMatrix):
    '''This take a variable length array and return it in a
    properly formed constant length array, to avoid some pyfits obscure bugs'''

    nrows = len(variableLengthMatrix)
    ncolumns = max([len(elem) for elem in variableLengthMatrix])
    matrix = np.zeros([ncolumns, nrows])

    for i in range(nrows):
        for j in range(ncolumns):

            try:

                matrix[j, i] = variableLengthMatrix[i][j]

            except:

                # This happens when the i,j element does not exist, which is not an error
                # We will just leave it to zero
                pass

    return matrix
<commit_msg>Added a method to view the count spectrum view_count_plot<commit_after>import collections
import os
import warnings

import astropy.io.fits as pyfits
import numpy as np
import scipy.integrate
from math import log

from threeML.minimizer import minimization
from threeML.plugin_prototype import PluginPrototype
from threeML.plugins.gammaln import logfactorial
from threeML.plugins.ogip import OGIPPHA
from threeML.io.gbm_channel_plot import gbm_channel_plot, excluded_channel_plot

from astromodels.parameter import Parameter

__instrument_name = "Fermi GBM (all detectors)"

class FermiGBMLike(PluginPrototype):

    def __init__(self, name, phafile, bkgfile, rspfile):
        '''
        If the input files are PHA2 files, remember to specify the spectrum number, for example:
        FermiGBMLike("GBM","spectrum.pha{2}","bkgfile.bkg{2}","rspfile.rsp{2}")
        to load the second spectrum, second background spectrum and second response.
        '''

        self.name = name

        # Check that all file exists
        notExistant = []

        if (not os.path.exists(phafile.split("{")[0])):
            notExistant.append(phafile.split("{")[0])

        if (not os.path.exists(bkgfile.split("{")[0])):
            notExistant.append(bkgfile.split("{")[0])

        if (not os.path.exists(rspfile.split("{")[0])):
            notExistant.append(rspfile.split("{")[0])

        if (len(notExistant) > 0):

            for nt in notExistant:
                print("File %s does not exists!" % (nt))

            raise IOError("One or more input file do not exist!")

        self.phafile = OGIPPHA(phafile, filetype='observed')
        self.exposure = self.phafile.getExposure()
        self.bkgfile = OGIPPHA(bkgfile, filetype="background")
        self.response = Response(rspfile)

        # Start with an empty mask (the user will overwrite it using the
        # setActiveMeasurement method)
        self.mask = np.asarray(
            np.ones(self.phafile.getRates().shape),
            np.bool)

        # Get the counts for this spectrum
        self.counts = (self.phafile.getRates()[self.mask]
                       * self.exposure)

        # Check that counts is positive
        idx = (self.counts < 0)

        if (np.sum(idx) > 0):
            warnings.warn("The observed spectrum for %s " % self.name +
                          "has negative channels! Fixing those to zero.",
                          RuntimeWarning)
            self.counts[idx] = 0

        # Get the background counts for this spectrum
        self.bkgCounts = (self.bkgfile.getRates()[self.mask]
                          * self.exposure)

        # Get the error on the background counts
        self.bkgErr = self.bkgfile.getRatesErrors()[self.mask] * self.exposure

        # Check that bkgCounts is positive
        idx = (self.bkgCounts < 0)

        if (np.sum(idx) > 0):
            warnings.warn("The background spectrum for %s " % self.name +
                          "has negative channels! Fixing those to zero.",
                          RuntimeWarning)
            self.bkgCounts[idx] = 0

        # Check that bkgErr is positive
        idx = (self.bkgCounts < 0)

        if (np.sum(idx) > 0):

            raise RuntimeError("The background spectrum for %s " % self.name +
                                "has negative errors! Fixing those to zero.")

        # Check that errors are zeros when the bkg is zero
        if np.any((self.bkgErr == 0) & (self.bkgCounts > 0)):

            raise RuntimeError("The background error is zero but the background counts are not zero for some bins. "
                               "Data are corrupted.")

        # Check that the observed counts are positive

        idx = self.counts < 0

        if np.sum(idx) > 0:
            raise RuntimeError("Negative counts in observed spectrum %s. Data are corrupted." % (phafile))

        # Keep a copy which will never be modified
        self.counts_backup = np.array(self.counts, copy=True)
        self.bkgCounts_backup = np.array(self.bkgCounts, copy=True)
        self.bkgErr_backup = np.array(self.bkgErr, copy=True)

        # Effective area correction is disabled by default, i.e.,
        # the nuisance parameter is fixed to 1
        self.nuisanceParameters = {}
        self.nuisanceParameters['InterCalib'] = Parameter("InterCalib", 1, min_value=0.9, max_value=1.1, delta=0.01)
        self.nuisanceParameters['InterCalib'].fix = True

    pass

    def use_intercalibration_constant(self, factorLowBound=0.9, factorHiBound=1.1):

        self.nuisanceParameters['InterCalib'].free()
        self.nuisanceParameters['InterCalib'].set_bounds(factorLowBound, factorHiBound)

        # Check that the parameter is within the provided bounds
        value = self.nuisanceParameters['InterCalib'].value

        if (value < factorLowBound):
            warnings.warn(
                "The intercalibration constant was %s, lower than the provided lower bound %s." % (value, factorLowBound) +
                " Setting it equal to the lower bound")

            self.nuisanceParameters['InterCalib'].setValue(float(factorLowBound))

        if (value > factorHiBound):
            warnings.warn(
                "The intercalibration constant was %s, larger than the provided hi bound %s." % (value, factorHiBound) +
                " Setting it equal to the hi bound")

            self.nuisanceParameters['InterCalib'].setValue(float(factorHiBound))

    def fix_intercalibration_constant(self, value=None):

        if (value is not None):
            # Fixing the constant to the provided value
            self.nuisanceParameters['InterCalib'].value = float(value)

        else:

            # Do nothing, i.e., leave the constant to the value
            # it currently has
            pass

        self.nuisanceParameters['InterCalib'].fix()

    def set_active_measurements(self, *args):
        '''Set the measurements to be used during the analysis.
        Use as many ranges as you need,
        specified as 'emin-emax'. Energies are in keV. Example:

        setActiveMeasurements('10-12.5','56.0-100.0')

        which will set the energy range 10-12.5 keV and 56-100 keV to be
        used in the analysis'''

        # To implelemnt this we will use an array of boolean index,
        # which will filter
        # out the non-used channels during the logLike

        # Now build the mask: values for which the mask is 0 will be masked
        mask = np.zeros(self.phafile.getRates().shape)

        for arg in args:
            ee = map(float, arg.replace(" ", "").split("-"))
            emin, emax = sorted(ee)
            idx1 = self.response.energy_to_channel(emin)
            idx2 = self.response.energy_to_channel(emax)
            mask[idx1:idx2 + 1] = True
        pass
        self.mask = np.array(mask, np.bool)

        self.counts = self.counts_backup[self.mask]
        self.bkgCounts = self.bkgCounts_backup[self.mask]
        self.bkgErr = self.bkgErr_backup[self.mask]

        print("Now using %s channels out of %s" % (np.sum(self.mask),
                                                   self.phafile.getRates().shape[0]
                                                   ))

    pass

    def get_name(self):
        '''
        Return a name for this dataset (likely set during the constructor)
        '''
        return self.name

    pass

    def set_model(self, likelihoodModel):
        '''
        Set the model to be used in the joint minimization.
        '''
        self.likelihoodModel = likelihoodModel

        nPointSources = self.likelihoodModel.get_number_of_point_sources()

        # This is a wrapper which iterates over all the point sources and get
        # the fluxes
        # We assume there are no extended sources, since the GBM cannot handle them

        def diffFlux(energies):
            fluxes = self.likelihoodModel.get_point_source_fluxes(0, energies)

            # If we have only one point source, this will never be executed
            for i in range(1, nPointSources):
                fluxes += self.likelihoodModel.get_point_source_fluxes(i, energies)

            return fluxes

        self.diffFlux = diffFlux

        # The following integrates the diffFlux function using Simpson's rule
        # This assume that the intervals e1,e2 are all small, which is guaranteed
        # for any reasonable response matrix, given that e1 and e2 are Monte-Carlo
        # energies. It also assumes that the function is smooth in the interval
        # e1 - e2 and twice-differentiable, again reasonable on small intervals for
        # decent models. It might fail for models with too sharp features, smaller
        # than the size of the monte carlo interval.

        def integral(e1, e2):
            # Simpson's rule

            return (e2 - e1) / 6.0 * (self.diffFlux(e1)
                                      + 4 * self.diffFlux((e1 + e2) / 2.0)
                                      + self.diffFlux(e2))

        self.response.set_function(diffFlux,
                                   integral)

    pass

    def inner_fit(self):

        return self.get_log_like()

    def get_folded_model(self):

        # Get the folded model for this spectrum
        # (this is the rate predicted, in cts/s)

        return self.response.convolve()[self.mask]

    def get_model_and_data(self):

        e1, e2 = (self.response.ebounds[:, 0],
                  self.response.ebounds[:, 1])

        return (self.response.convolve()[self.mask] * self.exposure
                + self.bkgCounts,
                e1[self.mask],
                e2[self.mask],
                self.counts)

    def _get_model_counts(self):

        # Get the folded model for this spectrum (this is the rate predicted,
        # in cts/s)

        folded = self.get_folded_model()

        # Model is folded+background (i.e., we assume negligible errors on the
        # background)
        modelCounts = self.nuisanceParameters['InterCalib'].value * folded * self.exposure

        # Put all negative predicted counts to zero
        clippedModelCounts = np.clip(modelCounts, 0, 1e10)

        return clippedModelCounts

    def _compute_log_like(self, modelCounts):

        # This loglike assume Gaussian errors on the background and Poisson uncertainties on the
        # observed counts. It is a profile likelihood.

        MB = self.bkgCounts + modelCounts
        s2 = self.bkgErr**2

        b = 0.5 * (np.sqrt( MB**2 - 2 * s2 * (MB - 2 * self.counts) + self.bkgErr**4)
                   + self.bkgCounts - modelCounts - s2)

        # Now there are two branches: when the background is 0 we are in the normal situation of a pure
        # Poisson likelihood, while when the background is not zero we use the profile likelihood

        # NOTE: In the constructor we enforced that bkgErr can be 0 only when also bkgCounts = 0
        # Also it is evident from the expression above that when bkgCounts = 0 and bkgErr=0 also b=0

        # Let's do the branch with background > 0 first

        idx = self.bkgCounts > 0

        log_likes = np.empty_like(modelCounts)

        log_likes[idx] = ( -(b[idx] - self.bkgCounts[idx])**2 / (2 * s2[idx])
                           + self.counts[idx] * np.log(b[idx] + modelCounts[idx])
                           - b[idx] - modelCounts[idx] - logfactorial(self.counts[idx])
                           - 0.5 * log(2*np.pi) - np.log(self.bkgErr[idx]))

        # Let's do the other branch

        nidx = ~idx

        # the 1e-100 in the log is to avoid zero divisions
        # This is the Poisson likelihood with no background
        log_likes[nidx] = self.counts[nidx] * np.log(modelCounts[nidx] + 1e-100) - modelCounts[nidx] \
                          - logfactorial(self.counts[nidx])

        return np.sum(log_likes)

    def get_log_like(self):
        '''
        Return the value of the log-likelihood with the current values for the
        parameters
        '''

        modelCounts = self._get_model_counts()

        logLike = self._compute_log_like(modelCounts)

        return logLike

    def get_nuisance_parameters(self):
        '''
        Return a list of nuisance parameter names. Return an empty list if there
        are no nuisance parameters
        '''
        return self.nuisanceParameters.keys()

    def view_count_spectrum(self):
        '''
        View the count and background spectrum. Useful to check energy selections.

        '''
        # First plot the counts
        _  = gbm_channel_plot(self.response.ebounds[:,0],self.response.ebounds[:,1],self.counts_backup,color='#377eb8',lw=2,alpha=1)
        ax = gbm_channel_plot(self.response.ebounds[:,0],self.response.ebounds[:,1],self.bkgCounts_backup,color='#e41a1c',alpha=.8)
        # Now fade the non-used channels
        excluded_channel_plot(self.response.ebounds[:,0],self.response.ebounds[:,1],self.mask,self.counts_backup,self.bkgCounts_backup,ax)

        ax.set_xlabel("Energy (keV)")
        ax.set_ylabel("Counts/keV")
        ax.set_xlim(left=self.response.ebounds[0,0],right=self.response.ebounds[-1,1])

class Response(object):

    def __init__(self, rspfile):

        rspNumber = 1

        # Check if we are dealing with a .rsp2 file (containing more than
        # one response). This is checked by looking for the syntax
        # [responseFile]{[responseNumber]}

        if ('{' in rspfile):
            tokens = rspfile.split("{")
            rspfile = tokens[0]
            rspNumber = int(tokens[-1].split('}')[0].replace(" ", ""))

        # Read the response
        with pyfits.open(rspfile) as f:

            try:

                # GBM typical response
                data = f['MATRIX', rspNumber].data

            except:
                # Other detectors might use the SPECRESP MATRIX name instead
                # Note that here we are not catching any exception, because
                # we have to fail if we cannot read the matrix

                data = f['SPECRESP MATRIX', rspNumber].data

            # Sometimes .rsp files contains a weird format featuring variable-length
            # arrays. Historically those confuse pyfits quite a lot, so we ensure
            # to transform them into standard numpy matrices to avoid issues

            self.matrix = variableToMatrix(data.field('MATRIX'))

            self.ebounds = np.vstack([f['EBOUNDS'].data.field("E_MIN"),
                                         f['EBOUNDS'].data.field("E_MAX")]).T

            self.mc_channels = np.vstack([data.field("ENERG_LO"),
                                             data.field("ENERG_HI")]).T

    def set_function(self, differentialFunction, integralFunction=None):
        '''
        Set the function to be used for the convolution
        '''

        self.differentialFunction = differentialFunction

        if (integralFunction == None):

            # This should never happen in 3ML, but we keep this functionality for
            # other uses. Print a warning anyway

            warnings.warn("Using the slow numerical integration in the GBM plugin!",
                          RuntimeWarning)

            def integral(x, y):
                return scipy.integrate.quad(self.differentialFunction, x, y)[0]

            # NB: vectorize is super slow!
            self.integralFunction = np.vectorize(integral, otypes=[np.float])

        else:

            self.integralFunction = integralFunction

    def convolve(self):

        trueFluxes = self.integralFunction(self.mc_channels[:, 0],
                                           self.mc_channels[:, 1])

        foldedCounts = np.dot(trueFluxes, self.matrix.T)

        return foldedCounts

    def get_counts_vector(self, e1, e2):

        trueFluxes = self.integralFunction(self.mc_channels[:, 0], self.mc_channels[:, 1])

    def energy_to_channel(self, energy):

        '''Finds the channel containing the provided energy.
        NOTE: returns the channel index (starting at zero),
        not the channel number (likely starting from 1)'''

        # Get the index of the first ebounds upper bound larger than energy

        try:

            idx = next(idx for idx,
                               value in enumerate(self.ebounds[:, 1])
                       if value >= energy)

        except StopIteration:

            # No values above the given energy, return the last channel
            return self.ebounds[:, 1].shape[0]

        return idx

def variableToMatrix(variableLengthMatrix):
    '''This take a variable length array and return it in a
    properly formed constant length array, to avoid some pyfits obscure bugs'''

    nrows = len(variableLengthMatrix)
    ncolumns = max([len(elem) for elem in variableLengthMatrix])
    matrix = np.zeros([ncolumns, nrows])

    for i in range(nrows):
        for j in range(ncolumns):

            try:

                matrix[j, i] = variableLengthMatrix[i][j]

            except:

                # This happens when the i,j element does not exist, which is not an error
                # We will just leave it to zero
                pass

    return matrix
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
# This file is part of PYBOSSA.
#
# Copyright (C) 2017 Scifabric LTD.
#
# PYBOSSA 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.
#
# PYBOSSA 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with PYBOSSA.  If not, see <http://www.gnu.org/licenses/>.

class AnnouncementAuth(object):
    _specific_actions = []

    def __init__(self):
        pass

    @property
    def specific_actions(self):
        return self._specific_actions

    def can(self, user, action, announcement=None):
        action = ''.join(['_', action])
        return getattr(self, action)(user, announcement)

    def _create(self, user, announcement=None):
        if not user.admin or (announcement is None):
            return False
        return True

    def _read(self, user, announcement=None):
        if announcement:
            if announcement is None:
                return False
        else:
            return True

    def _update(self, user, announcement):
        if not user.admin:
            return False
        return True

    def _delete(self, user, announcement):
        if not user.admin:
            return False
        return True
<commit_msg>remove auth pitfalls<commit_after># -*- coding: utf8 -*-
# This file is part of PYBOSSA.
#
# Copyright (C) 2017 Scifabric LTD.
#
# PYBOSSA 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.
#
# PYBOSSA 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 Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with PYBOSSA.  If not, see <http://www.gnu.org/licenses/>.

class AnnouncementAuth(object):
    _specific_actions = []

    def __init__(self):
        pass

    @property
    def specific_actions(self):
        return self._specific_actions

    def can(self, user, action, announcement=None):
        action = ''.join(['_', action])
        return getattr(self, action)(user, announcement)

    def _create(self, user, announcement=None):
        if user.is_anonymous() or (announcement is None):
            return False
        if not user.is_anonymous() and user.admin:
            return True
        return False

    def _read(self, user, announcement=None):
        if announcement is None:
            return False
        else:
            return True

    def _update(self, user, announcement):
        if user.is_anonymous() or (announcement is None):
            return False
        if not user.is_anonymous() and user.admin:
            return True
        return False

    def _delete(self, user, announcement):
        if user.is_anonymous() or (announcement is None):
            return False
        if not user.is_anonymous() and user.admin:
            return True
        return False
<|endoftext|>


<commit_before>"""
Reverse generated from CSS back to PY via pyg_css_convert.py
"""
from pygments.style import Style
from pygments.token import Keyword, Name, Comment, String, Error, Text, \
     Number, Operator, Generic, Whitespace, Punctuation, Other, Literal

class SolarizedStyle(Style):
    background_color = "#002b36"  # Background
    highlight_color = "#ffffcc"   # <-- Not defined; defaulted

    styles = {
        Comment:                     "#586e75",
        Comment.Multiline:           "#586e75",
        Comment.Preproc:             "#859900",
        Comment.Single:              "#586e75",
        Comment.Special:             "#859900",
        Error:                       "#93a1a1",
        Generic:                     "#93a1a1",
        Generic.Deleted:             "#2aa198",
        Generic.Emph:                "italic #93a1a1",
        Generic.Error:               "#dc322f",
        Generic.Heading:             "#cb4b16",
        Generic.Inserted:            "#859900",
        Generic.Output:              "#93a1a1",
        Generic.Prompt:              "#93a1a1",
        Generic.Strong:              "bold #93a1a1",
        Generic.Subheading:          "#cb4b16",
        Generic.Traceback:           "#93a1a1",
        Keyword:                     "#859900",
        Keyword.Constant:            "#cb4b16",
        Keyword.Declaration:         "#268bd2",
        Keyword.Namespace:           "#859900",
        Keyword.Pseudo:              "#859900",
        Keyword.Reserved:            "#268bd2",
        Keyword.Type:                "#dc322f",
        Literal:                     "#93a1a1",
        Literal.Date:                "#93a1a1",
        Literal.Number:              "#2aa198",
        Literal.Number.Float:        "#2aa198",
        Literal.Number.Hex:          "#2aa198",
        Literal.Number.Integer:      "#2aa198",
        Literal.Number.Integer.Long: "#2aa198",
        Literal.Number.Oct:          "#2aa198",
        Literal.String:              "#2aa198",
        Literal.String.Backtick:     "#586e75",
        Literal.String.Char:         "#2aa198",
        Literal.String.Doc:          "#93a1a1",
        Literal.String.Double:       "#2aa198",
        Literal.String.Escape:       "#cb4b16",
        Literal.String.Heredoc:      "#93a1a1",
        Literal.String.Interpol:     "#2aa198",
        Literal.String.Other:        "#2aa198",
        Literal.String.Regex:        "#dc322f",
        Literal.String.Single:       "#2aa198",
        Literal.String.Symbol:       "#2aa198",
        Name:                        "#93a1a1",
        Name.Attribute:              "#93a1a1",
        Name.Builtin:                "#b58900",
        Name.Builtin.Pseudo:         "#268bd2",
        Name.Class:                  "#268bd2",
        Name.Constant:               "#cb4b16",
        Name.Decorator:              "#268bd2",
        Name.Entity:                 "#cb4b16",
        Name.Exception:              "#cb4b16",
        Name.Function:               "#268bd2",
        Name.Label:                  "#93a1a1",
        Name.Namespace:              "#93a1a1",
        Name.Other:                  "#93a1a1",
        Name.Property:               "#93a1a1",
        Name.Tag:                    "#268bd2",
        Name.Variable:               "#268bd2",
        Name.Variable.Class:         "#268bd2",
        Name.Variable.Global:        "#268bd2",
        Name.Variable.Instance:      "#268bd2",
        Operator:                    "#859900",
        Operator.Word:               "#859900",
        Other:                       "#cb4b16",
        Punctuation:                 "#93a1a1",
        Text:                        "#93a1a1",  # Foreground
        Text.Whitespace:             "#93a1a1"

        # Below are classes that could not be resolved:

        # Below are invalid rules:
    }
<commit_msg>Remove accidental commit of solarized.py<commit_after><|endoftext|>


<commit_before>"""Implement a graph traversal using composition from the original Graph class."""
from graph import Graph

def dfs(graph, start):
    """Depth-first traversal method for Graph class."""
    visited = []
    visited.append(start)
    start = graph.nodes[start]
    for val in start:
        print(val)
        if val in visited:
            return
        visited.extend(dfs(graph, val))
    return visited
<commit_msg>refactored dfs function to include visited list<commit_after>"""Implement graph traversal with composition from the original Graph class."""
from graph import Graph

def dfs(graph, node):
    """Depth-first traversal method for Graph class."""
    visited = []

    def traverse(graph, node):
        layer = []
        visited.append(node)
        layer.append(node)
        node = graph.nodes[node]
        for val in node:
            if val not in visited:
                layer.extend(traverse(graph, val))
        return layer
    return traverse(graph, node)

if __name__ == '__main__':
    g = Graph()
    g.add_edge(1, 2)
    g.add_edge(2, 3)
    g.add_edge(2, 4)
    g.add_edge(4, 1)
    g.add_edge(4, 5)
    print(g.nodes)
    print(dfs(g, 1))
<|endoftext|>


<commit_before>#!/usr/bin/python
#
# Copyright (c) 2012 Mikkel Schubert <MSchubert@snm.ku.dk>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
import io
import os
import re
import subprocess
import collections
from itertools import \
    izip_longest

import pysam

from pypeline.node import \
    Node, \
    NodeError
from pypeline.common.fileutils import \
    describe_files, \
    make_dirs
from pypeline.common.utilities import \
    chain_sorted

import pypeline.common.formats.fastq as fastq
import pypeline.common.sampling as sampling
import pypeline.tools.factory as factory

class DetectInputDuplicationNode(Node):
    """Attempts to detect reads included multiple times as input based on the
    presence of reads with identical names AND sequences. This is compromise
    between sensitivity, specificity, and running time.

    A possible refinement would be to consider reads with the same name where
    one read is the prefix of the other (due to different amounts of trimming
    or collapsing of reads).
    """

    def __init__(self, input_files, output_file, dependencies=()):
        Node.__init__(self,
                      description="<Detect Input Duplication: %s>"
                      % (describe_files(input_files)),
                      input_files=input_files,
                      output_files=output_file,
                      dependencies=dependencies)

    def run(self, _):
        handles = []
        try:
            sequences = []
            for fpath in self.input_files:
                handle = pysam.Samfile(fpath)
                handles.append(handle)

                sequence = izip_longest(handle, (), fillvalue=fpath)
                sequences.append(sequence)

            position = 0
            records = chain_sorted(*sequences, key=self._key_by_tid_pos)
            observed_reads = collections.defaultdict(list)
            for (record, fpath) in records:
                if record.pos != position:
                    self._process_reads(observed_reads, self.output_files)
                    observed_reads.clear()
                    position = record.pos
                elif record.is_unmapped:
                    break

                # Ignore supplementary / secondary alignments
                if not record.flag & 0x900:
                    key = (record.qname, record.seq, record.qual)
                    observed_reads[key].append(fpath)
            self._process_reads(observed_reads, self.output_files)

            # Everything is ok, touch the output files
            for fpath in self.output_files:
                make_dirs(os.path.dirname(fpath))
                with open(fpath, "w"):
                    pass
        finally:
            for handle in handles:
                handle.close()

    @classmethod
    def _process_reads(cls, observed_reads, output_files):
        for ((name, _, _), fpaths) in observed_reads.iteritems():
            if len(fpaths) > 1:
                message = ["Read %r found in multiple files:" % (name,)]
                for fpath in fpaths:
                    message.append("  - %r" % (fpath,))
                message.append("")
                message.append("This indicates that the same data files have "
                               "been included multiple times in the project. "
                               "Please review the input files used in this "
                               "project, to ensure that each set of data is "
                               "included only once.\n\n"

                               "If this is not the case, then execute the "
                               "following command(s) to mark this test as having "
                               "succeeded:")

                for fpath in output_files:
                    message.append("$ touch '%s'" % (fpath,))

                raise NodeError("\n".join(message))

    @classmethod
    def _key_by_tid_pos(cls, record):
        return (record[0].tid, record[0].pos)

class ValidateFASTQFilesNode(Node):
    def __init__(self, input_files, output_file, offset, dependencies=()):
        self._offset = offset
        Node.__init__(self,
                      description="<Validate FASTQ Files: %s>"
                      % (describe_files(input_files)),
                      input_files=input_files,
                      output_files=output_file,
                      dependencies=dependencies)

    def _run(self, _config, _temp):
        check_fastq_files(self.input_files, self._offset, True)
        output_file = tuple(self.output_files)[0]
        make_dirs(os.path.dirname(output_file))
        with open(output_file, "w"):
            pass

class ValidateFASTAFilesNode(Node):
    def __init__(self, input_files, output_file, dependencies=()):
        Node.__init__(self,
                      description="<Validate FASTA Files: %s>"
                      % (describe_files(input_files)),
                      input_files=input_files,
                      output_files=output_file,
                      dependencies=dependencies)

        assert len(self.output_files) == 1, self.output_files

    def _run(self, _config, _temp):
        for filename in self.input_files:
            check_fasta_file(filename)
        output_file, = self.output_files
        make_dirs(os.path.dirname(output_file))
        with open(output_file, "w"):
            pass

def check_fastq_files(filenames, required_offset, allow_empty=False):
    for filename in filenames:
        qualities = _read_sequences(filename)
        offsets = fastq.classify_quality_strings(qualities)
        if offsets == fastq.OFFSET_BOTH:
            raise NodeError("FASTQ file contains quality scores with both "
                            "quality offsets (33 and 64); file may be "
                            "unexpected format or corrupt. Please ensure "
                            "that this file contains valid FASTQ reads from a "
                            "single source.\n    Filename = %r" % (filename,))
        elif offsets == fastq.OFFSET_MISSING:
            if allow_empty and not qualities:
                return

            raise NodeError("FASTQ file did not contain quality scores; file "
                            "may be unexpected format or corrupt. Ensure that "
                            "the file is a FASTQ file.\n    Filename = %r"
                            % (filename,))
        elif offsets not in (fastq.OFFSET_AMBIGIOUS, required_offset):
            raise NodeError("FASTQ file contains quality scores with wrong "
                            "quality score offset (%i); expected reads with "
                            "quality score offset %i. Ensure that the "
                            "'QualityOffset' specified in the makefile "
                            "corresponds to the input.\n    Filename = %s"
                            % (offsets, required_offset, filename))

def _read_sequences(filename):
    cat_call = factory.new("cat")
    cat_call.add_multiple_values((filename,))
    cat_call = cat_call.finalized_call

    cat = None
    try:
        cat = subprocess.Popen(cat_call,
                               bufsize=io.DEFAULT_BUFFER_SIZE,
                               stderr=subprocess.PIPE,
                               stdout=subprocess.PIPE)
        qualities = _collect_qualities(cat.stdout, filename)

        return sampling.reservoir_sampling(qualities, 100000)
    except:
        if cat:
            cat.kill()
            cat.wait()
            cat = None
        raise
    finally:
        rc_cat = cat.wait() if cat else 0
        if rc_cat:
            message = "Error running 'paleomix cat':\n" \
                      "  Unicat return-code = %i\n\n%s" \
                      % (rc_cat, cat.stderr.read())
            raise NodeError(message)

def _collect_qualities(handle, filename):
    header = handle.readline()
    while header:
        sequence = handle.readline()
        seperator = handle.readline()
        qualities = handle.readline()

        if not header.startswith("@"):
            if header.startswith(">"):
                raise NodeError("Input file appears to be in FASTA format "
                                "(header starts with '>', expected '@'), "
                                "but only FASTQ files are supported\n"
                                "Filename = %r" % (filename,))

            raise NodeError("Input file lacks FASTQ header (expected '@', "
                            "found %r), but only FASTQ files are supported\n"
                            "    Filename = %r" % (header[:1], filename))
        elif not qualities:
            raise NodeError("Partial record found; is not 4 lines long:\n"
                            "Filename = %r\n    Record = '%s'"
                            % (filename, header.rstrip()))
        elif not seperator.startswith("+"):
            raise NodeError("Input file lacks FASTQ seperator (expected '+', "
                            "found %r), but only FASTQ files are supported\n"
                            "    Filename = %r" % (seperator[:1], filename))
        elif len(sequence) != len(qualities):
            raise NodeError("Input file contains malformed FASTQ records; "
                            "length of sequence / qualities are not the "
                            "same.\n    Filename = %r\n    Record = '%s'"
                            % (filename, header.rstrip()))

        yield qualities
        header = handle.readline()

def check_fasta_file(filename):
    with open(filename) as handle:
        namecache = {}
        state, linelength, linelengthchanged = _NA, None, False
        for linenum, line in enumerate(handle, start=1):
            line = line.rstrip('\n\r')

            if not line:
                if state in (_NA, _IN_WHITESPACE):
                    continue
                elif state == _IN_HEADER:
                    raise NodeError("Expected FASTA sequence, found empty line"
                                    "\n    Filename = %r\n    Line = %r"
                                    % (filename, linenum))
                elif state == _IN_SEQUENCE:
                    state = _IN_WHITESPACE
                else:
                    assert False
            elif line.startswith(">"):
                if state in (_NA, _IN_SEQUENCE, _IN_WHITESPACE):
                    _validate_fasta_header(filename, linenum, line, namecache)
                    state = _IN_HEADER
                    linelength = None
                    linelengthchanged = False
                elif state == _IN_HEADER:
                    raise NodeError("Empty sequences not allowed\n"
                                    "    Filename = %r\n    Line = %r"
                                    % (filename, linenum - 1))
                else:
                    assert False
            else:
                if state == _NA:
                    raise NodeError("Expected FASTA header, found %r\n"
                                    "    Filename = %r\n    Line = %r"
                                    % (line, filename, linenum))
                elif state == _IN_HEADER:
                    _validate_fasta_line(filename, linenum, line)
                    linelength = len(line)
                    state = _IN_SEQUENCE
                elif state == _IN_SEQUENCE:
                    _validate_fasta_line(filename, linenum, line)
                    # If the length has changed, then that line must be the
                    # last line in the record, which may be shorter due to the
                    # sequence length. This is because the FAI index format
                    # expects that each line has the same length.
                    if linelengthchanged or (linelength < len(line)):
                        raise NodeError("Lines in FASTQ files must be of same "
                                        "length\n    Filename = %r\n"
                                        "    Line = %r" % (filename, linenum))
                    elif linelength != len(line):
                        linelengthchanged = True
                elif state == _IN_WHITESPACE:
                    raise NodeError("Empty lines not allowed in sequences\n"
                                    "    Filename = %r\n    Line = %r"
                                    % (filename, linenum))
                else:
                    assert False

        if state in (_NA, _IN_HEADER):
            raise NodeError("File does not contain any sequences"
                            "    Filename = %r" % (filename, ))

# Standard nucleotides + UIPAC codes
_VALID_CHARS_STR = "ACGTN" "RYSWKMBDHV"
_VALID_CHARS = frozenset(_VALID_CHARS_STR.upper() + _VALID_CHARS_STR.lower())
_NA, _IN_HEADER, _IN_SEQUENCE, _IN_WHITESPACE = range(4)

def _validate_fasta_header(filename, linenum, line, cache):
    name = line.split(" ", 1)[0][1:]
    if not name:
        raise NodeError("FASTA sequence must have non-empty name\n"
                        "    Filename = %r\n    Line = %r\n"
                        % (filename, linenum))
    elif not _RE_REF_NAME.match(name):
        raise NodeError("Invalid name for FASTA sequence: %r\n"
                        "    Filename = %r\n    Line = %r\n"
                        % (name, filename, linenum))
    elif name in cache:
        raise NodeError("FASTA sequences have identical name\n"
                        "    Filename = %r\n    Name = %r\n"
                        "    Line 1 = %r\n    Line 2 = %r\n"
                        % (filename, name, linenum, cache[name]))
    cache[name] = linenum
_RE_REF_NAME = re.compile("[!-()+-<>-~][!-~]*")

def _validate_fasta_line(filename, linenum, line):
    invalid_chars = frozenset(line) - _VALID_CHARS
    if invalid_chars:
        raise NodeError("FASTA sequence contains invalid characters\n"
                        "    Filename = %r\n    Line = %r\n"
                        "    Invalid characters = %r"
                        % (filename, linenum, "".join(invalid_chars)))
<commit_msg>Consider orientation when detecting putative duplicates.<commit_after>#!/usr/bin/python
#
# Copyright (c) 2012 Mikkel Schubert <MSchubert@snm.ku.dk>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
import io
import os
import re
import subprocess
import collections
from itertools import \
    izip_longest

import pysam

from pypeline.node import \
    Node, \
    NodeError
from pypeline.common.fileutils import \
    describe_files, \
    make_dirs
from pypeline.common.utilities import \
    chain_sorted

import pypeline.common.formats.fastq as fastq
import pypeline.common.sampling as sampling
import pypeline.tools.factory as factory

class DetectInputDuplicationNode(Node):
    """Attempts to detect reads included multiple times as input based on the
    presence of reads with identical names AND sequences. This is compromise
    between sensitivity, specificity, and running time.

    A possible refinement would be to consider reads with the same name where
    one read is the prefix of the other (due to different amounts of trimming
    or collapsing of reads).
    """

    def __init__(self, input_files, output_file, dependencies=()):
        Node.__init__(self,
                      description="<Detect Input Duplication: %s>"
                      % (describe_files(input_files)),
                      input_files=input_files,
                      output_files=output_file,
                      dependencies=dependencies)

    def run(self, _):
        handles = []
        try:
            sequences = []
            for fpath in self.input_files:
                handle = pysam.Samfile(fpath)
                handles.append(handle)

                sequence = izip_longest(handle, (), fillvalue=fpath)
                sequences.append(sequence)

            position = 0
            records = chain_sorted(*sequences, key=self._key_by_tid_pos)
            observed_reads = collections.defaultdict(list)
            for (record, fpath) in records:
                if record.pos != position:
                    self._process_reads(observed_reads, self.output_files)
                    observed_reads.clear()
                    position = record.pos
                elif record.is_unmapped:
                    break

                # Ignore supplementary / secondary alignments
                if not record.flag & 0x900:
                    key = (record.is_reverse, record.qname,
                           record.seq, record.qual)
                    observed_reads[key].append(fpath)
            self._process_reads(observed_reads, self.output_files)

            # Everything is ok, touch the output files
            for fpath in self.output_files:
                make_dirs(os.path.dirname(fpath))
                with open(fpath, "w"):
                    pass
        finally:
            for handle in handles:
                handle.close()

    @classmethod
    def _process_reads(cls, observed_reads, output_files):
        for ((_, name, _, _), fpaths) in observed_reads.iteritems():
            if len(fpaths) > 1:
                message = ["Read %r found in multiple files:" % (name,)]
                for fpath in fpaths:
                    message.append("  - %r" % (fpath,))
                message.append("")
                message.append("This indicates that the same data files have "
                               "been included multiple times in the project. "
                               "Please review the input files used in this "
                               "project, to ensure that each set of data is "
                               "included only once.\n\n"

                               "If this is not the case, then execute the "
                               "following command(s) to mark this test as having "
                               "succeeded:")

                for fpath in output_files:
                    message.append("$ touch '%s'" % (fpath,))

                raise NodeError("\n".join(message))

    @classmethod
    def _key_by_tid_pos(cls, record):
        return (record[0].tid, record[0].pos)

class ValidateFASTQFilesNode(Node):
    def __init__(self, input_files, output_file, offset, dependencies=()):
        self._offset = offset
        Node.__init__(self,
                      description="<Validate FASTQ Files: %s>"
                      % (describe_files(input_files)),
                      input_files=input_files,
                      output_files=output_file,
                      dependencies=dependencies)

    def _run(self, _config, _temp):
        check_fastq_files(self.input_files, self._offset, True)
        output_file = tuple(self.output_files)[0]
        make_dirs(os.path.dirname(output_file))
        with open(output_file, "w"):
            pass

class ValidateFASTAFilesNode(Node):
    def __init__(self, input_files, output_file, dependencies=()):
        Node.__init__(self,
                      description="<Validate FASTA Files: %s>"
                      % (describe_files(input_files)),
                      input_files=input_files,
                      output_files=output_file,
                      dependencies=dependencies)

        assert len(self.output_files) == 1, self.output_files

    def _run(self, _config, _temp):
        for filename in self.input_files:
            check_fasta_file(filename)
        output_file, = self.output_files
        make_dirs(os.path.dirname(output_file))
        with open(output_file, "w"):
            pass

def check_fastq_files(filenames, required_offset, allow_empty=False):
    for filename in filenames:
        qualities = _read_sequences(filename)
        offsets = fastq.classify_quality_strings(qualities)
        if offsets == fastq.OFFSET_BOTH:
            raise NodeError("FASTQ file contains quality scores with both "
                            "quality offsets (33 and 64); file may be "
                            "unexpected format or corrupt. Please ensure "
                            "that this file contains valid FASTQ reads from a "
                            "single source.\n    Filename = %r" % (filename,))
        elif offsets == fastq.OFFSET_MISSING:
            if allow_empty and not qualities:
                return

            raise NodeError("FASTQ file did not contain quality scores; file "
                            "may be unexpected format or corrupt. Ensure that "
                            "the file is a FASTQ file.\n    Filename = %r"
                            % (filename,))
        elif offsets not in (fastq.OFFSET_AMBIGIOUS, required_offset):
            raise NodeError("FASTQ file contains quality scores with wrong "
                            "quality score offset (%i); expected reads with "
                            "quality score offset %i. Ensure that the "
                            "'QualityOffset' specified in the makefile "
                            "corresponds to the input.\n    Filename = %s"
                            % (offsets, required_offset, filename))

def _read_sequences(filename):
    cat_call = factory.new("cat")
    cat_call.add_multiple_values((filename,))
    cat_call = cat_call.finalized_call

    cat = None
    try:
        cat = subprocess.Popen(cat_call,
                               bufsize=io.DEFAULT_BUFFER_SIZE,
                               stderr=subprocess.PIPE,
                               stdout=subprocess.PIPE)
        qualities = _collect_qualities(cat.stdout, filename)

        return sampling.reservoir_sampling(qualities, 100000)
    except:
        if cat:
            cat.kill()
            cat.wait()
            cat = None
        raise
    finally:
        rc_cat = cat.wait() if cat else 0
        if rc_cat:
            message = "Error running 'paleomix cat':\n" \
                      "  Unicat return-code = %i\n\n%s" \
                      % (rc_cat, cat.stderr.read())
            raise NodeError(message)

def _collect_qualities(handle, filename):
    header = handle.readline()
    while header:
        sequence = handle.readline()
        seperator = handle.readline()
        qualities = handle.readline()

        if not header.startswith("@"):
            if header.startswith(">"):
                raise NodeError("Input file appears to be in FASTA format "
                                "(header starts with '>', expected '@'), "
                                "but only FASTQ files are supported\n"
                                "Filename = %r" % (filename,))

            raise NodeError("Input file lacks FASTQ header (expected '@', "
                            "found %r), but only FASTQ files are supported\n"
                            "    Filename = %r" % (header[:1], filename))
        elif not qualities:
            raise NodeError("Partial record found; is not 4 lines long:\n"
                            "Filename = %r\n    Record = '%s'"
                            % (filename, header.rstrip()))
        elif not seperator.startswith("+"):
            raise NodeError("Input file lacks FASTQ seperator (expected '+', "
                            "found %r), but only FASTQ files are supported\n"
                            "    Filename = %r" % (seperator[:1], filename))
        elif len(sequence) != len(qualities):
            raise NodeError("Input file contains malformed FASTQ records; "
                            "length of sequence / qualities are not the "
                            "same.\n    Filename = %r\n    Record = '%s'"
                            % (filename, header.rstrip()))

        yield qualities
        header = handle.readline()

def check_fasta_file(filename):
    with open(filename) as handle:
        namecache = {}
        state, linelength, linelengthchanged = _NA, None, False
        for linenum, line in enumerate(handle, start=1):
            line = line.rstrip('\n\r')

            if not line:
                if state in (_NA, _IN_WHITESPACE):
                    continue
                elif state == _IN_HEADER:
                    raise NodeError("Expected FASTA sequence, found empty line"
                                    "\n    Filename = %r\n    Line = %r"
                                    % (filename, linenum))
                elif state == _IN_SEQUENCE:
                    state = _IN_WHITESPACE
                else:
                    assert False
            elif line.startswith(">"):
                if state in (_NA, _IN_SEQUENCE, _IN_WHITESPACE):
                    _validate_fasta_header(filename, linenum, line, namecache)
                    state = _IN_HEADER
                    linelength = None
                    linelengthchanged = False
                elif state == _IN_HEADER:
                    raise NodeError("Empty sequences not allowed\n"
                                    "    Filename = %r\n    Line = %r"
                                    % (filename, linenum - 1))
                else:
                    assert False
            else:
                if state == _NA:
                    raise NodeError("Expected FASTA header, found %r\n"
                                    "    Filename = %r\n    Line = %r"
                                    % (line, filename, linenum))
                elif state == _IN_HEADER:
                    _validate_fasta_line(filename, linenum, line)
                    linelength = len(line)
                    state = _IN_SEQUENCE
                elif state == _IN_SEQUENCE:
                    _validate_fasta_line(filename, linenum, line)
                    # If the length has changed, then that line must be the
                    # last line in the record, which may be shorter due to the
                    # sequence length. This is because the FAI index format
                    # expects that each line has the same length.
                    if linelengthchanged or (linelength < len(line)):
                        raise NodeError("Lines in FASTQ files must be of same "
                                        "length\n    Filename = %r\n"
                                        "    Line = %r" % (filename, linenum))
                    elif linelength != len(line):
                        linelengthchanged = True
                elif state == _IN_WHITESPACE:
                    raise NodeError("Empty lines not allowed in sequences\n"
                                    "    Filename = %r\n    Line = %r"
                                    % (filename, linenum))
                else:
                    assert False

        if state in (_NA, _IN_HEADER):
            raise NodeError("File does not contain any sequences"
                            "    Filename = %r" % (filename, ))

# Standard nucleotides + UIPAC codes
_VALID_CHARS_STR = "ACGTN" "RYSWKMBDHV"
_VALID_CHARS = frozenset(_VALID_CHARS_STR.upper() + _VALID_CHARS_STR.lower())
_NA, _IN_HEADER, _IN_SEQUENCE, _IN_WHITESPACE = range(4)

def _validate_fasta_header(filename, linenum, line, cache):
    name = line.split(" ", 1)[0][1:]
    if not name:
        raise NodeError("FASTA sequence must have non-empty name\n"
                        "    Filename = %r\n    Line = %r\n"
                        % (filename, linenum))
    elif not _RE_REF_NAME.match(name):
        raise NodeError("Invalid name for FASTA sequence: %r\n"
                        "    Filename = %r\n    Line = %r\n"
                        % (name, filename, linenum))
    elif name in cache:
        raise NodeError("FASTA sequences have identical name\n"
                        "    Filename = %r\n    Name = %r\n"
                        "    Line 1 = %r\n    Line 2 = %r\n"
                        % (filename, name, linenum, cache[name]))
    cache[name] = linenum
_RE_REF_NAME = re.compile("[!-()+-<>-~][!-~]*")

def _validate_fasta_line(filename, linenum, line):
    invalid_chars = frozenset(line) - _VALID_CHARS
    if invalid_chars:
        raise NodeError("FASTA sequence contains invalid characters\n"
                        "    Filename = %r\n    Line = %r\n"
                        "    Invalid characters = %r"
                        % (filename, linenum, "".join(invalid_chars)))
<|endoftext|>