<commit_before># Sam Greydanus and Luke Chang 2015
# Some code taken from nilearn searchlight implementation: https://github.com/nilearn/nilearn/blob/master/nilearn/decoding/searchlight.py

import os

import time
import sys
import warnings
from distutils.version import LooseVersion
import random

import cPickle
import numpy as np
import matplotlib.pyplot as plt
from nilearn import datasets
from nilearn import plotting

import nibabel as nib

import sklearn
from sklearn import neighbors
from sklearn.externals.joblib import Parallel, delayed, cpu_count
from sklearn import svm
from sklearn.cross_validation import cross_val_score
from sklearn.base import BaseEstimator
from sklearn import neighbors
from sklearn.svm import SVR

from nilearn import masking
from nilearn.input_data import NiftiMasker

from scipy.stats import multivariate_normal

from nltools.analysis import Predict
import glob

class Simulator:
    def __init__(self, brain_mask=None, output_dir = None): #no scoring param
        self.resource_folder = os.path.join(os.getcwd(),'resources')
        if output_dir is None:
            self.output_dir = os.path.join(os.getcwd())
        else:
            self.output_dir = output_dir

        if type(brain_mask) is str:
            brain_mask = nib.load(brain_mask)
        elif brain_mask is None:
            brain_mask = nib.load(os.path.join(self.resource_folder,'MNI152_T1_2mm_brain_mask_dil.nii.gz'))
        elif type(brain_mask) is not nib.nifti1.Nifti1Image:
            print(brain_mask)
            print(type(brain_mask))
            raise ValueError("brain_mask is not a string or a nibabel instance")
        self.brain_mask = brain_mask
        self.nifti_masker = NiftiMasker(mask_img=self.brain_mask)

    def gaussian(self, mu, sigma, i_tot):
        x, y, z = np.mgrid[0:self.brain_mask.shape[0], 0:self.brain_mask.shape[1], 0:self.brain_mask.shape[2]]

        # Need an (N, 3) array of (x, y) pairs.
        xyz = np.column_stack([x.flat, y.flat, z.flat])

        covariance = np.diag(sigma**2)
        g = multivariate_normal.pdf(xyz, mean=mu, cov=covariance)

        # Reshape back to a 3D grid.
        g = g.reshape(x.shape).astype(float)

        #select only the regions within the brain mask
        g = np.multiply(self.brain_mask.get_data(),g)
        #adjust total intensity of gaussian
        g = np.multiply(i_tot/np.sum(g),g)

        return g

    def sphere(self, r, p):
        dims = self.brain_mask.shape

        x, y, z = np.ogrid[-p[0]:dims[0]-p[0], -p[1]:dims[1]-p[1], -p[2]:dims[2]-p[2]]
        mask = x*x + y*y + z*z <= r*r

        activation = np.zeros(dims)
        activation[mask] = 1
        activation = np.multiply(activation, self.brain_mask.get_data())
        activation = nib.Nifti1Image(activation, affine=np.eye(4))

        #return the 3D numpy matrix of zeros containing the sphere as a region of ones
        return activation.get_data()

    def normal_noise(self, mu, sigma):
        vlength = np.sum(self.brain_mask.get_data())
        n = np.random.normal(mu, sigma, vlength)
        m = self.nifti_masker.inverse_transform(n)

        #return the 3D numpy matrix of zeros containing the brain mask filled with noise produced over a normal distribution
        return m.get_data()

    def constant_activation(self, c):
        vlength = np.sum(self.brain_mask.get_data())
        n = np.random.normal(mu, sigma, vlength)
        n = np.multiply(c, np.zeros_like(n))
        m = self.nifti_masker.inverse_transform(n)
        #return the 3D numpy matrix of zeros containing the brain mask filled a constant activation c
        return m.get_data()

    def to_nifti(self, m):
        if not (type(m) == np.ndarray and len(m.shape) == 3):
            raise ValueError("ERROR: need 3D np.ndarray matrix to create the nifti file")
        m = m.astype(np.float32)
        ni = nib.Nifti1Image(m, affine=np.eye(4))
        return ni

    def n_spheres(self, r, p_list):
        #initialize useful values
        dims = self.brain_mask.get_data().shape

        #generate and sum spheres of 0's and 1's
        A = np.zeros_like(self.brain_mask.get_data())
        for p in p_list:
            A = np.add(A, self.sphere(r, p))

        return A

    def collection_from_pattern(self, A, sigma, reps = 1, I = None, output_dir = None):
            if I is None:
                I = [sigma/10.0]

            levels = len(I)
            temp = I
            for i in xrange(reps - 1):
                I = I + temp

            #initialize useful values
            dims = self.brain_mask.get_data().shape

            #for each intensity
            A_list = []
            for i in I:
                A_list.append(np.multiply(A, i))

            #generate a different gaussian noise profile for each mask
            mu = 0 #values centered around 0
            N_list = []
            for i in xrange(len(I)):
                N_list.append(self.normal_noise(mu, sigma))

            #generate a constant activation mask
            c = 1.0 #values centered around 0
            C_list = []
            for i in xrange(len(I)):
                C_list.append(self.constant_activation(c))

            #add noise and signal together, then convert to nifti files
            NF_list = []
            for i in xrange(len(I)):
                NF_list.append(self.to_nifti(np.add(np.add(A_list[i],N_list[i]),C_list[i])))

            if output_dir is not None:
                if type(output_dir) is str:
                    for i in xrange(len(I)):
                        NF_list[i].to_filename(os.path.join(output_dir,'centered_sphere_' + str(i) + "_" + str(i%levels) + '.nii.gz'))
                else:
                    raise ValueError("ERROR. output_dir must be a string")

            return (NF_list, I)

    def collection_of_centered_spheres(self, r, sigma, reps = 1, I = None, output_dir = None):
        dims = self.brain_mask.get_data().shape
        p = [dims[0]/2, dims[1]/2, dims[2]/2]
        A = self.sphere(r, p)

        c = self.collection_from_pattern(A, sigma, reps = reps, I = I, output_dir = output_dir)

        return c

    # def getnifti(self, mu, sigma, i_tot=1):
    #     # coordinates of center of gaussian (center in brain mask)
    #     if mu == 'center':
    #         mu = np.array([self.brain_mask.shape[0]/2.0, self.brain_mask.shape[1]/2.0, self.brain_mask.shape[2]/2.0])

    #     g = self.gaussian(mu, sigma, i_tot)

    #     n = nib.Nifti1Image(g, affine=np.eye(4))
    #     n.to_filename(os.path.join(os.getcwd(),'data_3D.nii.gz'))

    #     return n
<commit_msg>Another simple bug fix<commit_after># Sam Greydanus and Luke Chang 2015
# Some code taken from nilearn searchlight implementation: https://github.com/nilearn/nilearn/blob/master/nilearn/decoding/searchlight.py

import os

import time
import sys
import warnings
from distutils.version import LooseVersion
import random

import cPickle
import numpy as np
import matplotlib.pyplot as plt
from nilearn import datasets
from nilearn import plotting

import nibabel as nib

import sklearn
from sklearn import neighbors
from sklearn.externals.joblib import Parallel, delayed, cpu_count
from sklearn import svm
from sklearn.cross_validation import cross_val_score
from sklearn.base import BaseEstimator
from sklearn import neighbors
from sklearn.svm import SVR

from nilearn import masking
from nilearn.input_data import NiftiMasker

from scipy.stats import multivariate_normal

from nltools.analysis import Predict
import glob

class Simulator:
    def __init__(self, brain_mask=None, output_dir = None): #no scoring param
        self.resource_folder = os.path.join(os.getcwd(),'resources')
        if output_dir is None:
            self.output_dir = os.path.join(os.getcwd())
        else:
            self.output_dir = output_dir

        if type(brain_mask) is str:
            brain_mask = nib.load(brain_mask)
        elif brain_mask is None:
            brain_mask = nib.load(os.path.join(self.resource_folder,'MNI152_T1_2mm_brain_mask_dil.nii.gz'))
        elif type(brain_mask) is not nib.nifti1.Nifti1Image:
            print(brain_mask)
            print(type(brain_mask))
            raise ValueError("brain_mask is not a string or a nibabel instance")
        self.brain_mask = brain_mask
        self.nifti_masker = NiftiMasker(mask_img=self.brain_mask)

    def gaussian(self, mu, sigma, i_tot):
        x, y, z = np.mgrid[0:self.brain_mask.shape[0], 0:self.brain_mask.shape[1], 0:self.brain_mask.shape[2]]

        # Need an (N, 3) array of (x, y) pairs.
        xyz = np.column_stack([x.flat, y.flat, z.flat])

        covariance = np.diag(sigma**2)
        g = multivariate_normal.pdf(xyz, mean=mu, cov=covariance)

        # Reshape back to a 3D grid.
        g = g.reshape(x.shape).astype(float)

        #select only the regions within the brain mask
        g = np.multiply(self.brain_mask.get_data(),g)
        #adjust total intensity of gaussian
        g = np.multiply(i_tot/np.sum(g),g)

        return g

    def sphere(self, r, p):
        dims = self.brain_mask.shape

        x, y, z = np.ogrid[-p[0]:dims[0]-p[0], -p[1]:dims[1]-p[1], -p[2]:dims[2]-p[2]]
        mask = x*x + y*y + z*z <= r*r

        activation = np.zeros(dims)
        activation[mask] = 1
        activation = np.multiply(activation, self.brain_mask.get_data())
        activation = nib.Nifti1Image(activation, affine=np.eye(4))

        #return the 3D numpy matrix of zeros containing the sphere as a region of ones
        return activation.get_data()

    def normal_noise(self, mu, sigma):
        vmask = self.nifti_masker.fit_transform(self.brain_mask)

        vlength = np.sum(self.brain_mask.get_data())
        n = np.random.normal(mu, sigma, vlength)
        m = self.nifti_masker.inverse_transform(n)

        #return the 3D numpy matrix of zeros containing the brain mask filled with noise produced over a normal distribution
        return m.get_data()

    def to_nifti(self, m):
        if not (type(m) == np.ndarray and len(m.shape) == 3):
            raise ValueError("ERROR: need 3D np.ndarray matrix to create the nifti file")
        m = m.astype(np.float32)
        ni = nib.Nifti1Image(m, affine=np.eye(4))
        return ni

    def n_spheres(self, r, p_list):
        #initialize useful values
        dims = self.brain_mask.get_data().shape

        #generate and sum spheres of 0's and 1's
        A = np.zeros_like(self.brain_mask.get_data())
        for p in p_list:
            A = np.add(A, self.sphere(r, p))

        return A

    def collection_from_pattern(self, A, sigma, reps = 1, I = None, output_dir = None):
            if I is None:
                I = [sigma/10.0]

            levels = len(I)
            temp = I
            for i in xrange(reps - 1):
                I = I + temp

            #initialize useful values
            dims = self.brain_mask.get_data().shape

            #for each intensity
            A_list = []
            for i in I:
                A_list.append(np.multiply(A, i))

            #generate a different gaussian noise profile for each mask
            mu = 0 #values centered around 0
            N_list = []
            for i in xrange(len(I)):
                N_list.append(self.normal_noise(mu, sigma))

            #add noise and signal together, then convert to nifti files
            NF_list = []
            for i in xrange(len(I)):
                c = 1.0
                const = self.brain_mask.get_data() * c
                NF_list.append(self.to_nifti(np.add( const,np.add(A_list[i],N_list[i])) ))

            if output_dir is not None:
                if type(output_dir) is str:
                    for i in xrange(len(I)):
                        NF_list[i].to_filename(os.path.join(output_dir,'centered_sphere_' + str(i) + "_" + str(i%levels) + '.nii.gz'))
                else:
                    raise ValueError("ERROR. output_dir must be a string")

            return (NF_list, I)

    def collection_of_centered_spheres(self, r, sigma, reps = 1, I = None, output_dir = None):
        dims = self.brain_mask.get_data().shape
        p = [dims[0]/2, dims[1]/2, dims[2]/2]
        A = self.sphere(r, p)

        c = self.collection_from_pattern(A, sigma, reps = reps, I = I, output_dir = output_dir)

        return c

    # def getnifti(self, mu, sigma, i_tot=1):
    #     # coordinates of center of gaussian (center in brain mask)
    #     if mu == 'center':
    #         mu = np.array([self.brain_mask.shape[0]/2.0, self.brain_mask.shape[1]/2.0, self.brain_mask.shape[2]/2.0])

    #     g = self.gaussian(mu, sigma, i_tot)

    #     n = nib.Nifti1Image(g, affine=np.eye(4))
    #     n.to_filename(os.path.join(os.getcwd(),'data_3D.nii.gz'))

    #     return n
<|endoftext|>


<commit_before>"""celery.worker"""
from carrot.connection import DjangoAMQPConnection
from celery.messaging import TaskConsumer
from celery.conf import DAEMON_CONCURRENCY, DAEMON_LOG_FILE
from celery.conf import QUEUE_WAKEUP_AFTER, EMPTY_MSG_EMIT_EVERY
from celery.log import setup_logger
from celery.registry import tasks
from celery.datastructures import TaskProcessQueue
from celery.models import PeriodicTaskMeta
from celery.backends import default_backend
from celery.timer import EventTimer
import multiprocessing
import simplejson
import traceback
import logging
import time

class EmptyQueue(Exception):
    """The message queue is currently empty."""

class UnknownTask(Exception):
    """Got an unknown task in the queue. The message is requeued and
    ignored."""

def jail(task_id, func, args, kwargs):
    """Wraps the task in a jail, which catches all exceptions, and
    saves the status and result of the task execution to the task
    meta backend.

    If the call was successful, it saves the result to the task result
    backend, and sets the task status to ``"DONE"``.

    If the call results in an exception, it saves the exception as the task
    result, and sets the task status to ``"FAILURE"``.

    :param task_id: The id of the task.

    :param func: Callable object to execute.

    :param args: List of positional args to pass on to the function.

    :param kwargs: Keyword arguments mapping to pass on to the function.

    :returns: the function return value on success.

    :returns: the exception instance on failure.

    """
    try:
        result = func(*args, **kwargs)
    except Exception, exc:
        default_backend.mark_as_failure(task_id, exc)
        return exc
    else:
        default_backend.mark_as_done(task_id, result)
        return result

class TaskWrapper(object):
    """Class wrapping a task to be run.

    :param task_name: see :attr:`task_name`.

    :param task_id: see :attr:`task_id`.

    :param task_func: see :attr:`task_func`

    :param args: see :attr:`args`

    :param kwargs: see :attr:`kwargs`.

    .. attribute:: task_name

        Kind of task. Must be a name registered in the task registry.

    .. attribute:: task_id

        UUID of the task.

    .. attribute:: task_func

        The tasks callable object.

    .. attribute:: args

        List of positional arguments to apply to the task.

    .. attribute:: kwargs

        Mapping of keyword arguments to apply to the task.

    """

    def __init__(self, task_name, task_id, task_func, args, kwargs):
        self.task_name = task_name
        self.task_id = task_id
        self.task_func = task_func
        self.args = args
        self.kwargs = kwargs

    @classmethod
    def from_message(cls, message):
        """Create a :class:`TaskWrapper` from a task message sent by
        :class:`celery.messaging.TaskPublisher`.

        :raises UnknownTask: if the message does not describe a task,
            the message is also rejected.

        :returns: :class:`TaskWrapper` instance.

        """
        message_data = simplejson.loads(message.body)
        task_name = message_data["task"]
        task_id = message_data["id"]
        args = message_data["args"]
        kwargs = message_data["kwargs"]
        if task_name not in tasks:
            message.reject()
            raise UnknownTask(task_name)
        task_func = tasks[task_name]
        return cls(task_name, task_id, task_func, args, kwargs)

    def extend_with_default_kwargs(self, loglevel, logfile):
        """Extend the tasks keyword arguments with standard task arguments.

        These are ``logfile``, ``loglevel``, ``task_id`` and ``task_name``.

        """
        task_func_kwargs = {"logfile": logfile,
                            "loglevel": loglevel,
                            "task_id": self.task_id,
                            "task_name": self.task_name}
        task_func_kwargs.update(self.kwargs)
        return task_func_kwargs

    def execute(self, loglevel=None, logfile=None):
        """Execute the task in a :func:`jail` and store return value
        and status in the task meta backend.

        :keyword loglevel: The loglevel used by the task.

        :keyword logfile: The logfile used by the task.

        """
        task_func_kwargs = self.extend_with_default_kwargs(loglevel, logfile)
        return jail(self.task_id, [
                        self.task_func, self.args, task_func_kwargs])

    def execute_using_pool(self, pool, loglevel=None, logfile=None):
        """Like :meth:`execute`, but using the :mod:`multiprocessing` pool.

        :param pool: A :class:`multiprocessing.Pool` instance.

        :keyword loglevel: The loglevel used by the task.

        :keyword logfile: The logfile used by the task.

        :returns :class:`multiprocessing.AsyncResult` instance.

        """
        task_func_kwargs = self.extend_kwargs_with_logging(loglevel, logfile)
        return pool.apply_async(jail, [self.task_id, self.task_func,
                                       self.args, task_func_kwargs])

class TaskDaemon(object):
    """Executes tasks waiting in the task queue.

    :param concurrency: see :attr:`concurrency`.

    :param logfile: see :attr:`logfile`.

    :param loglevel: see :attr:`loglevel`.

    :param queue_wakeup_after: see :attr:`queue_wakeup_after`.

    .. attribute:: concurrency

        The number of simultaneous processes doing work (default:
        :const:`celery.conf.DAEMON_CONCURRENCY`)

    .. attribute:: loglevel

        The loglevel used (default: :const:`logging.INFO`)

    .. attribute:: logfile

        The logfile used, if no logfile is specified it uses ``stderr``
        (default: :const:`celery.conf.DAEMON_LOG_FILE`).

    .. attribute:: queue_wakeup_after

        The time it takes for the daemon to wake up after the queue is empty,
        so it can check for more work
        (default: :const:`celery.conf.QUEUE_WAKEUP_AFTER`).

    .. attribute:: empty_msg_emit_every

        How often the daemon emits the ``"Waiting for queue..."`` message.
        If this is ``None``, the message will never be logged.
        (default: :const:`celery.conf.EMPTY_MSG_EMIT_EVERY`)

    .. attribute:: logger

        The :class:`logging.Logger` instance used for logging.

    .. attribute:: pool

        The :class:`multiprocessing.Pool` instance used.

    .. attribute:: task_consumer

        The :class:`celery.messaging.TaskConsumer` instance used.

    """
    loglevel = logging.ERROR
    concurrency = DAEMON_CONCURRENCY
    logfile = DAEMON_LOG_FILE
    queue_wakeup_after = QUEUE_WAKEUP_AFTER
    empty_msg_emit_every = EMPTY_MSG_EMIT_EVERY

    def __init__(self, concurrency=None, logfile=None, loglevel=None,
            queue_wakeup_after=None):
        self.loglevel = loglevel or self.loglevel
        self.concurrency = concurrency or self.concurrency
        self.logfile = logfile or self.logfile
        self.queue_wakeup_after = queue_wakeup_after or \
                                    self.queue_wakeup_after
        self.logger = setup_logger(loglevel, logfile)
        self.pool = multiprocessing.Pool(self.concurrency)
        self.task_consumer = None
        self.reset_connection()

    def reset_connection(self):
        """Reset the AMQP connection, and reinitialize the
        :class:`celery.messaging.TaskConsumer` instance.

        Resets the task consumer in :attr:`task_consumer`.

        """
        if self.task_consumer:
            self.task_consumer.connection.close()
        amqp_connection = DjangoAMQPConnection()
        self.task_consumer = TaskConsumer(connection=amqp_connection)

    def connection_diagnostics(self):
        """Diagnose the AMQP connection, and reset connection if
        necessary."""
        if not self.task_consumer.channel.connection:
            self.logger.info(
                    "AMQP Connection has died, restoring connection.")
            self.reset_connection()

    def receive_message(self):
        """Receive the next message from the message broker.

        Tries to reset the AMQP connection if not available.
        Returns ``None`` if no message is waiting on the queue.

        :rtype: :class:`carrot.messaging.Message` instance.

        """
        self.connection_diagnostics()
        message = self.task_consumer.fetch()
        if message is not None:
            message.ack()
        return message

    def fetch_next_task(self):
        """Fetch the next task from the AMQP broker.

        Raises :exc:`EmptyQueue` exception if there is no message
        waiting on the queue.

        :returns: :class:`TaskWrapper` instance.

        """
        message = self.receive_message()
        if message is None: # No messages waiting.
            raise EmptyQueue()

        task = TaskWrapper.from_message(message)
        self.logger.info("Got task from broker: %s[%s]" % (
                            task.task_name, task.task_id))

        return task, message

    def execute_next_task(self):
        """Execute the next task on the queue using the multiprocessing pool.

        Catches all exceptions and logs them with level
        :const:`logging.CRITICAL`.

        """
        task, message = self.fetch_next_task()

        try:
            result = task.execute_using_pool(self.pool, self.loglevel,
                                             self.logfile)
        except Exception, error:
            self.logger.critical("Worker got exception %s: %s\n%s" % (
                error.__class__, error, traceback.format_exc()))
            return

        return result, task.task_name, task.task_id

    def run_periodic_tasks(self):
        """Schedule all waiting periodic tasks for execution.

        :rtype: list of :class:`celery.models.PeriodicTaskMeta` objects.
        """
        waiting_tasks = PeriodicTaskMeta.objects.get_waiting_tasks()
        [waiting_task.delay()
                for waiting_task in waiting_tasks]
        return waiting_tasks

    def schedule_retry_tasks(self):
        """Reschedule all requeued tasks waiting for retry."""
        pass

    def run(self):
        """Starts the workers main loop."""
        results = TaskProcessQueue(self.concurrency, logger=self.logger,
                done_msg="Task %(name)s[%(id)s] processed: %(return_value)s")
        log_wait = lambda: self.logger.info("Waiting for queue...")
        ev_msg_waiting = EventTimer(log_wait, self.empty_msg_emit_every)
        events = [
            EventTimer(self.run_periodic_tasks, 1),
            EventTimer(self.schedule_retry_tasks, 2),
        ]

        while True:
            [event.tick() for event in events]
            try:
                result, task_name, task_id = self.execute_next_task()
            except ValueError:
                # execute_next_task didn't return a r/name/id tuple,
                # probably because it got an exception.
                continue
            except EmptyQueue:
                ev_msg_waiting.tick()
                time.sleep(self.queue_wakeup_after)
                continue
            except UnknownTask, e:
                self.logger.info("Unknown task ignored: %s" % (e))
                continue
            except Exception, e:
                self.logger.critical("Message queue raised %s: %s\n%s" % (
                             e.__class__, e, traceback.format_exc()))
                continue

            results.add(result, task_name, task_id)
<commit_msg>connection_diagnostics now supports carrot versions >= 0.4.0<commit_after>"""celery.worker"""
from carrot.connection import DjangoAMQPConnection
from celery.messaging import TaskConsumer
from celery.conf import DAEMON_CONCURRENCY, DAEMON_LOG_FILE
from celery.conf import QUEUE_WAKEUP_AFTER, EMPTY_MSG_EMIT_EVERY
from celery.log import setup_logger
from celery.registry import tasks
from celery.datastructures import TaskProcessQueue
from celery.models import PeriodicTaskMeta
from celery.backends import default_backend
from celery.timer import EventTimer
import multiprocessing
import simplejson
import traceback
import logging
import time

class EmptyQueue(Exception):
    """The message queue is currently empty."""

class UnknownTask(Exception):
    """Got an unknown task in the queue. The message is requeued and
    ignored."""

def jail(task_id, func, args, kwargs):
    """Wraps the task in a jail, which catches all exceptions, and
    saves the status and result of the task execution to the task
    meta backend.

    If the call was successful, it saves the result to the task result
    backend, and sets the task status to ``"DONE"``.

    If the call results in an exception, it saves the exception as the task
    result, and sets the task status to ``"FAILURE"``.

    :param task_id: The id of the task.

    :param func: Callable object to execute.

    :param args: List of positional args to pass on to the function.

    :param kwargs: Keyword arguments mapping to pass on to the function.

    :returns: the function return value on success.

    :returns: the exception instance on failure.

    """
    try:
        result = func(*args, **kwargs)
    except Exception, exc:
        default_backend.mark_as_failure(task_id, exc)
        return exc
    else:
        default_backend.mark_as_done(task_id, result)
        return result

class TaskWrapper(object):
    """Class wrapping a task to be run.

    :param task_name: see :attr:`task_name`.

    :param task_id: see :attr:`task_id`.

    :param task_func: see :attr:`task_func`

    :param args: see :attr:`args`

    :param kwargs: see :attr:`kwargs`.

    .. attribute:: task_name

        Kind of task. Must be a name registered in the task registry.

    .. attribute:: task_id

        UUID of the task.

    .. attribute:: task_func

        The tasks callable object.

    .. attribute:: args

        List of positional arguments to apply to the task.

    .. attribute:: kwargs

        Mapping of keyword arguments to apply to the task.

    """

    def __init__(self, task_name, task_id, task_func, args, kwargs):
        self.task_name = task_name
        self.task_id = task_id
        self.task_func = task_func
        self.args = args
        self.kwargs = kwargs

    @classmethod
    def from_message(cls, message):
        """Create a :class:`TaskWrapper` from a task message sent by
        :class:`celery.messaging.TaskPublisher`.

        :raises UnknownTask: if the message does not describe a task,
            the message is also rejected.

        :returns: :class:`TaskWrapper` instance.

        """
        message_data = simplejson.loads(message.body)
        task_name = message_data["task"]
        task_id = message_data["id"]
        args = message_data["args"]
        kwargs = message_data["kwargs"]
        if task_name not in tasks:
            message.reject()
            raise UnknownTask(task_name)
        task_func = tasks[task_name]
        return cls(task_name, task_id, task_func, args, kwargs)

    def extend_with_default_kwargs(self, loglevel, logfile):
        """Extend the tasks keyword arguments with standard task arguments.

        These are ``logfile``, ``loglevel``, ``task_id`` and ``task_name``.

        """
        task_func_kwargs = {"logfile": logfile,
                            "loglevel": loglevel,
                            "task_id": self.task_id,
                            "task_name": self.task_name}
        task_func_kwargs.update(self.kwargs)
        return task_func_kwargs

    def execute(self, loglevel=None, logfile=None):
        """Execute the task in a :func:`jail` and store return value
        and status in the task meta backend.

        :keyword loglevel: The loglevel used by the task.

        :keyword logfile: The logfile used by the task.

        """
        task_func_kwargs = self.extend_with_default_kwargs(loglevel, logfile)
        return jail(self.task_id, [
                        self.task_func, self.args, task_func_kwargs])

    def execute_using_pool(self, pool, loglevel=None, logfile=None):
        """Like :meth:`execute`, but using the :mod:`multiprocessing` pool.

        :param pool: A :class:`multiprocessing.Pool` instance.

        :keyword loglevel: The loglevel used by the task.

        :keyword logfile: The logfile used by the task.

        :returns :class:`multiprocessing.AsyncResult` instance.

        """
        task_func_kwargs = self.extend_kwargs_with_logging(loglevel, logfile)
        return pool.apply_async(jail, [self.task_id, self.task_func,
                                       self.args, task_func_kwargs])

class TaskDaemon(object):
    """Executes tasks waiting in the task queue.

    :param concurrency: see :attr:`concurrency`.

    :param logfile: see :attr:`logfile`.

    :param loglevel: see :attr:`loglevel`.

    :param queue_wakeup_after: see :attr:`queue_wakeup_after`.

    .. attribute:: concurrency

        The number of simultaneous processes doing work (default:
        :const:`celery.conf.DAEMON_CONCURRENCY`)

    .. attribute:: loglevel

        The loglevel used (default: :const:`logging.INFO`)

    .. attribute:: logfile

        The logfile used, if no logfile is specified it uses ``stderr``
        (default: :const:`celery.conf.DAEMON_LOG_FILE`).

    .. attribute:: queue_wakeup_after

        The time it takes for the daemon to wake up after the queue is empty,
        so it can check for more work
        (default: :const:`celery.conf.QUEUE_WAKEUP_AFTER`).

    .. attribute:: empty_msg_emit_every

        How often the daemon emits the ``"Waiting for queue..."`` message.
        If this is ``None``, the message will never be logged.
        (default: :const:`celery.conf.EMPTY_MSG_EMIT_EVERY`)

    .. attribute:: logger

        The :class:`logging.Logger` instance used for logging.

    .. attribute:: pool

        The :class:`multiprocessing.Pool` instance used.

    .. attribute:: task_consumer

        The :class:`celery.messaging.TaskConsumer` instance used.

    """
    loglevel = logging.ERROR
    concurrency = DAEMON_CONCURRENCY
    logfile = DAEMON_LOG_FILE
    queue_wakeup_after = QUEUE_WAKEUP_AFTER
    empty_msg_emit_every = EMPTY_MSG_EMIT_EVERY

    def __init__(self, concurrency=None, logfile=None, loglevel=None,
            queue_wakeup_after=None):
        self.loglevel = loglevel or self.loglevel
        self.concurrency = concurrency or self.concurrency
        self.logfile = logfile or self.logfile
        self.queue_wakeup_after = queue_wakeup_after or \
                                    self.queue_wakeup_after
        self.logger = setup_logger(loglevel, logfile)
        self.pool = multiprocessing.Pool(self.concurrency)
        self.task_consumer = None
        self.reset_connection()

    def reset_connection(self):
        """Reset the AMQP connection, and reinitialize the
        :class:`celery.messaging.TaskConsumer` instance.

        Resets the task consumer in :attr:`task_consumer`.

        """
        if self.task_consumer:
            self.task_consumer.connection.close()
        amqp_connection = DjangoAMQPConnection()
        self.task_consumer = TaskConsumer(connection=amqp_connection)

    def connection_diagnostics(self):
        """Diagnose the AMQP connection, and reset connection if
        necessary."""
        if hasattr(self.task_consumer.backend):
            connection = self.task_consumer.backend.channel.connection
        else:
            connection = self.task_consumer.channel.connection

        if not connection:
            self.logger.info(
                    "AMQP Connection has died, restoring connection.")
            self.reset_connection()

    def receive_message(self):
        """Receive the next message from the message broker.

        Tries to reset the AMQP connection if not available.
        Returns ``None`` if no message is waiting on the queue.

        :rtype: :class:`carrot.messaging.Message` instance.

        """
        self.connection_diagnostics()
        message = self.task_consumer.fetch()
        if message is not None:
            message.ack()
        return message

    def fetch_next_task(self):
        """Fetch the next task from the AMQP broker.

        Raises :exc:`EmptyQueue` exception if there is no message
        waiting on the queue.

        :returns: :class:`TaskWrapper` instance.

        """
        message = self.receive_message()
        if message is None: # No messages waiting.
            raise EmptyQueue()

        task = TaskWrapper.from_message(message)
        self.logger.info("Got task from broker: %s[%s]" % (
                            task.task_name, task.task_id))

        return task, message

    def execute_next_task(self):
        """Execute the next task on the queue using the multiprocessing pool.

        Catches all exceptions and logs them with level
        :const:`logging.CRITICAL`.

        """
        task, message = self.fetch_next_task()

        try:
            result = task.execute_using_pool(self.pool, self.loglevel,
                                             self.logfile)
        except Exception, error:
            self.logger.critical("Worker got exception %s: %s\n%s" % (
                error.__class__, error, traceback.format_exc()))
            return

        return result, task.task_name, task.task_id

    def run_periodic_tasks(self):
        """Schedule all waiting periodic tasks for execution.

        :rtype: list of :class:`celery.models.PeriodicTaskMeta` objects.
        """
        waiting_tasks = PeriodicTaskMeta.objects.get_waiting_tasks()
        [waiting_task.delay()
                for waiting_task in waiting_tasks]
        return waiting_tasks

    def schedule_retry_tasks(self):
        """Reschedule all requeued tasks waiting for retry."""
        pass

    def run(self):
        """Starts the workers main loop."""
        results = TaskProcessQueue(self.concurrency, logger=self.logger,
                done_msg="Task %(name)s[%(id)s] processed: %(return_value)s")
        log_wait = lambda: self.logger.info("Waiting for queue...")
        ev_msg_waiting = EventTimer(log_wait, self.empty_msg_emit_every)
        events = [
            EventTimer(self.run_periodic_tasks, 1),
            EventTimer(self.schedule_retry_tasks, 2),
        ]

        while True:
            [event.tick() for event in events]
            try:
                result, task_name, task_id = self.execute_next_task()
            except ValueError:
                # execute_next_task didn't return a r/name/id tuple,
                # probably because it got an exception.
                continue
            except EmptyQueue:
                ev_msg_waiting.tick()
                time.sleep(self.queue_wakeup_after)
                continue
            except UnknownTask, e:
                self.logger.info("Unknown task ignored: %s" % (e))
                continue
            except Exception, e:
                self.logger.critical("Message queue raised %s: %s\n%s" % (
                             e.__class__, e, traceback.format_exc()))
                continue

            results.add(result, task_name, task_id)
<|endoftext|>


<commit_before>from tastypie.authentication import Authentication
from tastypie.authorization import DjangoAuthorization
from tastypie.resources import ModelResource, ALL, ALL_WITH_RELATIONS

from hyquest.models import QuestSet, Quest, Task, UserTaskAction, UserQuestAction, UserQuestSetAction

class FeaturedQuestSetResource(ModelResource):
    class Meta:
        queryset = QuestSet.objects.filter(featured=True, active=True)
	resource_name = 'featured'
        # filtering = {"name": ALL}
    def dehydrate(self, bundle):
        try:
            ua = UserQuestSetAction.objects.get(questset=bundle.obj)
            bundle.data['complete'] = ua.complete
        except Exception:
            bundle.data['complete'] = False
        return bundle

class ActiveQuestSetResource(ModelResource):
    class Meta:
        resource_name = 'active'

    def get_object_list(self, request):
        return QuestSet.objects.filter(userquestsetaction__user=request.user, userquestsetaction__complete=False, active=True)

class CompleteQuestSetResource(ModelResource):
    class Meta:
        resource_name = 'complete'

    def get_object_list(self, request):
        return QuestSet.objects.filter(userquestsetaction__user=request.user, userquestsetaction__complete=True, active=True)

class QuestResource(ModelResource):
    class Meta:
        queryset = Quest.objects.all()
	resource_name = 'quest'
        filtering = {"quest_set": ALL}
    def dehydrate(self, bundle):
        try:
            ua = UserQuestAction.objects.get(quest=bundle.obj)
            bundle.data['complete'] = ua.complete
        except Exception:
            bundle.data['complete'] = False
        return bundle

class TaskResource(ModelResource):
    class Meta:
        queryset = Task.objects.all()
	resource_name = 'task'
        filtering = {"quest": ALL}
        fields = ['title', 'type', 'points']

    def dehydrate(self, bundle):
        try:
            ua = UserTaskAction.objects.get(task=bundle.obj)
            bundle.data['complete'] = ua.complete
        except Exception:
            bundle.data['complete'] = False
        return bundle
<commit_msg>masked private information about the API<commit_after>from tastypie.authentication import Authentication
from tastypie.authorization import DjangoAuthorization
from tastypie.resources import ModelResource, ALL, ALL_WITH_RELATIONS

from hyquest.models import QuestSet, Quest, Task, UserTaskAction, UserQuestAction, UserQuestSetAction

class FeaturedQuestSetResource(ModelResource):
    class Meta:
        queryset = QuestSet.objects.filter(featured=True, active=True)
	resource_name = 'featured'
        fields = ['title', 'description', 'points']
    def dehydrate(self, bundle):
        try:
            ua = UserQuestSetAction.objects.get(questset=bundle.obj)
            bundle.data['complete'] = ua.complete
        except Exception:
            bundle.data['complete'] = False
        return bundle

class ActiveQuestSetResource(ModelResource):
    class Meta:
        resource_name = 'active'
        fields = ['title', 'description', 'points']
    def get_object_list(self, request):
        return QuestSet.objects.filter(userquestsetaction__user=request.user, userquestsetaction__complete=False, active=True)

class CompleteQuestSetResource(ModelResource):
    class Meta:
        resource_name = 'complete'
        fields = ['title', 'description', 'points']
    def get_object_list(self, request):
        return QuestSet.objects.filter(userquestsetaction__user=request.user, userquestsetaction__complete=True, active=True)

class QuestResource(ModelResource):
    class Meta:
        queryset = Quest.objects.all()
	resource_name = 'quest'
        filtering = {"quest_set": ALL}
        fields = ['title', 'points']
    def dehydrate(self, bundle):
        try:
            ua = UserQuestAction.objects.get(quest=bundle.obj)
            bundle.data['complete'] = ua.complete
        except Exception:
            bundle.data['complete'] = False
        return bundle

class TaskResource(ModelResource):
    class Meta:
        queryset = Task.objects.all()
	resource_name = 'task'
        filtering = {"quest": ALL}
        fields = ['title', 'type', 'points']

    def dehydrate(self, bundle):
        try:
            ua = UserTaskAction.objects.get(task=bundle.obj)
            bundle.data['complete'] = ua.complete
        except Exception:
            bundle.data['complete'] = False
        return bundle
<|endoftext|>


<commit_before># Natural Language Toolkit: Corpus Readers
#
# Copyright (C) 2001-2014 NLTK Project
# Author: Steven Bird <stevenbird1@gmail.com>
#         Edward Loper <edloper@gmail.com>
# URL: <http://nltk.org/>
# For license information, see LICENSE.TXT

"""
NLTK corpus readers.  The modules in this package provide functions
that can be used to read corpus fileids in a variety of formats.  These
functions can be used to read both the corpus fileids that are
distributed in the NLTK corpus package, and corpus fileids that are part
of external corpora.

Corpus Reader Functions
=======================
Each corpus module defines one or more "corpus reader functions",
which can be used to read documents from that corpus.  These functions
take an argument, ``item``, which is used to indicate which document
should be read from the corpus:

- If ``item`` is one of the unique identifiers listed in the corpus
  module's ``items`` variable, then the corresponding document will
  be loaded from the NLTK corpus package.
- If ``item`` is a fileid, then that file will be read.

Additionally, corpus reader functions can be given lists of item
names; in which case, they will return a concatenation of the
corresponding documents.

Corpus reader functions are named based on the type of information
they return.  Some common examples, and their return types, are:

- words(): list of str
- sents(): list of (list of str)
- paras(): list of (list of (list of str))
- tagged_words(): list of (str,str) tuple
- tagged_sents(): list of (list of (str,str))
- tagged_paras(): list of (list of (list of (str,str)))
- chunked_sents(): list of (Tree w/ (str,str) leaves)
- parsed_sents(): list of (Tree with str leaves)
- parsed_paras(): list of (list of (Tree with str leaves))
- xml(): A single xml ElementTree
- raw(): unprocessed corpus contents

For example, to read a list of the words in the Brown Corpus, use
``nltk.corpus.brown.words()``:

    >>> from nltk.corpus import brown
    >>> print(", ".join(brown.words()))
    The, Fulton, County, Grand, Jury, said, ...

"""

from nltk.corpus.reader.plaintext import *
from nltk.corpus.reader.util import *
from nltk.corpus.reader.api import *
from nltk.corpus.reader.tagged import *
from nltk.corpus.reader.cmudict import *
from nltk.corpus.reader.conll import *
from nltk.corpus.reader.chunked import *
from nltk.corpus.reader.wordlist import *
from nltk.corpus.reader.xmldocs import *
from nltk.corpus.reader.ppattach import *
from nltk.corpus.reader.senseval import *
from nltk.corpus.reader.ieer import *
from nltk.corpus.reader.sinica_treebank import *
from nltk.corpus.reader.bracket_parse import *
from nltk.corpus.reader.indian import *
from nltk.corpus.reader.toolbox import *
from nltk.corpus.reader.timit import *
from nltk.corpus.reader.ycoe import *
from nltk.corpus.reader.rte import *
from nltk.corpus.reader.string_category import *
from nltk.corpus.reader.propbank import *
from nltk.corpus.reader.verbnet import *
from nltk.corpus.reader.bnc import *
from nltk.corpus.reader.nps_chat import *
from nltk.corpus.reader.wordnet import *
from nltk.corpus.reader.switchboard import *
from nltk.corpus.reader.dependency import *
from nltk.corpus.reader.nombank import *
from nltk.corpus.reader.ipipan import *
from nltk.corpus.reader.pl196x import *
from nltk.corpus.reader.knbc import *
from nltk.corpus.reader.chasen import *
from nltk.corpus.reader.childes import *
from nltk.corpus.reader.aligned import *
from nltk.corpus.reader.lin import *
from nltk.corpus.reader.semcor import *
from nltk.corpus.reader.framenet import *
from nltk.corpus.reader.udhr import *
from nltk.corpus.reader.bnc import *
from nltk.corpus.reader.sentiwordnet import *
from nltk.corpus.reader.twitter import *

# Make sure that nltk.corpus.reader.bracket_parse gives the module, not
# the function bracket_parse() defined in nltk.tree:
from nltk.corpus.reader import bracket_parse

__all__ = [
    'CorpusReader', 'CategorizedCorpusReader',
    'PlaintextCorpusReader', 'find_corpus_fileids',
    'TaggedCorpusReader', 'CMUDictCorpusReader',
    'ConllChunkCorpusReader', 'WordListCorpusReader',
    'PPAttachmentCorpusReader', 'SensevalCorpusReader',
    'IEERCorpusReader', 'ChunkedCorpusReader',
    'SinicaTreebankCorpusReader', 'BracketParseCorpusReader',
    'IndianCorpusReader', 'ToolboxCorpusReader',
    'TimitCorpusReader', 'YCOECorpusReader',
    'MacMorphoCorpusReader', 'SyntaxCorpusReader',
    'AlpinoCorpusReader', 'RTECorpusReader',
    'StringCategoryCorpusReader','EuroparlCorpusReader',
    'CategorizedBracketParseCorpusReader',
    'CategorizedTaggedCorpusReader',
    'CategorizedPlaintextCorpusReader',
    'PortugueseCategorizedPlaintextCorpusReader',
    'tagged_treebank_para_block_reader',
    'PropbankCorpusReader', 'VerbnetCorpusReader',
    'BNCCorpusReader', 'ConllCorpusReader',
    'XMLCorpusReader', 'NPSChatCorpusReader',
    'SwadeshCorpusReader', 'WordNetCorpusReader',
    'WordNetICCorpusReader', 'SwitchboardCorpusReader',
    'DependencyCorpusReader', 'NombankCorpusReader',
    'IPIPANCorpusReader', 'Pl196xCorpusReader',
    'TEICorpusView', 'KNBCorpusReader', 'ChasenCorpusReader',
    'CHILDESCorpusReader', 'AlignedCorpusReader',
    'TimitTaggedCorpusReader', 'LinThesaurusCorpusReader',
    'SemcorCorpusReader', 'FramenetCorpusReader', 'UdhrCorpusReader',
    'BNCCorpusReader', 'SentiWordNetCorpusReader', 'SentiSynset', 'TwitterCorpusReader'
]
<commit_msg>Renaming corpus readers for Twitter<commit_after># Natural Language Toolkit: Corpus Readers
#
# Copyright (C) 2001-2014 NLTK Project
# Author: Steven Bird <stevenbird1@gmail.com>
#         Edward Loper <edloper@gmail.com>
# URL: <http://nltk.org/>
# For license information, see LICENSE.TXT

"""
NLTK corpus readers.  The modules in this package provide functions
that can be used to read corpus fileids in a variety of formats.  These
functions can be used to read both the corpus fileids that are
distributed in the NLTK corpus package, and corpus fileids that are part
of external corpora.

Corpus Reader Functions
=======================
Each corpus module defines one or more "corpus reader functions",
which can be used to read documents from that corpus.  These functions
take an argument, ``item``, which is used to indicate which document
should be read from the corpus:

- If ``item`` is one of the unique identifiers listed in the corpus
  module's ``items`` variable, then the corresponding document will
  be loaded from the NLTK corpus package.
- If ``item`` is a fileid, then that file will be read.

Additionally, corpus reader functions can be given lists of item
names; in which case, they will return a concatenation of the
corresponding documents.

Corpus reader functions are named based on the type of information
they return.  Some common examples, and their return types, are:

- words(): list of str
- sents(): list of (list of str)
- paras(): list of (list of (list of str))
- tagged_words(): list of (str,str) tuple
- tagged_sents(): list of (list of (str,str))
- tagged_paras(): list of (list of (list of (str,str)))
- chunked_sents(): list of (Tree w/ (str,str) leaves)
- parsed_sents(): list of (Tree with str leaves)
- parsed_paras(): list of (list of (Tree with str leaves))
- xml(): A single xml ElementTree
- raw(): unprocessed corpus contents

For example, to read a list of the words in the Brown Corpus, use
``nltk.corpus.brown.words()``:

    >>> from nltk.corpus import brown
    >>> print(", ".join(brown.words()))
    The, Fulton, County, Grand, Jury, said, ...

"""

from nltk.corpus.reader.plaintext import *
from nltk.corpus.reader.util import *
from nltk.corpus.reader.api import *
from nltk.corpus.reader.tagged import *
from nltk.corpus.reader.cmudict import *
from nltk.corpus.reader.conll import *
from nltk.corpus.reader.chunked import *
from nltk.corpus.reader.wordlist import *
from nltk.corpus.reader.xmldocs import *
from nltk.corpus.reader.ppattach import *
from nltk.corpus.reader.senseval import *
from nltk.corpus.reader.ieer import *
from nltk.corpus.reader.sinica_treebank import *
from nltk.corpus.reader.bracket_parse import *
from nltk.corpus.reader.indian import *
from nltk.corpus.reader.toolbox import *
from nltk.corpus.reader.timit import *
from nltk.corpus.reader.ycoe import *
from nltk.corpus.reader.rte import *
from nltk.corpus.reader.string_category import *
from nltk.corpus.reader.propbank import *
from nltk.corpus.reader.verbnet import *
from nltk.corpus.reader.bnc import *
from nltk.corpus.reader.nps_chat import *
from nltk.corpus.reader.wordnet import *
from nltk.corpus.reader.switchboard import *
from nltk.corpus.reader.dependency import *
from nltk.corpus.reader.nombank import *
from nltk.corpus.reader.ipipan import *
from nltk.corpus.reader.pl196x import *
from nltk.corpus.reader.knbc import *
from nltk.corpus.reader.chasen import *
from nltk.corpus.reader.childes import *
from nltk.corpus.reader.aligned import *
from nltk.corpus.reader.lin import *
from nltk.corpus.reader.semcor import *
from nltk.corpus.reader.framenet import *
from nltk.corpus.reader.udhr import *
from nltk.corpus.reader.bnc import *
from nltk.corpus.reader.sentiwordnet import *
from nltk.corpus.reader.tweets import *

# Make sure that nltk.corpus.reader.bracket_parse gives the module, not
# the function bracket_parse() defined in nltk.tree:
from nltk.corpus.reader import bracket_parse

__all__ = [
    'CorpusReader', 'CategorizedCorpusReader',
    'PlaintextCorpusReader', 'find_corpus_fileids',
    'TaggedCorpusReader', 'CMUDictCorpusReader',
    'ConllChunkCorpusReader', 'WordListCorpusReader',
    'PPAttachmentCorpusReader', 'SensevalCorpusReader',
    'IEERCorpusReader', 'ChunkedCorpusReader',
    'SinicaTreebankCorpusReader', 'BracketParseCorpusReader',
    'IndianCorpusReader', 'ToolboxCorpusReader',
    'TimitCorpusReader', 'YCOECorpusReader',
    'MacMorphoCorpusReader', 'SyntaxCorpusReader',
    'AlpinoCorpusReader', 'RTECorpusReader',
    'StringCategoryCorpusReader','EuroparlCorpusReader',
    'CategorizedBracketParseCorpusReader',
    'CategorizedTaggedCorpusReader',
    'CategorizedPlaintextCorpusReader',
    'PortugueseCategorizedPlaintextCorpusReader',
    'tagged_treebank_para_block_reader',
    'PropbankCorpusReader', 'VerbnetCorpusReader',
    'BNCCorpusReader', 'ConllCorpusReader',
    'XMLCorpusReader', 'NPSChatCorpusReader',
    'SwadeshCorpusReader', 'WordNetCorpusReader',
    'WordNetICCorpusReader', 'SwitchboardCorpusReader',
    'DependencyCorpusReader', 'NombankCorpusReader',
    'IPIPANCorpusReader', 'Pl196xCorpusReader',
    'TEICorpusView', 'KNBCorpusReader', 'ChasenCorpusReader',
    'CHILDESCorpusReader', 'AlignedCorpusReader',
    'TimitTaggedCorpusReader', 'LinThesaurusCorpusReader',
    'SemcorCorpusReader', 'FramenetCorpusReader', 'UdhrCorpusReader',
    'BNCCorpusReader', 'SentiWordNetCorpusReader', 'SentiSynset', 'TwitterCorpusReader'
]
<|endoftext|>


<commit_before>import json
import logging
import random
import Algorithmia

from textblob import TextBlob
from text_corpus import TextCorpus
from aylienapiclient import textapi

logger = logging.getLogger(__name__)

class RtmEventHandler(object):
    def __init__(self, slack_clients, msg_writer, trump_corpus):
        self.clients = slack_clients
        self.msg_writer = msg_writer
        self.trump_corpus = trump_corpus

    def handle(self, event):

        if 'type' in event:
            self._handle_by_type(event['type'], event)

    def _handle_by_type(self, event_type, event):
        # See https://api.slack.com/rtm for a full list of events
        # if event_type == 'error':
            # error
            # ignore self.msg_writer.write_error(event['channel'], json.dumps(event))
        if event_type == 'message':
            # message was sent to channel
            self._handle_message(event)
        elif event_type == 'channel_joined':
            # you joined a channel
            self.msg_writer.send_message(event['channel'], "Let's make America great again!")
        elif event_type == 'group_joined':
            # you joined a private group
            self.msg_writer.write_help_message(event['channel'])
        else:
            pass

    def _handle_message(self, event):
        # Filter out messages from the bot itself
        if not self.clients.is_message_from_me(event['user']):

            msg_txt = event['text']

            if self.clients.is_bot_mention(msg_txt) or self._is_direct_message(event['channel']):
                # e.g. user typed: "@pybot tell me a joke!"
                if '/sentiment' in msg_txt:
                    self.msg_writer.write_analytics(event['channel'], msg_txt)
                    #if self.clients.is_bot_mention(msg_txt) or self._is_direct_message(event['channel']):
#                   txt_b = TextBlob(msg_txt)
#                   response = txt_b.tags
#                   self.msg_writer.send_message(event['channel'], response)
                    client = Algorithmia.client('sim3x6PzEv6m2icRR+23rqTTcOo1')
                    #response = txt_b.tags
                    #algo = client.algo('StanfordNLP/NamedEntityRecognition/0.2.0')
                    #entities = algo.pipe(msg_txt)
                    algo = client.algo('nlp/AutoTag/1.0.0')
                    response2 = algo.pipe(msg_txt)
                    response = response2.result[0]
                    client = textapi.Client("a19bb245", "2623b77754833e2711998a0b0bdad9db")
                    sentiment = client.Sentiment({"text": msg_txt})
                    self.msg_writer.send_message(event['channel'], sentiment['polarity'])

                elif '/classify' in msg_txt:

                    classifications = client.ClassifyByTaxonomy({"text": msg_txt, "taxonomy": "iab-qag"})
                    sent_str = ""
                    for category in classifications['categories']:
                        sent_str += category['label'] + ", "
                    sent_str = sent_str[:-1]
                    response = sent_str
                    self.msg_writer.send_message(event['channel'], response)

                else:
                    self.msg_writer.send_message(event['channel'], "What do you expect me to do?")

                return

        def _is_direct_message(self, channel_id):
            return channel_id.startswith('D')
<commit_msg>Update event_handler.py<commit_after>import json
import logging
import random
import Algorithmia

from textblob import TextBlob
from text_corpus import TextCorpus
from aylienapiclient import textapi

logger = logging.getLogger(__name__)

class RtmEventHandler(object):
    def __init__(self, slack_clients, msg_writer, trump_corpus):
        self.clients = slack_clients
        self.msg_writer = msg_writer
        self.trump_corpus = trump_corpus

    def handle(self, event):

        if 'type' in event:
            self._handle_by_type(event['type'], event)

    def _handle_by_type(self, event_type, event):
        # See https://api.slack.com/rtm for a full list of events
        # if event_type == 'error':
            # error
            # ignore self.msg_writer.write_error(event['channel'], json.dumps(event))
        if event_type == 'message':
            # message was sent to channel
            self._handle_message(event)
        elif event_type == 'channel_joined':
            # you joined a channel
            self.msg_writer.send_message(event['channel'], "Let's make America great again!")
        elif event_type == 'group_joined':
            # you joined a private group
            self.msg_writer.write_help_message(event['channel'])
        else:
            pass

    def _handle_message(self, event):
        # Filter out messages from the bot itself
        if not self.clients.is_message_from_me(event['user']):

            msg_txt = event['text']

            if self.clients.is_bot_mention(msg_txt) or self._is_direct_message(event['channel']):
                # e.g. user typed: "@pybot tell me a joke!"
                if '/sentiment' in msg_txt:
                    client = textapi.Client("a19bb245", "2623b77754833e2711998a0b0bdad9db")
                    algo = client.algo('nlp/AutoTag/1.0.0')
                    response2 = algo.pipe(msg_txt)
                    response = response2.result[0]
                    sentiment = client.Sentiment({"text": msg_txt})
                    self.msg_writer.send_message(event['channel'], sentiment['polarity'])

                elif '/tag' in msg_txt:
                    self.msg_writer.write_analytics(event['channel'], msg_txt)
                    #if self.clients.is_bot_mention(msg_txt) or self._is_direct_message(event['channel']):
                    txt_b = TextBlob(msg_txt)
#                   response = txt_b.tags
#                   self.msg_writer.send_message(event['channel'], response)
                    client = Algorithmia.client('sim3x6PzEv6m2icRR+23rqTTcOo1')
                    response = txt_b.tags
                    #algo = client.algo('StanfordNLP/NamedEntityRecognition/0.2.0')
                    #entities = algo.pipe(msg_txt)
                elif '/classify' in msg_txt:
                    client = textapi.Client("a19bb245", "2623b77754833e2711998a0b0bdad9db")
                    classifications = client.ClassifyByTaxonomy({"text": msg_txt, "taxonomy": "iab-qag"})
                    sent_str = ""
                    for category in classifications['categories']:
                        sent_str += category['label'] + ", "
                    sent_str = sent_str[:-1]
                    response = sent_str
                    self.msg_writer.send_message(event['channel'], response)

                else:
                    self.msg_writer.send_message(event['channel'], "What do you expect me to do?")

                return

        def _is_direct_message(self, channel_id):
            return channel_id.startswith('D')
<|endoftext|>


<commit_before>
import numpy as np
import astropy.units as u
from spectral_cube import SpectralCube
# from astropy.utils.console import ProgressBar
import sys
from warnings import warn

from bubble_segment2D import BubbleFinder2D
from bubble_objects import Bubble3D
from clustering import cluster_and_clean, cluster_brute_force
from utils import sig_clip
from progressbar import ProgressBar

class BubbleFinder(object):
    """docstring for BubbleFinder"""
    def __init__(self, cube, wcs=None, mask=None, sigma=None, empty_channel=0,
                 keep_threshold_mask=False):
        super(BubbleFinder, self).__init__()

        if not isinstance(cube, SpectralCube):
            if wcs is None:
                raise TypeError("When cube is not a SpectralCube, wcs must be"
                                " given.")
            cube = SpectralCube(cube, wcs)
            if mask is not None:
                cube = cube.with_mask(mask)

        self.cube = cube

        self.empty_channel = empty_channel

        if sigma is None:
            self.estimate_sigma()
        else:
            self.sigma = sigma

        self.keep_threshold_mask = keep_threshold_mask
        self._mask = None

    @property
    def cube(self):
        return self._cube

    @cube.setter
    def cube(self, input_cube):
        if input_cube.ndim != 3:
            raise TypeError("Given cube must be 3D.")

        self._cube = input_cube

    @property
    def mask(self):
        if not self.keep_threshold_mask:
            warn("Enable keep_threshold_mask to access the entire mask.")

        return self._mask

    def estimate_sigma(self, nsig=10):
        '''
        Use empty channels to estimate sigma. Uses iterative sigma clipping
        to obtain a robust estimate.
        '''

        self.sigma = sig_clip(self.cube[self.empty_channel], nsig=nsig)

    @property
    def sigma(self):
        return self._sigma

    @sigma.setter
    def sigma(self, val):
        if val < 0.0 or ~np.isfinite(val):
            raise ValueError("sigma must be positive and finite.")

        self._sigma = val

    def get_bubbles(self, verbose=True, overlap_frac=0.9, min_channels=3,
                    use_cube_mask=False, nsig=2., refit=False, distance=None,
                    cube_linewidth=None, multiprocess=True, nprocesses=None,
                    **kwargs):
        '''
        Perform segmentation on each channel, then cluster the results to find
        bubbles.
        '''

        if verbose:
            output = sys.stdout
        else:
            output = None

        if verbose:
            print("Running bubble finding plane-by-plane.")
        twod_results = \
            ProgressBar.map(_region_return,
                            ((self.cube[i],
                              self.cube.mask.include(view=(i, ))
                              if use_cube_mask else None,
                              i, self.sigma, nsig, overlap_frac,
                              self.keep_threshold_mask, distance) for i in
                             xrange(self.cube.shape[0])),
                            multiprocess=multiprocess,
                            nprocesses=nprocesses,
                            file=output,
                            step=self.cube.shape[0],
                            item_len=self.cube.shape[0])

        twod_regions = []
        if self.keep_threshold_mask:
            self._mask = np.zeros(self.cube.shape, dtype=np.bool)

        for out in twod_results:
            if self.keep_threshold_mask:
                chan, regions, mask_slice = out

                self._mask[chan] = mask_slice
            else:
                chan, regions = out

            twod_regions.extend(regions)

        self._bubbles = []
        self._unclustered_regions = []

        if len(twod_regions) == 0:
            warn("No bubbles found in the given cube.")
            return self

        bubble_props = np.vstack([bub.params for bub in twod_regions])

        # cluster_idx = cluster_and_clean(bubble_props, **kwargs)
        cluster_idx = cluster_brute_force(bubble_props, **kwargs)

        # Add the unclustered ones first
        unclusts = [twod_regions[idx] for idx in np.where(cluster_idx == 0)[0]]
        for reg in unclusts:
            self._unclustered_regions.append([reg])

        good_clusters = []
        for idx in np.unique(cluster_idx[cluster_idx > 0]):
            regions = [twod_regions[idx] for idx in
                       np.where(cluster_idx == idx)[0]]

            if len(regions) < min_channels:
                self._unclustered_regions.append(regions)
                continue

            chans = np.array([reg.channel_center for reg in regions])
            if chans.max() + 1 - chans.min() >= min_channels:
                good_clusters.append(regions)
            else:
                self._unclustered_regions.append(regions)

        if verbose:
            print("Creating bubbles and finding their properties.")
        # We need to pass a linewidth array for the cube, but don't want to
        # have to recompute it multiple times. Make sure it's using the mask
        if cube_linewidth is None:
            cube_linewidth = self.cube.with_mask(self.mask).linewidth_fwhm()
        # Now create the bubble objects and find their respective properties
        self._bubbles = ProgressBar.map(_make_bubble,
                                        ((regions, refit, self.cube, self.mask,
                                          distance, self.sigma, cube_linewidth)
                                         for regions in good_clusters),
                                        multiprocess=multiprocess,
                                        nprocesses=nprocesses,
                                        file=output,
                                        step=1,
                                        item_len=len(good_clusters))

        return self

    @property
    def bubbles(self):
        return self._bubbles

    @property
    def num_bubbles(self):
        return len(self.bubbles)

    @property
    def unclustered_regions(self):
        return self._unclustered_regions

    def visualize_bubbles(self, show=True, edges=False, ax=None,
                          moment0=None, region_col='b', edge_col='g',
                          log_scale=False, plot_threeD_shapes=True,
                          plot_twoD_shapes=False, plot_unclustered=False):
        '''
        Show the location of the bubbles on the moment 0 array.
        '''
        if len(self.bubbles) == 0:
            warn("No bubbles were found. Nothing to show.")
            return

        if moment0 is None:
            # Create the moment array from the cube
            moment0 = self.cube.moment0().value

        import matplotlib.pyplot as p

        if ax is None:
            ax = p.subplot(111)

        if log_scale:
            ax.imshow(np.log10(moment0), cmap='afmhot', origin='lower')
        else:
            ax.imshow(moment0, cmap='afmhot', origin='lower')

        if plot_unclustered:
            for clust in self.unclustered_regions:
                for reg in clust:
                    ax.add_patch(reg.as_patch(color=region_col, fill=False,
                                              linewidth=2))
                    ax.plot(reg.x, reg.y, region_col + 'D')
                    if edges:
                        ax.plot(reg.shell_coords[:, 1], reg.shell_coords[:, 0],
                                edge_col + "o")
        else:
            for bub in self.bubbles:
                if plot_twoD_shapes:
                    for reg in bub.twoD_regions:
                        ax.add_patch(reg.as_patch(color=region_col,
                                                  fill=False,
                                                  linewidth=2))
                        ax.plot(reg.x, reg.y, region_col + "D")
                        if edges:
                            ax.plot(reg.shell_coords[:, 1],
                                    reg.shell_coords[:, 0], edge_col + "o")
                if plot_threeD_shapes:
                    if not plot_twoD_shapes:
                        linewidth = 2
                        linestyle = "-"
                        color = region_col
                    else:
                        linewidth = 3
                        linestyle = "--"
                        color = 'm'
                    ax.add_patch(bub.as_patch(color=color, fill=False,
                                              linewidth=linewidth,
                                              linestyle=linestyle))
                    ax.plot(bub.x, bub.y, region_col + 'D')
                    if edges:
                        ax.plot(bub.shell_coords[:, 2], bub.shell_coords[:, 1],
                                edge_col + "o")

        p.xlim([0, self.cube.shape[2]])
        p.ylim([0, self.cube.shape[1]])

        if show:
            p.show()
        else:
            return ax

    def visualize_channel_maps(self, all_chans=False, subplot=False,
                               edges=False, plot_unclustered=False,
                               interactive_plot=True,
                               show_mask_contours=False):
        '''
        Plot each channel optionally overlaid with the regions and/or the
        edges.
        '''

        # Loop through all, or just those that have at least one region
        if all_chans:
            chans = np.arange(self.cube.shape[0], dtype=np.int)
        else:
            chans = []
            for bub in self.bubbles:
                chans.extend(list(bub.twoD_region_params()[:, -1]))
            chans = np.unique(np.array(chans, dtype=np.int))

        import matplotlib.pyplot as p
        if interactive_plot:
            p.ion()

        for chan in chans:
            if subplot:
                raise NotImplementedError("Need subplots at some point.")
            else:
                ax = p.subplot(111)

            ax.imshow(self.cube[chan].value, cmap='afmhot', origin='lower')

            if show_mask_contours:
                if self.keep_threshold_mask:
                    ax.contour(self.mask[chan], colors='k')
                else:
                    warn("'keep_threshold_mask' must be enabled to plot mask "
                         "contours.")

            for bub in self.bubbles:
                if chan < bub.channel_start and chan > bub.channel_end:
                    continue
                for twod in bub.twoD_regions:
                    if twod.channel_center != chan:
                        continue
                    ax.add_patch(twod.as_patch(color='b', fill=False,
                                               linewidth=2))
                    ax.plot(twod.x, twod.y, 'bD')
                    if edges:
                        ax.plot(twod.shell_coords[:, 1],
                                twod.shell_coords[:, 0], "go")

            # Also show the unclustered regions in each channel
            if plot_unclustered:
                for clust in self.unclustered_regions:
                    for reg in clust:
                        if reg.channel_center != chan:
                            continue
                        ax.add_patch(reg.as_patch(color='r', fill=False,
                                                  linewidth=2))
                        ax.plot(reg.x, reg.y, 'rD')
                        if edges:
                            ax.plot(reg.shell_coords[:, 1],
                                    reg.shell_coords[:, 0], "ro")

            p.xlim([0, self.cube.shape[2]])
            p.ylim([0, self.cube.shape[1]])

            if interactive_plot:
                p.draw()
                raw_input("Channel {}".format(chan))
                p.clf()
            else:
                p.show()

def _region_return(imps):
    arr, mask, i, sigma, nsig, overlap_frac, return_mask, distance = imps
    bubs = BubbleFinder2D(arr, channel=i,
                          mask=mask, sigma=sigma).\
        multiscale_bubblefind(nsig=nsig,
                              overlap_frac=overlap_frac,
                              distance=distance)
    if return_mask:
        return i, bubs.regions, bubs.mask

    return i, bubs.regions

def _make_bubble(imps):
    regions, refit, cube, mask, distance, sigma, lwidth = imps
    return Bubble3D.from_2D_regions(regions, refit=refit,
                                    cube=cube, mask=mask,
                                    distance=distance,
                                    sigma=sigma, linewidth=lwidth)
<commit_msg>Set galaxy properties<commit_after>
import numpy as np
import astropy.units as u
from spectral_cube import SpectralCube
# from astropy.utils.console import ProgressBar
from astropy.coordinates import SkyCoord, Angle
import sys
from warnings import warn

from bubble_segment2D import BubbleFinder2D
from bubble_objects import Bubble3D
from bubble_catalog import PPV_Catalog
from clustering import cluster_and_clean, cluster_brute_force
from utils import sig_clip
from progressbar import ProgressBar

GALAXY_KEYS = ["inclination", "position_angle", "center_coord",
               "scale_height"]

class BubbleFinder(object):
    """docstring for BubbleFinder"""
    def __init__(self, cube, wcs=None, mask=None, sigma=None, empty_channel=0,
                 keep_threshold_mask=False, distance=None, galaxy_props={}):
        super(BubbleFinder, self).__init__()

        if not isinstance(cube, SpectralCube):
            if wcs is None:
                raise TypeError("When cube is not a SpectralCube, wcs must be"
                                " given.")
            cube = SpectralCube(cube, wcs)
            if mask is not None:
                cube = cube.with_mask(mask)

        self.cube = cube

        self.empty_channel = empty_channel

        if sigma is None:
            self.estimate_sigma()
        else:
            self.sigma = sigma

        self.keep_threshold_mask = keep_threshold_mask
        self._mask = None
        self.distance = distance
        self.galaxy_props = galaxy_props

    @property
    def cube(self):
        return self._cube

    @cube.setter
    def cube(self, input_cube):
        if input_cube.ndim != 3:
            raise TypeError("Given cube must be 3D.")

        self._cube = input_cube

    @property
    def mask(self):
        if not self.keep_threshold_mask:
            warn("Enable keep_threshold_mask to access the entire mask.")

        return self._mask

    def estimate_sigma(self, nsig=10):
        '''
        Use empty channels to estimate sigma. Uses iterative sigma clipping
        to obtain a robust estimate.
        '''

        self.sigma = sig_clip(self.cube[self.empty_channel], nsig=nsig)

    @property
    def sigma(self):
        return self._sigma

    @sigma.setter
    def sigma(self, val):
        if val < 0.0 or ~np.isfinite(val):
            raise ValueError("sigma must be positive and finite.")

        self._sigma = val

    @property
    def galaxy_props(self):
        return self._galaxy_props

    @galaxy_props.setter
    def galaxy_props(self, input_dict):

        # Make sure all of the kwargs are given.
        in_input = [True if key in GALAXY_KEYS else False
                    for key in input_dict]
        if not np.all(in_input):
            missing = list(set(GALAXY_KEYS) - set(in_input))
            raise KeyError("galaxy_props is missing these keys: {}"
                           .format(missing))

        if not isinstance(input_dict["center_coord"], SkyCoord):
            raise TypeError("center_coords must be a SkyCoord.")

        if not input_dict["distance"].unit.is_equivalent(u.pc):
            raise u.UnitsError("distance must have a unit of distance")

        if not input_dict["scale_height"].unit.is_equivalent(u.pc):
            raise u.UnitsError("scale_height must have a unit of distance")

        if not isinstance(input_dict["inclination"], Angle):
            raise TypeError("inclination must be an Angle.")

        if not isinstance(input_dict["position_angle"], Angle):
            raise TypeError("position_angle must be an Angle.")

        self._galaxy_props = input_dict

    def get_bubbles(self, verbose=True, overlap_frac=0.9, min_channels=3,
                    use_cube_mask=False, nsig=2., refit=False, distance=None,
                    cube_linewidth=None, multiprocess=True, nprocesses=None,
                    **kwargs):
        '''
        Perform segmentation on each channel, then cluster the results to find
        bubbles.
        '''

        if verbose:
            output = sys.stdout
        else:
            output = None

        if verbose:
            print("Running bubble finding plane-by-plane.")
        twod_results = \
            ProgressBar.map(_region_return,
                            ((self.cube[i],
                              self.cube.mask.include(view=(i, ))
                              if use_cube_mask else None,
                              i, self.sigma, nsig, overlap_frac,
                              self.keep_threshold_mask, distance) for i in
                             xrange(self.cube.shape[0])),
                            multiprocess=multiprocess,
                            nprocesses=nprocesses,
                            file=output,
                            step=self.cube.shape[0],
                            item_len=self.cube.shape[0])

        twod_regions = []
        if self.keep_threshold_mask:
            self._mask = np.zeros(self.cube.shape, dtype=np.bool)

        for out in twod_results:
            if self.keep_threshold_mask:
                chan, regions, mask_slice = out

                self._mask[chan] = mask_slice
            else:
                chan, regions = out

            twod_regions.extend(regions)

        self._bubbles = []
        self._unclustered_regions = []

        if len(twod_regions) == 0:
            warn("No bubbles found in the given cube.")
            return self

        bubble_props = np.vstack([bub.params for bub in twod_regions])

        # cluster_idx = cluster_and_clean(bubble_props, **kwargs)
        cluster_idx = cluster_brute_force(bubble_props, **kwargs)

        # Add the unclustered ones first
        unclusts = [twod_regions[idx] for idx in np.where(cluster_idx == 0)[0]]
        for reg in unclusts:
            self._unclustered_regions.append([reg])

        good_clusters = []
        for idx in np.unique(cluster_idx[cluster_idx > 0]):
            regions = [twod_regions[idx] for idx in
                       np.where(cluster_idx == idx)[0]]

            if len(regions) < min_channels:
                self._unclustered_regions.append(regions)
                continue

            chans = np.array([reg.channel_center for reg in regions])
            if chans.max() + 1 - chans.min() >= min_channels:
                good_clusters.append(regions)
            else:
                self._unclustered_regions.append(regions)

        if verbose:
            print("Creating bubbles and finding their properties.")
        # We need to pass a linewidth array for the cube, but don't want to
        # have to recompute it multiple times. Make sure it's using the mask
        if cube_linewidth is None:
            cube_linewidth = self.cube.with_mask(self.mask).linewidth_fwhm()
        # Now create the bubble objects and find their respective properties
        self._bubbles = ProgressBar.map(_make_bubble,
                                        ((regions, refit, self.cube, self.mask,
                                          distance, self.sigma, cube_linewidth)
                                         for regions in good_clusters),
                                        multiprocess=False,
                                        nprocesses=nprocesses,
                                        file=output,
                                        step=1,
                                        item_len=len(good_clusters))

        return self

    @property
    def bubbles(self):
        return self._bubbles

    @property
    def num_bubbles(self):
        return len(self.bubbles)

    @property
    def unclustered_regions(self):
        return self._unclustered_regions

    def to_catalog(self):
        '''
        Returns a PPV_Catalog to explore the population properties.
        '''
        return PPV_Catalog(self.bubbles)

    def visualize_bubbles(self, show=True, edges=False, ax=None,
                          moment0=None, region_col='b', edge_col='g',
                          log_scale=False, plot_threeD_shapes=True,
                          plot_twoD_shapes=False, plot_unclustered=False):
        '''
        Show the location of the bubbles on the moment 0 array.
        '''
        if len(self.bubbles) == 0:
            warn("No bubbles were found. Nothing to show.")
            return

        if moment0 is None:
            # Create the moment array from the cube
            moment0 = self.cube.moment0().value

        import matplotlib.pyplot as p

        if ax is None:
            ax = p.subplot(111)

        if log_scale:
            ax.imshow(np.log10(moment0), cmap='afmhot', origin='lower')
        else:
            ax.imshow(moment0, cmap='afmhot', origin='lower')

        if plot_unclustered:
            for clust in self.unclustered_regions:
                for reg in clust:
                    ax.add_patch(reg.as_patch(color=region_col, fill=False,
                                              linewidth=2))
                    ax.plot(reg.x, reg.y, region_col + 'D')
                    if edges:
                        ax.plot(reg.shell_coords[:, 1], reg.shell_coords[:, 0],
                                edge_col + "o")
        else:
            for bub in self.bubbles:
                if plot_twoD_shapes:
                    for reg in bub.twoD_regions:
                        ax.add_patch(reg.as_patch(color=region_col,
                                                  fill=False,
                                                  linewidth=2))
                        ax.plot(reg.x, reg.y, region_col + "D")
                        if edges:
                            ax.plot(reg.shell_coords[:, 1],
                                    reg.shell_coords[:, 0], edge_col + "o")
                if plot_threeD_shapes:
                    if not plot_twoD_shapes:
                        linewidth = 2
                        linestyle = "-"
                        color = region_col
                    else:
                        linewidth = 3
                        linestyle = "--"
                        color = 'm'
                    ax.add_patch(bub.as_patch(color=color, fill=False,
                                              linewidth=linewidth,
                                              linestyle=linestyle))
                    ax.plot(bub.x, bub.y, region_col + 'D')
                    if edges:
                        ax.plot(bub.shell_coords[:, 2], bub.shell_coords[:, 1],
                                edge_col + "o")

        p.xlim([0, self.cube.shape[2]])
        p.ylim([0, self.cube.shape[1]])

        if show:
            p.show()
        else:
            return ax

    def visualize_channel_maps(self, all_chans=False, subplot=False,
                               edges=False, plot_unclustered=False,
                               interactive_plot=True,
                               show_mask_contours=False):
        '''
        Plot each channel optionally overlaid with the regions and/or the
        edges.
        '''

        # Loop through all, or just those that have at least one region
        if all_chans:
            chans = np.arange(self.cube.shape[0], dtype=np.int)
        else:
            chans = []
            for bub in self.bubbles:
                chans.extend(list(bub.twoD_region_params()[:, -1]))
            chans = np.unique(np.array(chans, dtype=np.int))

        import matplotlib.pyplot as p
        if interactive_plot:
            p.ion()

        for chan in chans:
            if subplot:
                raise NotImplementedError("Need subplots at some point.")
            else:
                ax = p.subplot(111)

            ax.imshow(self.cube[chan].value, cmap='afmhot', origin='lower')

            if show_mask_contours:
                if self.keep_threshold_mask:
                    ax.contour(self.mask[chan], colors='k')
                else:
                    warn("'keep_threshold_mask' must be enabled to plot mask "
                         "contours.")

            for bub in self.bubbles:
                if chan < bub.channel_start and chan > bub.channel_end:
                    continue
                for twod in bub.twoD_regions:
                    if twod.channel_center != chan:
                        continue
                    ax.add_patch(twod.as_patch(color='b', fill=False,
                                               linewidth=2))
                    ax.plot(twod.x, twod.y, 'bD')
                    if edges:
                        ax.plot(twod.shell_coords[:, 1],
                                twod.shell_coords[:, 0], "go")

            # Also show the unclustered regions in each channel
            if plot_unclustered:
                for clust in self.unclustered_regions:
                    for reg in clust:
                        if reg.channel_center != chan:
                            continue
                        ax.add_patch(reg.as_patch(color='r', fill=False,
                                                  linewidth=2))
                        ax.plot(reg.x, reg.y, 'rD')
                        if edges:
                            ax.plot(reg.shell_coords[:, 1],
                                    reg.shell_coords[:, 0], "ro")

            p.xlim([0, self.cube.shape[2]])
            p.ylim([0, self.cube.shape[1]])

            if interactive_plot:
                p.draw()
                raw_input("Channel {}".format(chan))
                p.clf()
            else:
                p.show()

def _region_return(imps):
    arr, mask, i, sigma, nsig, overlap_frac, return_mask, distance = imps
    bubs = BubbleFinder2D(arr, channel=i,
                          mask=mask, sigma=sigma).\
        multiscale_bubblefind(nsig=nsig,
                              overlap_frac=overlap_frac,
                              distance=distance)
    if return_mask:
        return i, bubs.regions, bubs.mask

    return i, bubs.regions

def _make_bubble(imps):
    regions, refit, cube, mask, distance, sigma, lwidth = imps
    return Bubble3D.from_2D_regions(regions, refit=refit,
                                    cube=cube, mask=mask,
                                    distance=distance,
                                    sigma=sigma, linewidth=lwidth)
<|endoftext|>