<commit_before>from raysect.optical import World, Node, translate, rotate, Point3D, d65_white, ConstantSF, InterpolatedSF
from raysect.optical.material import Lambert, UniformSurfaceEmitter
from raysect.optical.library import *
from raysect.primitive import Sphere, Box
from matplotlib.pyplot import *
from numpy import array

"""
Cornell Box Demo
================

This demo renders a variant of the classic Cornell Box scene.

To the original Cornell Box see:

    http://www.graphics.cornell.edu/online/box/data.html

The wall colours and light spectrum used in this demo are the values measured
for the physical Cornell Box.
"""

# define reflectivity for box surfaces
wavelengths = array(
    [400, 404, 408, 412, 416, 420, 424, 428, 432, 436, 440, 444, 448, 452, 456, 460, 464, 468, 472, 476, 480, 484, 488,
     492, 496, 500, 504, 508, 512, 516, 520, 524, 528, 532, 536, 540, 544, 548, 552, 556, 560, 564, 568, 572, 576, 580,
     584, 588, 592, 596, 600, 604, 608, 612, 616, 620, 624, 628, 632, 636, 640, 644, 648, 652, 656, 660, 664, 668, 672,
     676, 680, 684, 688, 692, 696, 700])

white = array(
    [0.343, 0.445, 0.551, 0.624, 0.665, 0.687, 0.708, 0.723, 0.715, 0.71, 0.745, 0.758, 0.739, 0.767, 0.777, 0.765,
     0.751, 0.745, 0.748, 0.729, 0.745, 0.757, 0.753, 0.75, 0.746, 0.747, 0.735, 0.732, 0.739, 0.734, 0.725, 0.721,
     0.733, 0.725, 0.732, 0.743, 0.744, 0.748, 0.728, 0.716, 0.733, 0.726, 0.713, 0.74, 0.754, 0.764, 0.752, 0.736,
     0.734, 0.741, 0.74, 0.732, 0.745, 0.755, 0.751, 0.744, 0.731, 0.733, 0.744, 0.731, 0.712, 0.708, 0.729, 0.73,
     0.727, 0.707, 0.703, 0.729, 0.75, 0.76, 0.751, 0.739, 0.724, 0.73, 0.74, 0.737])

green = array(
    [0.092, 0.096, 0.098, 0.097, 0.098, 0.095, 0.095, 0.097, 0.095, 0.094, 0.097, 0.098, 0.096, 0.101, 0.103, 0.104,
     0.107, 0.109, 0.112, 0.115, 0.125, 0.14, 0.16, 0.187, 0.229, 0.285, 0.343, 0.39, 0.435, 0.464, 0.472, 0.476, 0.481,
     0.462, 0.447, 0.441, 0.426, 0.406, 0.373, 0.347, 0.337, 0.314, 0.285, 0.277, 0.266, 0.25, 0.23, 0.207, 0.186,
     0.171, 0.16, 0.148, 0.141, 0.136, 0.13, 0.126, 0.123, 0.121, 0.122, 0.119, 0.114, 0.115, 0.117, 0.117, 0.118, 0.12,
     0.122, 0.128, 0.132, 0.139, 0.144, 0.146, 0.15, 0.152, 0.157, 0.159])

red = array(
    [0.04, 0.046, 0.048, 0.053, 0.049, 0.05, 0.053, 0.055, 0.057, 0.056, 0.059, 0.057, 0.061, 0.061, 0.06, 0.062, 0.062,
     0.062, 0.061, 0.062, 0.06, 0.059, 0.057, 0.058, 0.058, 0.058, 0.056, 0.055, 0.056, 0.059, 0.057, 0.055, 0.059,
     0.059, 0.058, 0.059, 0.061, 0.061, 0.063, 0.063, 0.067, 0.068, 0.072, 0.08, 0.09, 0.099, 0.124, 0.154, 0.192,
     0.255, 0.287, 0.349, 0.402, 0.443, 0.487, 0.513, 0.558, 0.584, 0.62, 0.606, 0.609, 0.651, 0.612, 0.61, 0.65, 0.638,
     0.627, 0.62, 0.63, 0.628, 0.642, 0.639, 0.657, 0.639, 0.635, 0.642])

white_reflectivity = InterpolatedSF(wavelengths, white)
red_reflectivity = InterpolatedSF(wavelengths, red)
green_reflectivity = InterpolatedSF(wavelengths, green)

# define light spectrum
light_spectrum = InterpolatedSF(array([400, 500, 600, 700]), array([0.0, 8.0, 15.6, 18.4]))

# set-up scenegraph
world = World()

# enclosing box
enclosure = Node(world)

e_back = Box(Point3D(-1, -1, 0), Point3D(1, 1, 0),
             parent=enclosure,
             transform=translate(0, 0, 1) * rotate(0, 0, 0),
             material=Lambert(white_reflectivity))

e_bottom = Box(Point3D(-1, -1, 0), Point3D(1, 1, 0),
               parent=enclosure,
               transform=translate(0, -1, 0) * rotate(0, -90, 0),
               # material=m)
               material=Lambert(white_reflectivity))

e_top = Box(Point3D(-1, -1, 0), Point3D(1, 1, 0),
            parent=enclosure,
            transform=translate(0, 1, 0) * rotate(0, 90, 0),
            material=Lambert(white_reflectivity))

e_left = Box(Point3D(-1, -1, 0), Point3D(1, 1, 0),
             parent=enclosure,
             transform=translate(1, 0, 0) * rotate(-90, 0, 0),
             material=Lambert(red_reflectivity))

e_right = Box(Point3D(-1, -1, 0), Point3D(1, 1, 0),
              parent=enclosure,
              transform=translate(-1, 0, 0) * rotate(90, 0, 0),
              material=Lambert(green_reflectivity))

# ceiling light
light = Box(Point3D(-0.4, -0.4, -0.01), Point3D(0.4, 0.4, 0.0),
            parent=enclosure,
            transform=translate(0, 1, 0) * rotate(0, 90, 0),
            material=UniformSurfaceEmitter(light_spectrum, 2))

# light = Box(Point3D(-0.4, -0.4, -0.01), Point3D(0.4, 0.4, 0.0),
#             parent=enclosure,
#             transform=translate(0, 1, 0) * rotate(0, 90, 0),
#             material=UniformSurfaceEmitter(d65_white, 2))

# back_light = Sphere(0.1,
#     parent=enclosure,
#     transform=translate(0.80, -0.85, 0.80)*rotate(0, 0, 0),
#     material=UniformSurfaceEmitter(light_spectrum, 10.0))

# objects in enclosure
box = Box(Point3D(-0.4, 0, -0.4), Point3D(0.3, 1.4, 0.3),
          parent=world,
          transform=translate(0.4, -1 + 1e-6, 0.4)*rotate(30, 0, 0),
          # material=RoughTungsten(0.5))
          # material=Lambert())
          material=schott("N-BK7"))

sphere = Sphere(0.4,
    parent=world,
    transform=translate(-0.4, -0.6 + 1e-6, -0.4)*rotate(0, 0, 0),
    # material=RoughGold(0.1))
    # material=Lambert())
    # material=Titanium())
    material=schott("N-BK7"))

from raysect.optical.observer import PinholeCamera, CCDArray, FibreOptic
from raysect.optical.observer import SpectralPipeline0D, PowerPipeline0D
from raysect.core.workflow import SerialEngine

power = PowerPipeline0D()

camera = FibreOptic(parent=world, transform=translate(0, 0, -3.3) * rotate(0, 0, 0), pipelines=[power])
camera.acceptance_angle = 10
camera.pixel_samples = 100
camera.spectral_bins = 1
camera.spectral_rays = 1
camera.ray_importance_sampling = True
camera.ray_important_path_weight = 0.25
camera.ray_max_depth = 500
camera.ray_extinction_min_depth = 3
camera.ray_extinction_prob = 0.01
camera.quiet = True

camera.render_engine = SerialEngine()

for i in range(1000):
    camera.observe()

print(power.value.mean, power.value.error())<commit_msg>Removed file accidentally added to commit.<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python
# coding: utf-8
#
# Copyright 2010 Alexandre Fiori
# based on the original Tornado by Facebook
#
# 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 collections
import functools
import sys

import cyclone.web
import cyclone.redis

from twisted.python import log
from twisted.internet import defer, reactor

class Application(cyclone.web.Application):
    def __init__(self):
        handlers = [
            (r"/text/(.+)", TextHandler),
            (r"/queue/(.+)", QueueHandler),
        ]
        settings = dict(
            debug=True,
            static_path="./frontend/static",
            template_path="./frontend/template",
        )
        RedisMixin.setup("127.0.0.1", 6379, 0, 10)
        cyclone.web.Application.__init__(self, handlers, **settings)

class RedisMixin(object):
    dbconn = None
    psconn = None
    channels = collections.defaultdict(lambda: [])

    @classmethod
    def setup(self, host, port, dbid, poolsize):
        # PubSub client connection
        qf = cyclone.redis.SubscriberFactory()
        qf.maxDelay = 20
        qf.protocol = QueueProtocol
        reactor.connectTCP(host, port, qf)

        # Normal client connection
        RedisMixin.dbconn = cyclone.redis.lazyConnectionPool(host, port,
                                                             dbid, poolsize)

    def subscribe(self, channel):
        if RedisMixin.psconn is None:
            raise cyclone.web.HTTPError(503)  # Service Unavailable

        if channel not in RedisMixin.channels:
            log.msg("Subscribing entire server to %s" % channel)
            if "*" in channel:
                RedisMixin.psconn.psubscribe(channel)
            else:
                RedisMixin.psconn.subscribe(channel)

        RedisMixin.channels[channel].append(self)
        log.msg("Client %s subscribed to %s" % \
                (self.request.remote_ip, channel))

    def unsubscribe_all(self, ign):
        # Unsubscribe peer from all channels
        for channel, peers in RedisMixin.channels.iteritems():
            try:
                peers.pop(peers.index(self))
            except:
                continue

            log.msg("Client %s unsubscribed from %s" % \
                    (self.request.remote_ip, channel))

            # Unsubscribe from channel if no peers are listening
            if not len(peers) and RedisMixin.psconn is not None:
                log.msg("Unsubscribing entire server from %s" % channel)
                if "*" in channel:
                    RedisMixin.psconn.punsubscribe(channel)
                else:
                    RedisMixin.psconn.unsubscribe(channel)

    def broadcast(self, pattern, channel, message):
        peers = self.channels.get(pattern or channel)
        if not peers:
            return

        # Broadcast the message to all peers in channel
        for peer in peers:
            # peer is an HTTP client, RequestHandler
            peer.write("%s: %s\r\n" % (channel, message))
            peer.flush()

# Provide GET, SET and DELETE redis operations via HTTP
class TextHandler(cyclone.web.RequestHandler, RedisMixin):
    @defer.inlineCallbacks
    def get(self, key):
        try:
            value = yield self.dbconn.get(key)
        except Exception, e:
            log.err("Redis failed to get('%s'): %s" % (key, str(e)))
            raise cyclone.web.HTTPError(503)

        self.set_header("Content-Type", "text/plain")
        self.finish("%s=%s\r\n" % (key, value))

    @defer.inlineCallbacks
    def post(self, key):
        value = self.get_argument("value")
        try:
            yield self.dbconn.set(key, value)
        except Exception, e:
            log.err("Redis failed to set('%s', '%s'): %s" % \
                    (key, value, str(e)))
            raise cyclone.web.HTTPError(503)

        self.set_header("Content-Type", "text/plain")
        self.finish("%s=%s\r\n" % (key, value))

    @defer.inlineCallbacks
    def delete(self, key):
        try:
            n = yield self.dbconn.delete(key)
        except Exception, e:
            log.err("Redis failed to del('%s'): %s" % (key, str(e)))
            raise cyclone.web.HTTPError(503)

        self.set_header("Content-Type", "text/plain")
        self.finish("DEL %s=%d\r\n" % (key, n))

# GET will subscribe to channels or patterns
# POST will (obviously) post messages to channels
class QueueHandler(cyclone.web.RequestHandler, RedisMixin):
    @cyclone.web.asynchronous
    def get(self, channels):
        try:
            channels = channels.split(",")
        except Exception, e:
            log.err("Could not split channel names: %s" % str(e))
            raise cyclone.web.HTTPError(400, str(e))

        self.set_header("Content-Type", "text/plain")
        self.notifyFinish().addCallback(
            functools.partial(RedisMixin.unsubscribe_all, self))

        for channel in channels:
            self.subscribe(channel)
            self.write("subscribed to %s\r\n" % channel)
        self.flush()

    @defer.inlineCallbacks
    def post(self, channel):
        message = self.get_argument("message")

        try:
            n = yield self.dbconn.publish(channel, message.encode("utf-8"))
        except Exception, e:
            log.msg("Redis failed to publish('%s', '%s'): %s" % \
                    (channel, repr(message), str(e)))
            raise cyclone.web.HTTPError(503)

        self.set_header("Content-Type", "text/plain")
        self.finish("OK %d\r\n" % n)

class QueueProtocol(cyclone.redis.SubscriberProtocol, RedisMixin):
    def messageReceived(self, pattern, channel, message):
        # When new messages are published to Redis channels or patterns,
        # they are broadcasted to all HTTP clients subscribed to those
        # channels.
        RedisMixin.broadcast(self, pattern, channel, message)

    def connectionMade(self):
        RedisMixin.psconn = self

        # If we lost connection with Redis during operation, we
        # re-subscribe to all channels once the connection is re-established.
        for channel in self.channels:
            if "*" in channel:
                self.psubscribe(channel)
            else:
                self.subscribe(channel)

    def connectionLost(self, why):
        RedisMixin.psconn = None

def main():
    log.startLogging(sys.stdout)
    reactor.listenTCP(8888, Application(), interface="127.0.0.1")
    reactor.run()

if __name__ == "__main__":
    main()
<commit_msg>review, +updates<commit_after>#!/usr/bin/env python
# coding: utf-8
#
# Copyright 2010 Alexandre Fiori
# based on the original Tornado by Facebook
#
# 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 collections
import functools
import sys

import cyclone.web
import cyclone.redis

from twisted.python import log
from twisted.internet import defer, reactor

class Application(cyclone.web.Application):
    def __init__(self):
        handlers = [
            (r"/", IndexHandler),
            (r"/text/(.+)", TextHandler),
            (r"/queue/(.+)", QueueHandler),
        ]
        settings = dict(
            debug=True,
        )
        RedisMixin.setup("127.0.0.1", 6379, 0, 10)
        cyclone.web.Application.__init__(self, handlers, **settings)

class RedisMixin(object):
    dbconn = None
    psconn = None
    channels = collections.defaultdict(lambda: [])

    @classmethod
    def setup(self, host, port, dbid, poolsize):
        # PubSub client connection
        qf = cyclone.redis.SubscriberFactory()
        qf.maxDelay = 20
        qf.protocol = QueueProtocol
        reactor.connectTCP(host, port, qf)

        # Normal client connection
        RedisMixin.dbconn = cyclone.redis.lazyConnectionPool(host, port,
                                                             dbid, poolsize)

    def subscribe(self, channel):
        if RedisMixin.psconn is None:
            raise cyclone.web.HTTPError(503)  # Service Unavailable

        if channel not in RedisMixin.channels:
            log.msg("Subscribing entire server to %s" % channel)
            if "*" in channel:
                RedisMixin.psconn.psubscribe(channel)
            else:
                RedisMixin.psconn.subscribe(channel)

        RedisMixin.channels[channel].append(self)
        log.msg("Client %s subscribed to %s" %
                (self.request.remote_ip, channel))

    def unsubscribe_all(self, ign):
        # Unsubscribe peer from all channels
        for channel, peers in RedisMixin.channels.iteritems():
            try:
                peers.pop(peers.index(self))
            except:
                continue

            log.msg("Client %s unsubscribed from %s" %
                    (self.request.remote_ip, channel))

            # Unsubscribe from channel if no peers are listening
            if not len(peers) and RedisMixin.psconn is not None:
                log.msg("Unsubscribing entire server from %s" % channel)
                if "*" in channel:
                    RedisMixin.psconn.punsubscribe(channel)
                else:
                    RedisMixin.psconn.unsubscribe(channel)

    def broadcast(self, pattern, channel, message):
        peers = self.channels.get(pattern or channel)
        if not peers:
            return

        # Broadcast the message to all peers in channel
        for peer in peers:
            # peer is an HTTP client, RequestHandler
            peer.write("%s: %s\r\n" % (channel, message))
            peer.flush()

class IndexHandler(cyclone.web.RequestHandler):
    readme = open("README").read().replace("\n", "\r\n")

    def get(self):
        self.set_header("Content-Type", "text/plain")
        self.write(self.readme)

# Provide GET, SET and DELETE redis operations via HTTP
class TextHandler(cyclone.web.RequestHandler, RedisMixin):
    @defer.inlineCallbacks
    def get(self, key):
        try:
            value = yield self.dbconn.get(key)
        except Exception, e:
            log.err("Redis failed to get('%s'): %s" % (key, str(e)))
            raise cyclone.web.HTTPError(503)

        self.set_header("Content-Type", "text/plain")
        self.finish("%s=%s\r\n" % (key, value))

    @defer.inlineCallbacks
    def post(self, key):
        value = self.get_argument("value")
        try:
            yield self.dbconn.set(key, value)
        except Exception, e:
            log.err("Redis failed to set('%s', '%s'): %s" %
                    (key, value, str(e)))
            raise cyclone.web.HTTPError(503)

        self.set_header("Content-Type", "text/plain")
        self.finish("%s=%s\r\n" % (key, value))

    @defer.inlineCallbacks
    def delete(self, key):
        try:
            n = yield self.dbconn.delete(key)
        except Exception, e:
            log.err("Redis failed to del('%s'): %s" % (key, str(e)))
            raise cyclone.web.HTTPError(503)

        self.set_header("Content-Type", "text/plain")
        self.finish("DEL %s=%d\r\n" % (key, n))

# GET will subscribe to channels or patterns
# POST will (obviously) post messages to channels
class QueueHandler(cyclone.web.RequestHandler, RedisMixin):
    @cyclone.web.asynchronous
    def get(self, channels):
        try:
            channels = channels.split(",")
        except Exception, e:
            log.err("Could not split channel names: %s" % str(e))
            raise cyclone.web.HTTPError(400, str(e))

        self.set_header("Content-Type", "text/plain")
        self.notifyFinish().addCallback(
            functools.partial(RedisMixin.unsubscribe_all, self))

        for channel in channels:
            self.subscribe(channel)
            self.write("subscribed to %s\r\n" % channel)
        self.flush()

    @defer.inlineCallbacks
    def post(self, channel):
        message = self.get_argument("message")

        try:
            n = yield self.dbconn.publish(channel, message.encode("utf-8"))
        except Exception, e:
            log.msg("Redis failed to publish('%s', '%s'): %s" %
                    (channel, repr(message), str(e)))
            raise cyclone.web.HTTPError(503)

        self.set_header("Content-Type", "text/plain")
        self.finish("OK %d\r\n" % n)

class QueueProtocol(cyclone.redis.SubscriberProtocol, RedisMixin):
    def messageReceived(self, pattern, channel, message):
        # When new messages are published to Redis channels or patterns,
        # they are broadcasted to all HTTP clients subscribed to those
        # channels.
        RedisMixin.broadcast(self, pattern, channel, message)

    def connectionMade(self):
        RedisMixin.psconn = self

        # If we lost connection with Redis during operation, we
        # re-subscribe to all channels once the connection is re-established.
        for channel in self.channels:
            if "*" in channel:
                self.psubscribe(channel)
            else:
                self.subscribe(channel)

    def connectionLost(self, why):
        RedisMixin.psconn = None

def main():
    log.startLogging(sys.stdout)
    reactor.listenTCP(8888, Application(), interface="127.0.0.1")
    reactor.run()

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


<commit_before>from __future__ import print_function
import os, sys
sys.path.append('..')
import base.prepare
from base.prepare import prepare
from base.utils import fix_names
from base.titer_model import TiterModel
from datetime import datetime, timedelta, date
import argparse
from dengue_subsampling import dengue_subsampling
import numpy as np

def collect_args():
    """Returns a dengue-specific argument parser.
    """
    parser = base.prepare.collect_args()

    parser.add_argument('-s', '--serotypes', '--lineage', choices=["all", "denv1", "denv2", "denv3", "denv4", "multiple"], default="multiple", type=str, nargs='+', help="Serotype(s) to prepare. \"multiple\" will run them all.")
    parser.add_argument('-y', '--years_back', type=int, default=100, help="Years back in time to sample from")
    parser.add_argument('--titers', default='../../fauna/data/dengue_titers.tsv', help="tab-delimited file of titer strains and values from fauna (e.g., dengue_titers.tsv)")
    parser.set_defaults(
        viruses_per_month=3,
        file_prefix="dengue"
    )

    return parser

dropped_strains = [
    'DENV1/VIETNAM/BIDV992/2006', 'DENV1/FRANCE/00475/2008', 'DENV1/VIETNAM/BIDV3990/2008', 'DENV2/HAITI/DENGUEVIRUS2HOMOSAPIENS1/2016', # Probable recombinants
    'DENV2/AUSTRALIA/QML22/2015' # Suspiciously far diverged
]

references = {
    "denv1": {"metadata": {'strain': "DENV1/NAURUISLAND/REFERENCE/1997", "accession": "NC_001477", "date": "1997-XX-XX", 'host': "NA", 'country': "Nauru", 'region': "oceania"}},
    "denv2": {"metadata": {'strain': "DENV2/THAILAND/REFERENCE/1964", "accession": "NC_001474", "date": "1964-XX-XX", 'host': "NA", 'country': "Thailand", "region": "southeast_asia"}},
    "denv3": {"metadata": {'strain': "DENV3/SRI_LANKA/REFERENCE/2000", "accession": "NC_001475", "date": "2000-XX-XX", 'host': "NA", 'country': "Sri Lanka", "region": "south_asia"}},
    "denv4": {"metadata": {'strain': "DENV4/NA/REFERENCE/2003", "accession": "NC_002640", "date": "2003-XX-XX", 'host': "NA", 'country': "NA"}},
}

for key in references:
    references[key]["genes"] = ['C', 'M', 'E', 'NS1', 'NS2A', 'NS2B', 'NS3', 'NS4A', '2K', 'NS4B', 'NS5']
    references[key]["path"] = "metadata/dengue_{}_outgroup.gb".format(key)
    references[key]["include"] = 0
references["all"] = references["denv4"]

##### Utils #####
def select_serotype(infile, path, serotype):
    '''
    Takes all-serotype fasta file, path to save output, and desired serotype.
    Writes appropriate subset of sequences to dengue_serotype.fasta
    Returns path to output file as string.
    '''
    from Bio import SeqIO
    sequences = [ i for i in SeqIO.parse(infile, 'fasta') if i.description.split('/')[0] == serotype.upper() ]
    SeqIO.write(sequences, path+'dengue_%s.fasta'%serotype, 'fasta')
    return path+'dengue_%s.fasta'%serotype

##### Set up config #####
# For dengue, config is a function so it is applicable for multiple lineages
def make_config(serotype, params):
    if params.file_prefix is not None:
        file_prefix = params.file_prefix
    else:
        file_prefix = "dengue_%s" % serotype

    if params.sequences is not None:
        input_paths = params.sequences
    elif os.path.isfile("../../fauna/data/dengue_%s.fasta"%serotype): #is file: # Look for a serotype-specific fasta
        input_paths = ["../../fauna/data/dengue_%s.fasta"%serotype]
    else: # If it doesn't exist, try to pull serotype-specific sequences out of the all-serotype fasta (warn the user of this behavior)
        input_paths = select_serotype('../fauna/data/dengue_all.fasta', '../fauna/data/', serotype)
        print('WARNING: Did not find serotype-specific fasta file.\nPulled sequences with serotype %s from all-serotype fasta file %s\nWrote these to file %s'%(serotype, '../fauna/data/dengue.fasta', input_paths))

    years_back = params.years_back
    time_interval = [datetime.today().date(), (datetime.today()  - timedelta(days=365.25 * years_back)).date()]

    if params.titers is not None:
        if not os.path.isfile(params.titers):
            params.titers = '../../fauna/data/%s'%params.titers
        titer_values, strains, sources = TiterModel.load_from_file(params.titers)
    else:
        titer_values, strains, sources = None, None, None

    config = {
        "dir": "dengue",
        "lineage": serotype,
        "file_prefix": file_prefix,
        "input_paths": input_paths,
        "header_fields": {0:'strain', 1:'accession', 2:'date', 3:'region', 4:'country',
                        5:'division', 6: 'location', 7: 'authors', 8: 'url'},
        "filters": (("Dropped Strains", lambda s: s.id not in [fix_names(x) for x in dropped_strains]),
            # ("Sequence Length", lambda s: len(s.seq)>=5000),
            ("Bad Region", lambda s: s.attributes['region'] not in ['', '?']) ),

        "subsample": dengue_subsampling(params, years_back, titer_values),

        # "subsample": {
        #     "category": lambda x: (x.attributes['region'],
        #                           x.attributes['date'].year,
        #                           x.attributes['date'].month),
        #     "priority": priority,
        #     "threshold": 3 },

        "add_urls": {
            "prefix": "https://www.ncbi.nlm.nih.gov/nuccore/%s",
            "attr": "accession" },

        "colors": ["region"],
        "color_defs": ["./colors.tsv"],
        "lat_longs": ["region"],
        "lat_long_defs": '../../fauna/source-data/geo_lat_long.tsv',
        "reference": references[serotype],
        "time_interval": time_interval,
        "titers": titer_values,
        "sources": sources
    }
    return config

if __name__=="__main__":
    parser = collect_args()
    params = parser.parse_args()

    if 'multiple' in params.serotypes:
        params.serotypes = ['denv1', 'denv2', 'denv3', 'denv4', 'all']

    # modify config file according to serotype
    for serotype in params.serotypes:
        print("Preparing serotype %s"%serotype)

        config = make_config(serotype, params)

        if params.viruses_per_month == 0:
            config["subsample"] = False
        else:
            config["subsample"]["threshold"] = params.viruses_per_month

        runner = prepare(config)
        runner.load_references()
        # runner.applyFilters()
        runner.subsample()
        runner.colors()
        runner.latlongs()
        runner.write_to_json()
<commit_msg>do not make titers required by default for dengue.<commit_after>from __future__ import print_function
import os, sys
sys.path.append('..')
import base.prepare
from base.prepare import prepare
from base.utils import fix_names
from base.titer_model import TiterModel
from datetime import datetime, timedelta, date
import argparse
from dengue_subsampling import dengue_subsampling
import numpy as np

def collect_args():
    """Returns a dengue-specific argument parser.
    """
    parser = base.prepare.collect_args()

    parser.add_argument('-s', '--serotypes', '--lineage', choices=["all", "denv1", "denv2", "denv3", "denv4", "multiple"], default="multiple", type=str, nargs='+', help="Serotype(s) to prepare. \"multiple\" will run them all.")
    parser.add_argument('-y', '--years_back', type=int, default=100, help="Years back in time to sample from")
    parser.add_argument('--titers', help="tab-delimited file of titer strains and values from fauna (e.g., dengue_titers.tsv)")
    parser.set_defaults(
        viruses_per_month=3,
        file_prefix="dengue"
    )

    return parser

dropped_strains = [
    'DENV1/VIETNAM/BIDV992/2006', 'DENV1/FRANCE/00475/2008', 'DENV1/VIETNAM/BIDV3990/2008', 'DENV2/HAITI/DENGUEVIRUS2HOMOSAPIENS1/2016', # Probable recombinants
    'DENV2/AUSTRALIA/QML22/2015' # Suspiciously far diverged
]

references = {
    "denv1": {"metadata": {'strain': "DENV1/NAURUISLAND/REFERENCE/1997", "accession": "NC_001477", "date": "1997-XX-XX", 'host': "NA", 'country': "Nauru", 'region': "oceania"}},
    "denv2": {"metadata": {'strain': "DENV2/THAILAND/REFERENCE/1964", "accession": "NC_001474", "date": "1964-XX-XX", 'host': "NA", 'country': "Thailand", "region": "southeast_asia"}},
    "denv3": {"metadata": {'strain': "DENV3/SRI_LANKA/REFERENCE/2000", "accession": "NC_001475", "date": "2000-XX-XX", 'host': "NA", 'country': "Sri Lanka", "region": "south_asia"}},
    "denv4": {"metadata": {'strain': "DENV4/NA/REFERENCE/2003", "accession": "NC_002640", "date": "2003-XX-XX", 'host': "NA", 'country': "NA"}},
}

for key in references:
    references[key]["genes"] = ['C', 'M', 'E', 'NS1', 'NS2A', 'NS2B', 'NS3', 'NS4A', '2K', 'NS4B', 'NS5']
    references[key]["path"] = "metadata/dengue_{}_outgroup.gb".format(key)
    references[key]["include"] = 0
references["all"] = references["denv4"]

##### Utils #####
def select_serotype(infile, path, serotype):
    '''
    Takes all-serotype fasta file, path to save output, and desired serotype.
    Writes appropriate subset of sequences to dengue_serotype.fasta
    Returns path to output file as string.
    '''
    from Bio import SeqIO
    sequences = [ i for i in SeqIO.parse(infile, 'fasta') if i.description.split('/')[0] == serotype.upper() ]
    SeqIO.write(sequences, path+'dengue_%s.fasta'%serotype, 'fasta')
    return path+'dengue_%s.fasta'%serotype

##### Set up config #####
# For dengue, config is a function so it is applicable for multiple lineages
def make_config(serotype, params):
    if params.file_prefix is not None:
        file_prefix = params.file_prefix
    else:
        file_prefix = "dengue_%s" % serotype

    if params.sequences is not None:
        input_paths = params.sequences
    elif os.path.isfile("../../fauna/data/dengue_%s.fasta"%serotype): #is file: # Look for a serotype-specific fasta
        input_paths = ["../../fauna/data/dengue_%s.fasta"%serotype]
    else: # If it doesn't exist, try to pull serotype-specific sequences out of the all-serotype fasta (warn the user of this behavior)
        input_paths = select_serotype('../fauna/data/dengue_all.fasta', '../fauna/data/', serotype)
        print('WARNING: Did not find serotype-specific fasta file.\nPulled sequences with serotype %s from all-serotype fasta file %s\nWrote these to file %s'%(serotype, '../fauna/data/dengue.fasta', input_paths))

    years_back = params.years_back
    time_interval = [datetime.today().date(), (datetime.today()  - timedelta(days=365.25 * years_back)).date()]

    if params.titers is not None:
        if not os.path.isfile(params.titers):
            params.titers = '../../fauna/data/%s'%params.titers
        titer_values, strains, sources = TiterModel.load_from_file(params.titers)
    else:
        titer_values, strains, sources = None, None, None

    config = {
        "dir": "dengue",
        "lineage": serotype,
        "file_prefix": file_prefix,
        "input_paths": input_paths,
        "header_fields": {0:'strain', 1:'accession', 2:'date', 3:'region', 4:'country',
                        5:'division', 6: 'location', 7: 'authors', 8: 'url'},
        "filters": (("Dropped Strains", lambda s: s.id not in [fix_names(x) for x in dropped_strains]),
            # ("Sequence Length", lambda s: len(s.seq)>=5000),
            ("Bad Region", lambda s: s.attributes['region'] not in ['', '?']) ),

        "subsample": dengue_subsampling(params, years_back, titer_values),

        # "subsample": {
        #     "category": lambda x: (x.attributes['region'],
        #                           x.attributes['date'].year,
        #                           x.attributes['date'].month),
        #     "priority": priority,
        #     "threshold": 3 },

        "add_urls": {
            "prefix": "https://www.ncbi.nlm.nih.gov/nuccore/%s",
            "attr": "accession" },

        "colors": ["region"],
        "color_defs": ["./colors.tsv"],
        "lat_longs": ["region"],
        "lat_long_defs": '../../fauna/source-data/geo_lat_long.tsv',
        "reference": references[serotype],
        "time_interval": time_interval,
        "titers": titer_values,
        "sources": sources
    }
    return config

if __name__=="__main__":
    parser = collect_args()
    params = parser.parse_args()

    if 'multiple' in params.serotypes:
        params.serotypes = ['denv1', 'denv2', 'denv3', 'denv4', 'all']

    # modify config file according to serotype
    for serotype in params.serotypes:
        print("Preparing serotype %s"%serotype)

        config = make_config(serotype, params)

        if params.viruses_per_month == 0:
            config["subsample"] = False
        else:
            config["subsample"]["threshold"] = params.viruses_per_month

        runner = prepare(config)
        runner.load_references()
        # runner.applyFilters()
        runner.subsample()
        runner.colors()
        runner.latlongs()
        runner.write_to_json()
<|endoftext|>


<commit_before>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six
from jsonschema import validate as js_validate
import warnings
import uuid
import time as ttime
import pandas as pd
from ..utils import sanitize_np, apply_to_dict_recursively

class DatumNotFound(Exception):
    """
    Raised if a Datum id is not found.
    """
    def __init__(self, datum_id, *args, **kwargs):
        super(*args, **kwargs)
        self.datum_id = datum_id

    def __str__(self):
        return f"No datum found with datum id {self.datum_id}"

class EventDatumNotFound(Exception):
    """
    Raised if an Event document is found to have an unknown Datum id.
    """
    def __init__(self, event_uid, datum_id, *args, **kwargs):
        super(*args, **kwargs)
        self.event_uid = event_uid
        self.datum_id = datum_id

    def __str__(self):
        return (
            f"Event with uid {self.event_uid} references "
            f"unknown Datum with datum id {self.datum_id}"
        )

def doc_or_uid_to_uid(doc_or_uid):
    """Given Document or uid return the uid

    Parameters
    ----------
    doc_or_uid : dict or str
        If str, then assume uid and pass through, if not, return
        the 'uid' field

    Returns
    -------
    uid : str
        A string version of the uid of the given document

    """
    if not isinstance(doc_or_uid, six.string_types):
        doc_or_uid = doc_or_uid['uid']
    return doc_or_uid

def _get_datum_from_datum_id(col, datum_id, datum_cache, logger):
    try:
        datum = datum_cache[datum_id]
    except KeyError:
        # find the current document
        edoc = col.find_one({'datum_id': datum_id})
        if edoc is None:
            raise DatumNotFound(datum_id=datum_id)
        # save it for later
        datum = dict(edoc)

        res = edoc['resource']
        count = 0
        for dd in col.find({'resource': res}):
            count += 1
            d_id = dd['datum_id']
            if d_id not in datum_cache:
                datum_cache[d_id] = dict(dd)
        if count > datum_cache.max_size:
            logger.warn("More datum in a resource than your "
                        "datum cache can hold.")

    datum.pop('_id', None)
    return datum

def retrieve(col, datum_id, datum_cache, get_spec_handler, logger):
    datum = _get_datum_from_datum_id(col, datum_id, datum_cache, logger)
    handler = get_spec_handler(datum['resource'])
    return handler(**datum['datum_kwargs'])

def resource_given_datum_id(col, datum_id, datum_cache, logger):
    datum_id = doc_or_uid_to_uid(datum_id)
    datum = _get_datum_from_datum_id(col, datum_id, datum_cache, logger)
    res = datum['resource']
    return res

def resource_given_uid(col, resource):
    uid = doc_or_uid_to_uid(resource)
    ret = col.find_one({'uid': uid})
    ret.pop('_id', None)
    ret['id'] = ret['uid']
    return ret

def bulk_insert_datum(col, resource, datum_ids,
                      datum_kwarg_list):

    resource_id = doc_or_uid_to_uid(resource)

    def datum_factory():
        for d_id, d_kwargs in zip(datum_ids, datum_kwarg_list):
            datum = dict(resource=resource_id,
                         datum_id=str(d_id),
                         datum_kwargs=dict(d_kwargs))
            apply_to_dict_recursively(datum, sanitize_np)
            yield datum

    col.insert(datum_factory())

def bulk_register_datum_table(datum_col,
                              resource_uid,
                              dkwargs_table,
                              validate):
    if validate:
        raise

    d_ids = [str(uuid.uuid4()) for j in range(len(dkwargs_table))]
    dkwargs_table = pd.DataFrame(dkwargs_table)
    bulk_insert_datum(datum_col, resource_uid, d_ids, [
        dict(r) for _, r in dkwargs_table.iterrows()])
    return d_ids

def register_datum(col, resource_uid, datum_kwargs):
    datum_uid = str(uuid.uuid4())
    datum = insert_datum(col, resource_uid, datum_uid, datum_kwargs, {}, None)
    return datum['datum_id']

def insert_datum(col, resource, datum_id, datum_kwargs, known_spec,
                 resource_col, ignore_duplicate_error=False,
                 duplicate_exc=None):
    if ignore_duplicate_error:
        assert duplicate_exc is not None

    if duplicate_exc is None:
        class _PrivateException(Exception):
            pass
        duplicate_exc = _PrivateException
    try:
        resource['spec']
        spec = resource['spec']

        if spec in known_spec:
            js_validate(datum_kwargs, known_spec[spec]['datum'])
    except (AttributeError, TypeError):
        pass

    resource_uid = doc_or_uid_to_uid(resource)

    datum = dict(resource=resource_uid,
                 datum_id=str(datum_id),
                 datum_kwargs=dict(datum_kwargs))
    apply_to_dict_recursively(datum, sanitize_np)
    # We are transitioning from ophyd objects inserting directly into a
    # Registry to ophyd objects passing documents to the RunEngine which in
    # turn inserts them into a Registry. During the transition period, we allow
    # an ophyd object to attempt BOTH so that configuration files are
    # compatible with both the new model and the old model. Thus, we need to
    # ignore the second attempt to insert.
    try:
        col.insert_one(datum)
    except duplicate_exc:
        if ignore_duplicate_error:
            warnings.warn("Ignoring attempt to insert Resource with duplicate "
                          "uid, assuming that both ophyd and bluesky "
                          "attempted to insert this document. Remove the "
                          "Registry (`reg` parameter) from your ophyd "
                          "instance to remove this warning.")
        else:
            raise
    # do not leak mongo objectID
    datum.pop('_id', None)

    return datum

def insert_resource(col, spec, resource_path, resource_kwargs,
                    known_spec, root, path_semantics='posix', uid=None,
                    run_start=None, id=None,
                    ignore_duplicate_error=False, duplicate_exc=None):
    """Insert resource into a databroker.

    Parameters
    ----------
    col : pymongo.Collection instance
        Collection to insert data into
    spec : str
        The resource data spec
    resource_path : str
        The path to the resource files
    resource_kwargs : dict
        The kwargs for the resource
    known_spec : set
        The known specs
    root : str
        The root of the file path
    path_semantics : str, optional
        The name of the path semantics, e.g. ``posix`` for Linux systems
    uid : str, optional
        The unique ID for the resource
    run_start : str, optional
        The unique ID for the start document the resource is associated with
    id : str, optional
        Dummy variable so that we round trip resources, same as ``uid``

    Returns
    -------
    resource_object : dict
        The resource
    """
    if ignore_duplicate_error:
        assert duplicate_exc is not None
    if duplicate_exc is None:
        class _PrivateException(Exception):
            pass
        duplicate_exc = _PrivateException
    resource_kwargs = dict(resource_kwargs)
    if spec in known_spec:
        js_validate(resource_kwargs, known_spec[spec]['resource'])
    if uid is None:
        uid = str(uuid.uuid4())

    resource_object = dict(spec=str(spec),
                           resource_path=str(resource_path),
                           root=str(root),
                           resource_kwargs=resource_kwargs,
                           path_semantics=path_semantics,
                           uid=uid)
    # This is special-cased because it was added later.
    # Someday this may be required and no longer special-cased.
    if run_start is not None:
        resource_object['run_start'] = run_start
    # We are transitioning from ophyd objects inserting directly into a
    # Registry to ophyd objects passing documents to the RunEngine which in
    # turn inserts them into a Registry. During the transition period, we allow
    # an ophyd object to attempt BOTH so that configuration files are
    # compatible with both the new model and the old model. Thus, we need to
    # ignore the second attempt to insert.
    try:
        col.insert_one(resource_object)
    except duplicate_exc:
        if ignore_duplicate_error:
            warnings.warn("Ignoring attempt to insert Datum with duplicate "
                          "datum_id, assuming that both ophyd and bluesky "
                          "attempted to insert this document. Remove the "
                          "Registry (`reg` parameter) from your ophyd "
                          "instance to remove this warning.")
        else:
            raise
    resource_object['id'] = resource_object['uid']
    resource_object.pop('_id', None)
    return resource_object

def update_resource(update_col, resource_col, old, new, cmd, cmd_kwargs):
    '''Update a resource document

    Parameters
    ----------
    update_col : Collection
        The collection to record audit trail in
    resource_col : Collection
        The resource collection

    old : dict
        The old resource document

    new : dict
        The new resource document

    cmd : str
        The name of the operation which generated this update

    cmd_kwargs : dict
        The arguments that went into the update (excluding the resource id)

    Returns
    -------
    ret : dict
        The new resource document

    log_object : dict
        The history object inserted (with oid removed)
    '''
    if old['uid'] != new['uid']:
        raise RuntimeError('must not change the resource uid')
    uid = old['uid']
    log_object = {'resource': uid,
                  'old': old,
                  'new': new,
                  'time': ttime.time(),
                  'cmd': cmd,
                  'cmd_kwargs': cmd_kwargs}

    update_col.insert_one(log_object)
    result = resource_col.replace_one({'uid': uid}, new)
    ret = resource_given_uid(resource_col, uid)
    # TODO look inside of result
    del result
    log_object.pop('_id', None)
    return ret, log_object

def get_resource_history(col, resource):
    uid = doc_or_uid_to_uid(resource)
    cursor = col.find({'resource': uid})
    for doc in cursor:
        for k in ['new', 'old']:
            d = doc[k]
            d.pop('_id', None)
            d['id'] = d['uid']
            doc[k] = d
        doc.pop('_id', None)
        yield doc

def get_datumkw_by_resuid_gen(datum_col, resource_uid):
    '''Given a resource uid, get all datum_kwargs

    No order is guaranteed.

    Internally the result of this is passed to the `get_file_list` method
    of the handler object in `change_root`

    Parameters
    ----------
    datum_col : Collection
        The Datum collection

    resource_uid : dict or str
       The resource to work on

    Yields
    ------
    datum_kwarg : dict
    '''
    resource_uid = doc_or_uid_to_uid(resource_uid)
    cur = datum_col.find({'resource': resource_uid})

    for d in cur:
        yield d['datum_kwargs']

def get_datum_by_res_gen(datum_col, resource_uid):
    '''Given a resource uid, get all datums

    No order is guaranteed.

    Internally the result of this is passed to the `get_file_list` method
    of the handler object in `change_root`

    Parameters
    ----------
    datum_col : Collection
        The Datum collection

    resource_uid : dict or str
       The resource to work on

    Yields
    ------
    datum : dict
    '''
    resource_uid = doc_or_uid_to_uid(resource_uid)
    cur = datum_col.find({'resource': resource_uid})

    for d in cur:
        yield d

def get_file_list(resource, datum_kwarg_gen, get_spec_handler):
    """
    Given a resource and an iterable of datum kwargs, get a list of
    associated files.

    DO NOT USE FOR COPYING OR MOVING. This is for debugging only.
    See the methods for moving and copying on the Registry object.
    """
    handler = get_spec_handler(resource['uid'])
    return handler.get_file_list(datum_kwarg_gen)
<commit_msg>use super() correctly<commit_after>from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import six
from jsonschema import validate as js_validate
import warnings
import uuid
import time as ttime
import pandas as pd
from ..utils import sanitize_np, apply_to_dict_recursively

class DatumNotFound(Exception):
    """
    Raised if a Datum id is not found.
    """
    def __init__(self, datum_id, *args):
        super().__init__(*args)
        self.datum_id = datum_id

    def __str__(self):
        return f"No datum found with datum id {self.datum_id}"

class EventDatumNotFound(Exception):
    """
    Raised if an Event document is found to have an unknown Datum id.
    """
    def __init__(self, event_uid, datum_id, *args):
        super().__init__(*args)
        self.event_uid = event_uid
        self.datum_id = datum_id

    def __str__(self):
        return (
            f"Event with uid {self.event_uid} references "
            f"unknown Datum with datum id {self.datum_id}"
        )

def doc_or_uid_to_uid(doc_or_uid):
    """Given Document or uid return the uid

    Parameters
    ----------
    doc_or_uid : dict or str
        If str, then assume uid and pass through, if not, return
        the 'uid' field

    Returns
    -------
    uid : str
        A string version of the uid of the given document

    """
    if not isinstance(doc_or_uid, six.string_types):
        doc_or_uid = doc_or_uid['uid']
    return doc_or_uid

def _get_datum_from_datum_id(col, datum_id, datum_cache, logger):
    try:
        datum = datum_cache[datum_id]
    except KeyError:
        # find the current document
        edoc = col.find_one({'datum_id': datum_id})
        if edoc is None:
            raise DatumNotFound(datum_id=datum_id)
        # save it for later
        datum = dict(edoc)

        res = edoc['resource']
        count = 0
        for dd in col.find({'resource': res}):
            count += 1
            d_id = dd['datum_id']
            if d_id not in datum_cache:
                datum_cache[d_id] = dict(dd)
        if count > datum_cache.max_size:
            logger.warn("More datum in a resource than your "
                        "datum cache can hold.")

    datum.pop('_id', None)
    return datum

def retrieve(col, datum_id, datum_cache, get_spec_handler, logger):
    datum = _get_datum_from_datum_id(col, datum_id, datum_cache, logger)
    handler = get_spec_handler(datum['resource'])
    return handler(**datum['datum_kwargs'])

def resource_given_datum_id(col, datum_id, datum_cache, logger):
    datum_id = doc_or_uid_to_uid(datum_id)
    datum = _get_datum_from_datum_id(col, datum_id, datum_cache, logger)
    res = datum['resource']
    return res

def resource_given_uid(col, resource):
    uid = doc_or_uid_to_uid(resource)
    ret = col.find_one({'uid': uid})
    ret.pop('_id', None)
    ret['id'] = ret['uid']
    return ret

def bulk_insert_datum(col, resource, datum_ids,
                      datum_kwarg_list):

    resource_id = doc_or_uid_to_uid(resource)

    def datum_factory():
        for d_id, d_kwargs in zip(datum_ids, datum_kwarg_list):
            datum = dict(resource=resource_id,
                         datum_id=str(d_id),
                         datum_kwargs=dict(d_kwargs))
            apply_to_dict_recursively(datum, sanitize_np)
            yield datum

    col.insert(datum_factory())

def bulk_register_datum_table(datum_col,
                              resource_uid,
                              dkwargs_table,
                              validate):
    if validate:
        raise

    d_ids = [str(uuid.uuid4()) for j in range(len(dkwargs_table))]
    dkwargs_table = pd.DataFrame(dkwargs_table)
    bulk_insert_datum(datum_col, resource_uid, d_ids, [
        dict(r) for _, r in dkwargs_table.iterrows()])
    return d_ids

def register_datum(col, resource_uid, datum_kwargs):
    datum_uid = str(uuid.uuid4())
    datum = insert_datum(col, resource_uid, datum_uid, datum_kwargs, {}, None)
    return datum['datum_id']

def insert_datum(col, resource, datum_id, datum_kwargs, known_spec,
                 resource_col, ignore_duplicate_error=False,
                 duplicate_exc=None):
    if ignore_duplicate_error:
        assert duplicate_exc is not None

    if duplicate_exc is None:
        class _PrivateException(Exception):
            pass
        duplicate_exc = _PrivateException
    try:
        resource['spec']
        spec = resource['spec']

        if spec in known_spec:
            js_validate(datum_kwargs, known_spec[spec]['datum'])
    except (AttributeError, TypeError):
        pass

    resource_uid = doc_or_uid_to_uid(resource)

    datum = dict(resource=resource_uid,
                 datum_id=str(datum_id),
                 datum_kwargs=dict(datum_kwargs))
    apply_to_dict_recursively(datum, sanitize_np)
    # We are transitioning from ophyd objects inserting directly into a
    # Registry to ophyd objects passing documents to the RunEngine which in
    # turn inserts them into a Registry. During the transition period, we allow
    # an ophyd object to attempt BOTH so that configuration files are
    # compatible with both the new model and the old model. Thus, we need to
    # ignore the second attempt to insert.
    try:
        col.insert_one(datum)
    except duplicate_exc:
        if ignore_duplicate_error:
            warnings.warn("Ignoring attempt to insert Resource with duplicate "
                          "uid, assuming that both ophyd and bluesky "
                          "attempted to insert this document. Remove the "
                          "Registry (`reg` parameter) from your ophyd "
                          "instance to remove this warning.")
        else:
            raise
    # do not leak mongo objectID
    datum.pop('_id', None)

    return datum

def insert_resource(col, spec, resource_path, resource_kwargs,
                    known_spec, root, path_semantics='posix', uid=None,
                    run_start=None, id=None,
                    ignore_duplicate_error=False, duplicate_exc=None):
    """Insert resource into a databroker.

    Parameters
    ----------
    col : pymongo.Collection instance
        Collection to insert data into
    spec : str
        The resource data spec
    resource_path : str
        The path to the resource files
    resource_kwargs : dict
        The kwargs for the resource
    known_spec : set
        The known specs
    root : str
        The root of the file path
    path_semantics : str, optional
        The name of the path semantics, e.g. ``posix`` for Linux systems
    uid : str, optional
        The unique ID for the resource
    run_start : str, optional
        The unique ID for the start document the resource is associated with
    id : str, optional
        Dummy variable so that we round trip resources, same as ``uid``

    Returns
    -------
    resource_object : dict
        The resource
    """
    if ignore_duplicate_error:
        assert duplicate_exc is not None
    if duplicate_exc is None:
        class _PrivateException(Exception):
            pass
        duplicate_exc = _PrivateException
    resource_kwargs = dict(resource_kwargs)
    if spec in known_spec:
        js_validate(resource_kwargs, known_spec[spec]['resource'])
    if uid is None:
        uid = str(uuid.uuid4())

    resource_object = dict(spec=str(spec),
                           resource_path=str(resource_path),
                           root=str(root),
                           resource_kwargs=resource_kwargs,
                           path_semantics=path_semantics,
                           uid=uid)
    # This is special-cased because it was added later.
    # Someday this may be required and no longer special-cased.
    if run_start is not None:
        resource_object['run_start'] = run_start
    # We are transitioning from ophyd objects inserting directly into a
    # Registry to ophyd objects passing documents to the RunEngine which in
    # turn inserts them into a Registry. During the transition period, we allow
    # an ophyd object to attempt BOTH so that configuration files are
    # compatible with both the new model and the old model. Thus, we need to
    # ignore the second attempt to insert.
    try:
        col.insert_one(resource_object)
    except duplicate_exc:
        if ignore_duplicate_error:
            warnings.warn("Ignoring attempt to insert Datum with duplicate "
                          "datum_id, assuming that both ophyd and bluesky "
                          "attempted to insert this document. Remove the "
                          "Registry (`reg` parameter) from your ophyd "
                          "instance to remove this warning.")
        else:
            raise
    resource_object['id'] = resource_object['uid']
    resource_object.pop('_id', None)
    return resource_object

def update_resource(update_col, resource_col, old, new, cmd, cmd_kwargs):
    '''Update a resource document

    Parameters
    ----------
    update_col : Collection
        The collection to record audit trail in
    resource_col : Collection
        The resource collection

    old : dict
        The old resource document

    new : dict
        The new resource document

    cmd : str
        The name of the operation which generated this update

    cmd_kwargs : dict
        The arguments that went into the update (excluding the resource id)

    Returns
    -------
    ret : dict
        The new resource document

    log_object : dict
        The history object inserted (with oid removed)
    '''
    if old['uid'] != new['uid']:
        raise RuntimeError('must not change the resource uid')
    uid = old['uid']
    log_object = {'resource': uid,
                  'old': old,
                  'new': new,
                  'time': ttime.time(),
                  'cmd': cmd,
                  'cmd_kwargs': cmd_kwargs}

    update_col.insert_one(log_object)
    result = resource_col.replace_one({'uid': uid}, new)
    ret = resource_given_uid(resource_col, uid)
    # TODO look inside of result
    del result
    log_object.pop('_id', None)
    return ret, log_object

def get_resource_history(col, resource):
    uid = doc_or_uid_to_uid(resource)
    cursor = col.find({'resource': uid})
    for doc in cursor:
        for k in ['new', 'old']:
            d = doc[k]
            d.pop('_id', None)
            d['id'] = d['uid']
            doc[k] = d
        doc.pop('_id', None)
        yield doc

def get_datumkw_by_resuid_gen(datum_col, resource_uid):
    '''Given a resource uid, get all datum_kwargs

    No order is guaranteed.

    Internally the result of this is passed to the `get_file_list` method
    of the handler object in `change_root`

    Parameters
    ----------
    datum_col : Collection
        The Datum collection

    resource_uid : dict or str
       The resource to work on

    Yields
    ------
    datum_kwarg : dict
    '''
    resource_uid = doc_or_uid_to_uid(resource_uid)
    cur = datum_col.find({'resource': resource_uid})

    for d in cur:
        yield d['datum_kwargs']

def get_datum_by_res_gen(datum_col, resource_uid):
    '''Given a resource uid, get all datums

    No order is guaranteed.

    Internally the result of this is passed to the `get_file_list` method
    of the handler object in `change_root`

    Parameters
    ----------
    datum_col : Collection
        The Datum collection

    resource_uid : dict or str
       The resource to work on

    Yields
    ------
    datum : dict
    '''
    resource_uid = doc_or_uid_to_uid(resource_uid)
    cur = datum_col.find({'resource': resource_uid})

    for d in cur:
        yield d

def get_file_list(resource, datum_kwarg_gen, get_spec_handler):
    """
    Given a resource and an iterable of datum kwargs, get a list of
    associated files.

    DO NOT USE FOR COPYING OR MOVING. This is for debugging only.
    See the methods for moving and copying on the Registry object.
    """
    handler = get_spec_handler(resource['uid'])
    return handler.get_file_list(datum_kwarg_gen)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""Displays GPU name, temperature and memory usage.

Parameters:
   * nvidiagpu.format: Format string (defaults to '{name}: {temp}°C %{usedmem}/{totalmem} MiB')
     Available values are: {name} {temp} {mem_used} {mem_total} {fanspeed} {clock_gpu} {clock_mem}

Requires nvidia-smi

contributed by `RileyRedpath <https://github.com/RileyRedpath>`_ - many thanks!
"""

import core.module
import core.widget

import util.cli
import util.format

class Module(core.module.Module):
    def __init__(self, config, theme):
        super().__init__(config, theme, core.widget.Widget(self.utilization))

        self.__utilization = "Not found: 0 0/0"

    def utilization(self, widget):
        return self.__utilization

    def hidden(self):
        return "not found" in self.__utilization

    def update(self):
        sp = util.cli.execute("nvidia-smi -q", ignore_errors=True)

        title = ""
        usedMem = ""
        totalMem = ""
        temp = ""
        name = "not found"
        clockMem = ""
        clockGpu = ""
        fanspeed = ""
        for item in sp.split("\n"):
            try:
                key, val = item.split(":")
                key, val = key.strip(), val.strip()
                if title == "Clocks":
                    if key == "Graphics":
                        clockGpu = val.split(" ")[0]
                    elif key == "Memory":
                        clockMem = val.split(" ")[0]
                if title == "FB Memory Usage":
                    if key == "Total":
                        totalMem = val.split(" ")[0]
                    elif key == "Used":
                        usedMem = val.split(" ")[0]
                elif key == "GPU Current Temp":
                    temp = val.split(" ")[0]
                elif key == "Product Name":
                    name = val
                elif key == "Fan Speed":
                    fanspeed = val.split(" ")[0]

            except:
                title = item.strip()

        str_format = self.parameter(
            "format", "{name}: {temp}°C {mem_used}/{mem_total} MiB"
        )
        self.__utilization = str_format.format(
            name=name,
            temp=temp,
            mem_used=usedMem,
            mem_total=totalMem,
            clock_gpu=clockGpu,
            clock_mem=clockMem,
            fanspeed=fanspeed,
        )

# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4
<commit_msg>feat: add GPU usage % and GPU memory usage % to nvidiagpu<commit_after># -*- coding: utf-8 -*-

"""Displays GPU name, temperature and memory usage.

Parameters:
   * nvidiagpu.format: Format string (defaults to '{name}: {temp}°C %{usedmem}/{totalmem} MiB')
     Available values are: {name} {temp} {mem_used} {mem_total} {fanspeed} {clock_gpu} {clock_mem} {gpu_usage_pct} {mem_usage_pct}

Requires nvidia-smi

contributed by `RileyRedpath <https://github.com/RileyRedpath>`_ - many thanks!
"""

import core.module
import core.widget

import util.cli
import util.format

class Module(core.module.Module):
    def __init__(self, config, theme):
        super().__init__(config, theme, core.widget.Widget(self.utilization))

        self.__utilization = "Not found: 0 0/0"

    def utilization(self, widget):
        return self.__utilization

    def hidden(self):
        return "not found" in self.__utilization

    def update(self):
        sp = util.cli.execute("nvidia-smi -q", ignore_errors=True)

        title = ""
        usedMem = ""
        totalMem = ""
        temp = ""
        name = "not found"
        clockMem = ""
        clockGpu = ""
        fanspeed = ""
        gpuUsagePct = ""
        memPct = ""
        for item in sp.split("\n"):
            try:
                key, val = item.split(":")
                key, val = key.strip(), val.strip()
                if title == "Clocks":
                    if key == "Graphics":
                        clockGpu = val.split(" ")[0]
                    elif key == "Memory":
                        clockMem = val.split(" ")[0]
                if title == "FB Memory Usage":
                    if key == "Total":
                        totalMem = val.split(" ")[0]
                    elif key == "Used":
                        usedMem = val.split(" ")[0]
                elif key == "GPU Current Temp":
                    temp = val.split(" ")[0]
                elif key == "Product Name":
                    name = val
                elif key == "Fan Speed":
                    fanspeed = val.split(" ")[0]
                elif title == "Utilization":
                    if key == "Gpu":
                        gpuUsagePct = val.split(" ")[0]
                    elif key == "Memory":
                        memPct = val.split(" ")[0]

            except:
                title = item.strip()

        str_format = self.parameter(
            "format", "{name}: {temp}°C {mem_used}/{mem_total} MiB"
        )
        self.__utilization = str_format.format(
            name=name,
            temp=temp,
            mem_used=usedMem,
            mem_total=totalMem,
            clock_gpu=clockGpu,
            clock_mem=clockMem,
            fanspeed=fanspeed,
            gpu_usage_pct=gpuUsagePct,
            mem_usage_pct=memPct,
        )

# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4
<|endoftext|>


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

#   Copyright 2014-2018 PUNCH Cyber Analytics Group
#
#   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.
"""
Overview
========

Parse SMTP sessions

"""

from configparser import ConfigParser
from typing import Dict, Optional
import pyzmail
from bs4 import UnicodeDammit

from stoq import (
    Payload, RequestMeta, WorkerResponse,
    DispatcherResponse, ExtractedPayload, PayloadMeta)
from stoq.plugins import WorkerPlugin

class SMTPPlugin(WorkerPlugin):

    def __init__(self, config: ConfigParser,
                 plugin_opts: Optional[Dict]) -> None:
        super().__init__(config, plugin_opts)

        if plugin_opts and "omit_body" in plugin_opts:
            self.omit_body = plugin_opts["omit_body"]
        elif config.has_option("options", "omit_body"):
            self.omit_body = config.getboolean("options", "omit_body")
        else:
            self.omit_body = False

    def scan(
            self,
            payload: Payload,
            request_meta: RequestMeta,
    ) -> WorkerResponse:

        message_json = {}
        attachments = []
        errors = []
        email_session = UnicodeDammit(payload.content).unicode_markup
        message = pyzmail.message_from_string(email_session)

        # Create a dict of the SMTP headers
        for header in message.keys():
            curr_header = header.lower()
            if curr_header in message_json:
                # If the header key already exists, let's join them
                message_json[curr_header] += f"\n{message.get_decoded_header(header)}"
            else:
                message_json[curr_header] = message.get_decoded_header(header)

        # Extract the e-mail body, to include HTML if available
        message_json['body'] = '' if message.text_part is None else UnicodeDammit(
            message.text_part.get_payload()).unicode_markup
        message_json['body_html'] = '' if message.html_part is None else UnicodeDammit(
            message.html_part.get_payload()).unicode_markup

        # Handle attachments
        for mailpart in message.mailparts:
            # Skip if the attachment is a body part
            if mailpart.is_body:
                continue
            try:
                attachment_meta = PayloadMeta(extra_data={
                    'charset': mailpart.charset,
                    'content-description': mailpart.part.get('Content-Description'),
                    'content-id': mailpart.content_id,
                    'disposition': mailpart.disposition,
                    'filename': mailpart.filename,
                    'type': mailpart.type
                    })
                attachment = ExtractedPayload(mailpart.get_payload(), attachment_meta)
                attachments.append(attachment)
            except Exception as err:
                errors.append(f'Failed extracting attachment: {err}')

        # Make sure we delete the body and body_html keys if they are to
        # be omitted
        if self.omit_body:
            message_json.pop('body', None)
            message_json.pop('body_html', None)

        return WorkerResponse(message_json, errors=errors, extracted=attachments)
<commit_msg>More minor cleanup for smtp plugin<commit_after>#!/usr/bin/env python3

#   Copyright 2014-2018 PUNCH Cyber Analytics Group
#
#   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.
"""
Overview
========

Parse SMTP sessions

"""

from configparser import ConfigParser
from typing import Dict, Optional
import pyzmail
from bs4 import UnicodeDammit

from stoq import (
    Payload, RequestMeta, WorkerResponse,
    DispatcherResponse, ExtractedPayload, PayloadMeta)
from stoq.plugins import WorkerPlugin

class SMTPPlugin(WorkerPlugin):

    def __init__(self, config: ConfigParser,
                 plugin_opts: Optional[Dict]) -> None:
        super().__init__(config, plugin_opts)

        if plugin_opts and 'omit_body' in plugin_opts:
            self.omit_body = plugin_opts['omit_body']
        elif config.has_option('options', 'omit_body'):
            self.omit_body = config.getboolean('options', 'omit_body')
        else:
            self.omit_body = False

    def scan(
            self,
            payload: Payload,
            request_meta: RequestMeta,
    ) -> WorkerResponse:

        message_json = {}
        attachments = []
        errors = []
        email_session = UnicodeDammit(payload.content).unicode_markup
        message = pyzmail.message_from_string(email_session)

        # Create a dict of the SMTP headers
        for header in message.keys():
            curr_header = header.lower()
            if curr_header in message_json:
                # If the header key already exists, let's join them
                message_json[curr_header] += f'\n{message.get_decoded_header(header)}'
            else:
                message_json[curr_header] = message.get_decoded_header(header)

        if not self.omit_body:
            # Extract the e-mail body, to include HTML if available
            message_json['body'] = '' if message.text_part is None else UnicodeDammit(
                message.text_part.get_payload()).unicode_markup
            message_json['body_html'] = '' if message.html_part is None else UnicodeDammit(
                message.html_part.get_payload()).unicode_markup

        # Handle attachments
        for mailpart in message.mailparts:
            # Skip if the attachment is a body part
            if mailpart.is_body:
                continue
            try:
                attachment_meta = PayloadMeta(extra_data={
                    'charset': mailpart.charset,
                    'content-description': mailpart.part.get('Content-Description'),
                    'content-id': mailpart.content_id,
                    'disposition': mailpart.disposition,
                    'filename': mailpart.filename if mailpart.filename else mailpart.sanitized_filename,
                    'type': mailpart.type
                    })
                attachment = ExtractedPayload(mailpart.get_payload(), attachment_meta)
                attachments.append(attachment)
            except Exception as err:
                errors.append(f'Failed extracting attachment: {err}')

        return WorkerResponse(message_json, errors=errors, extracted=attachments)
<|endoftext|>


<commit_before>from django.core.management.base import BaseCommand, CommandError

import django
import os
import datetime
from django.utils import timezone
from ...models import Course, CourseResult
from gensim import corpora, models, similarities
from collections import defaultdict
from nltk.corpus import stopwords

class Command(BaseCommand):
    help = 'Does some magical work'

    def handle(self, *args, **options):
        """ Do your work here """
        # self.stdout.write('There are {} things!'.format(Course.objects.count()))
        course_count = Course.objects.count()
        courses = Course.objects.all()

        documents = []
        documents_names = []
        document_courses = []

        # words not to take into consideration
        # parametar 1: remove keywords from decriptions
        stoplist = stopwords.words('english')
        stoplist = stoplist + ['is', 'how', 'or', 'to', 'of', 'the',
                        'in', 'for', 'on', 'will', 'a', 'advanced', 'an', 'and', 'are', 'as', 'be', 'by', 'course',
                        'with', 'some', 'student', 'students', 'systems', 'system', 'basic',
                        'this', 'knowledge', 'use', 'using', 'well', 'hours;', 'four']

        # parametar 1: remove keywords from names
        stoplist_names = stopwords.words('english')

        print "Created stop list."

        for course in courses:
            documents.append(course.description + course.name)
            documents_names.append(course.name)
            document_courses.append(course)

        print "Fetched course descriptions."
        texts = [[word.replace(".", "").lower() for word in document.split()
                  if word.replace(".", "").lower() not in stoplist]
                 for document in documents]

        # to add remove list for names
        texts_names = [[word.lower() for word in document_name.split()
                       if word.replace(".", "").lower() not in stoplist_names]
                       for document_name in documents_names]

        frequency = defaultdict(int)
        for text in texts:
            for token in text:
                frequency[token] += 1
        texts = [[token for token in text if frequency[token] > 1]
                 for text in texts]

        dictionary = corpora.Dictionary(texts)
        dictionary.filter_n_most_frequent(15)
        dictionary.save_as_text("dictionary.txt", sort_by_word=False)

        frequency_names = defaultdict(int)
        for text in texts_names:
            for token in text:
                frequency_names[token] += 1
        texts_names = [[token for token in text if frequency_names[token] > 1]
                       for text in texts_names]

        # save to fale keywords sorted by appereance
        dictionary_names = corpora.Dictionary(texts_names)

        print "Created a dictionary of word-bag."
        # doc2bow counts the number of occurences of each distinct word,
        # converts the word to its integer word id and returns the result
        # as a sparse vector
        corpus = [dictionary.doc2bow(text) for text in texts]

        #print(corpus)
        # parametar 2: num_topics
        lsi = models.LsiModel(corpus, id2word=dictionary, num_topics=120)

        #for i in lsi.show_topics():
        #    print i[0], i[1]

        # parametar 2: num_topics
        corpus_names = [dictionary_names.doc2bow(text) for text in texts_names]
        lsi_names = models.LsiModel(corpus_names, id2word=dictionary_names, num_topics=125)

        j = 0
        CourseResult.objects.all().delete()
        for course in courses:
            print "Working on course " + str(j + 1) + " of " + str(course_count)
            doc = course.description + course.name
            vec_bow = dictionary.doc2bow(doc.lower().split())

            doc_name = course.name
            vec_bow_name = dictionary_names.doc2bow(doc_name.lower().split())

            # convert the query to LSI space
            vec_lsi = lsi[vec_bow]
            index = similarities.MatrixSimilarity(lsi[corpus])

            vec_lsi_name = lsi_names[vec_bow_name]
            index_name = similarities.MatrixSimilarity(lsi_names[corpus_names])

            # perform a similarity query against the corpus
            sims = index[vec_lsi]
            sims = sorted(enumerate(sims), key=lambda item: -item[1])

            sims_names = index_name[vec_lsi_name]
            sims_names = sorted(enumerate(sims_names), key=lambda item: -item[1])

            j = j + 1
            i = 0
            for similarity in sims:
                #print similarity
                if i != 0:
                    # parametar 3: ratio descriptions vs names
                    CourseResult.objects.create(first_course_id=document_courses[sims[0][0]].id,
                                                second_course_id=document_courses[sims[i][0]].id,
                                                result=(1 * sims[i][1] + 0.0 * sims_names[i][1]))
                i = i + 1

                # print sims[1:11]
                # print document_courses[sims[1][0]].description
                # print document_courses[sims[2][0]].description
                # print document_courses[sims[3][0]].description

                #    for course_temp in courses:
                #        if course.id != course_temp.id and course_temp not in skip_list:
                #            print "compare"
                #    skip_list.append(course.id)
<commit_msg>course similarity without names<commit_after>from django.core.management.base import BaseCommand, CommandError

import django
import os
import datetime
from django.utils import timezone
from ...models import Course, CourseResult
from gensim import corpora, models, similarities
from collections import defaultdict
from nltk.corpus import stopwords

class Command(BaseCommand):
    help = 'Does some magical work'

    def handle(self, *args, **options):
        """ Do your work here """
        # self.stdout.write('There are {} things!'.format(Course.objects.count()))
        course_count = Course.objects.count()
        courses = Course.objects.all()

        documents = []
        documents_names = []
        document_courses = []

        # words not to take into consideration
        # parametar 1: remove keywords from decriptions
        stoplist = stopwords.words('english')
        #stoplist = stoplist + ['is', 'how', 'or', 'to', 'of', 'the',
        #                'in', 'for', 'on', 'will', 'a', 'advanced', 'an', 'and', 'are', 'as', 'be', 'by', 'course',
        #                'with', 'some', 'student', 'students', 'systems', 'system', 'basic',
        #                'this', 'knowledge', 'use', 'using', 'well', 'hours;', 'four']

        # parametar 1: remove keywords from names
        stoplist_names = stopwords.words('english')

        print "Created stop list."

        for course in courses:
            documents.append(course.description + course.name)
            documents_names.append(course.name)
            document_courses.append(course)

        print "Fetched course descriptions."
        texts = [[word.replace(".", "").lower() for word in document.split()
                  if word.replace(".", "").lower() not in stoplist]
                 for document in documents]

        # to add remove list for names
        #texts_names = [[word.lower() for word in document_name.split()
        #               if word.replace(".", "").lower() not in stoplist_names]
        #               for document_name in documents_names]

        frequency = defaultdict(int)
        for text in texts:
            for token in text:
                frequency[token] += 1
        texts = [[token for token in text if frequency[token] > 1]
                 for text in texts]

        dictionary = corpora.Dictionary(texts)
        #dictionary.filter_n_most_frequent(50)
        #dictionary.save_as_text("dictionary.txt", sort_by_word=False)

        #frequency_names = defaultdict(int)
        #for text in texts_names:
        #    for token in text:
        #        frequency_names[token] += 1
        #texts_names = [[token for token in text if frequency_names[token] > 1]
        #               for text in texts_names]

        # save to fale keywords sorted by appereance
        #dictionary_names = corpora.Dictionary(texts_names)

        print "Created a dictionary of word-bag."
        # doc2bow counts the number of occurences of each distinct word,
        # converts the word to its integer word id and returns the result
        # as a sparse vector
        corpus = [dictionary.doc2bow(text) for text in texts]

        #print(corpus)
        # parametar 2: num_topics
        lsi = models.LsiModel(corpus, id2word=dictionary, num_topics=120)

        #for i in lsi.show_topics():
        #    print i[0], i[1]

        # parametar 2: num_topics
        #corpus_names = [dictionary_names.doc2bow(text) for text in texts_names]
        #lsi_names = models.LsiModel(corpus_names, id2word=dictionary_names, num_topics=125)

        j = 0
        CourseResult.objects.all().delete()
        for course in courses:
            print "Working on course " + str(j + 1) + " of " + str(course_count)
            doc = course.description + course.name
            vec_bow = dictionary.doc2bow(doc.lower().split())

            #doc_name = course.name
            #vec_bow_name = dictionary_names.doc2bow(doc_name.lower().split())

            # convert the query to LSI space
            vec_lsi = lsi[vec_bow]
            index = similarities.MatrixSimilarity(lsi[corpus])

            #vec_lsi_name = lsi_names[vec_bow_name]
            #index_name = similarities.MatrixSimilarity(lsi_names[corpus_names])

            # perform a similarity query against the corpus
            sims = index[vec_lsi]
            sims = sorted(enumerate(sims), key=lambda item: -item[1])

            #sims_names = index_name[vec_lsi_name]
            #sims_names = sorted(enumerate(sims_names), key=lambda item: -item[1])

            j = j + 1
            i = 0
            for similarity in sims:
                #print similarity
                if i != 0:
                    # parametar 3: ratio descriptions vs names
                    CourseResult.objects.create(first_course_id=document_courses[sims[0][0]].id,
                                                second_course_id=document_courses[sims[i][0]].id,
                                                result=(1 * sims[i][1]))
                i = i + 1

                # print sims[1:11]
                # print document_courses[sims[1][0]].description
                # print document_courses[sims[2][0]].description
                # print document_courses[sims[3][0]].description

                #    for course_temp in courses:
                #        if course.id != course_temp.id and course_temp not in skip_list:
                #            print "compare"
                #    skip_list.append(course.id)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Base management command that provides common functionality for the other commands in this app.
"""
import sys
import codecs
import locale
import logging
import requests
from re import sub
from datetime import datetime
from email.utils import parsedate
from django.utils import timezone
from django.utils.six.moves import input
from django.utils.termcolors import colorize
from django.core.management.base import BaseCommand
logger = logging.getLogger(__name__)

class CalAccessCommand(BaseCommand):
    """
    Base management command that provides common functionality for the other commands in this app.
    """
    url = 'http://campaignfinance.cdn.sos.ca.gov/dbwebexport.zip'

    def handle(self, *args, **options):
        """
        Sets options common to all commands.

        Any command subclassing this object should implement its own
        handle method, as is standard in Django, and run this method
        via a super call to inherit its functionality.
        """
        # Set global options
        self.verbosity = options.get("verbosity")
        self.no_color = options.get("no_color")

        # Start the clock
        self.start_datetime = datetime.now()

    def get_download_metadata(self):
        """
        Returns a dict with metadata about the current CAL-ACCESS snapshot.
        """
        request = requests.head(self.url)
        last_modified = request.headers['last-modified']
        dt = datetime(*parsedate(last_modified)[:6])
        # content length is a string, need to convert
        try:
            # long int type is big enough for double the current size of the zip
            length = long(request.headers['content-length'])
        except NameError:
            # in py3, no long(), instead int will suffice
            length = int(request.headers['content-length'])
        return {
            # should prob not call int here, can this remain a string until writing to db?
            'content-length': length,
            'last-modified': timezone.utc.localize(dt)
        }

    #
    # Logging methods
    #

    def header(self, string):
        """
        Writes out a string to stdout formatted to look like a header.
        """
        logger.debug(string)
        if not self.no_color:
            string = colorize(string, fg="cyan", opts=("bold",))
        self.stdout.write(string)

    def log(self, string):
        """
        Writes out a string to stdout formatted to look like a standard line.
        """
        logger.debug(string)
        if not self.no_color:
            string = colorize("%s" % string, fg="white")
        self.stdout.write(string)

    def success(self, string):
        """
        Writes out a string to stdout formatted green to communicate success.
        """
        logger.debug(string)
        if not self.no_color:
            string = colorize(string, fg="green")
        self.stdout.write(string)

    def failure(self, string):
        """
        Writes out a string to stdout formatted red to communicate failure.
        """
        logger.debug(string)
        if not self.no_color:
            string = colorize(string, fg="red")
        self.stdout.write(string)

    def duration(self):
        """
        Calculates how long the command has been running and writes it to stdout.
        """
        duration = datetime.now() - self.start_datetime
        self.stdout.write('Duration: {}'.format(str(duration)))

    def confirm_proceed(self, prompt):
        """
        Prompts the user for yes/no confirmation to proceed.
        """
        # Ensure stdout can handle Unicode data: http://bit.ly/1C3l4eV
        locale_encoding = locale.getpreferredencoding()
        old_stdout = sys.stdout
        sys.stdout = codecs.getwriter(locale_encoding)(sys.stdout)

        # Send the confirmation prompt out to the user
        user_input = input(prompt)

        confirm = None

        while confirm is None:
            if user_input.lower() in ['y', 'yes']:
                confirm = True
            elif user_input.lower() in ['n', 'no']:
                confirm = False
            else:
                user_input = input("Invalid input. Please type 'yes', 'no', 'y' or 'n':\n")

        # Set things back to the way they were before continuing.
        sys.stdout = old_stdout

        # Pass back what the user typed
        return confirm

    def __str__(self):
        return sub(r'(.+\.)*', '', self.__class__.__module__)
<commit_msg>Log duration<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Base management command that provides common functionality for the other commands in this app.
"""
import sys
import codecs
import locale
import logging
import requests
from re import sub
from datetime import datetime
from email.utils import parsedate
from django.utils import timezone
from django.utils.six.moves import input
from django.utils.termcolors import colorize
from django.core.management.base import BaseCommand
logger = logging.getLogger(__name__)

class CalAccessCommand(BaseCommand):
    """
    Base management command that provides common functionality for the other commands in this app.
    """
    url = 'http://campaignfinance.cdn.sos.ca.gov/dbwebexport.zip'

    def handle(self, *args, **options):
        """
        Sets options common to all commands.

        Any command subclassing this object should implement its own
        handle method, as is standard in Django, and run this method
        via a super call to inherit its functionality.
        """
        # Set global options
        self.verbosity = options.get("verbosity")
        self.no_color = options.get("no_color")

        # Start the clock
        self.start_datetime = datetime.now()

    def get_download_metadata(self):
        """
        Returns a dict with metadata about the current CAL-ACCESS snapshot.
        """
        request = requests.head(self.url)
        last_modified = request.headers['last-modified']
        dt = datetime(*parsedate(last_modified)[:6])
        # content length is a string, need to convert
        try:
            # long int type is big enough for double the current size of the zip
            length = long(request.headers['content-length'])
        except NameError:
            # in py3, no long(), instead int will suffice
            length = int(request.headers['content-length'])
        return {
            # should prob not call int here, can this remain a string until writing to db?
            'content-length': length,
            'last-modified': timezone.utc.localize(dt)
        }

    #
    # Logging methods
    #

    def header(self, string):
        """
        Writes out a string to stdout formatted to look like a header.
        """
        logger.debug(string)
        if not self.no_color:
            string = colorize(string, fg="cyan", opts=("bold",))
        self.stdout.write(string)

    def log(self, string):
        """
        Writes out a string to stdout formatted to look like a standard line.
        """
        logger.debug(string)
        if not self.no_color:
            string = colorize("%s" % string, fg="white")
        self.stdout.write(string)

    def success(self, string):
        """
        Writes out a string to stdout formatted green to communicate success.
        """
        logger.debug(string)
        if not self.no_color:
            string = colorize(string, fg="green")
        self.stdout.write(string)

    def failure(self, string):
        """
        Writes out a string to stdout formatted red to communicate failure.
        """
        logger.debug(string)
        if not self.no_color:
            string = colorize(string, fg="red")
        self.stdout.write(string)

    def duration(self):
        """
        Calculates how long the command has been running and writes it to stdout.
        """
        logger.debug(string)
        duration = datetime.now() - self.start_datetime
        self.stdout.write('Duration: {}'.format(str(duration)))

    def confirm_proceed(self, prompt):
        """
        Prompts the user for yes/no confirmation to proceed.
        """
        # Ensure stdout can handle Unicode data: http://bit.ly/1C3l4eV
        locale_encoding = locale.getpreferredencoding()
        old_stdout = sys.stdout
        sys.stdout = codecs.getwriter(locale_encoding)(sys.stdout)

        # Send the confirmation prompt out to the user
        user_input = input(prompt)

        confirm = None

        while confirm is None:
            if user_input.lower() in ['y', 'yes']:
                confirm = True
            elif user_input.lower() in ['n', 'no']:
                confirm = False
            else:
                user_input = input("Invalid input. Please type 'yes', 'no', 'y' or 'n':\n")

        # Set things back to the way they were before continuing.
        sys.stdout = old_stdout

        # Pass back what the user typed
        return confirm

    def __str__(self):
        return sub(r'(.+\.)*', '', self.__class__.__module__)
<|endoftext|>


<commit_before>from pipes import quote
import subprocess
import yaml

config = yaml.load(open("settings.yml"))

def runSSH(ip, command):
    command = 'ssh -i %s %s@%s %s' % (config["private_key_path"],
            config["remote_server_username"],
            ip, " ".join(map(quote, command)))
    return subprocess.check_output(command, shell=True)

def pushConfig(text, instances):
    f = open("master_address", "w")
    f.write(text)
    f.close()

    print "Written master address"
    for (externalIP, internalIP) in instances:
        runSSH(externalIP, ['sudo', 'mkdir', '-p', '/etc/flocker'])

        f = open("my_address", "w")
        f.write(externalIP)
        f.close()

        for f in ('master_address', 'my_address'):
            scp = ("scp -i %(private_key_path)s %(filename)s "
                   "%(remote_server_username)s@%(external_ip)s:/tmp/%(filename)s") % dict(
                        private_key_path=config["private_key_path"],
                        remote_server_username=config["remote_server_username"],
                        external_ip=externalIP, filename=f)
            subprocess.check_output(scp, shell=True)
            runSSH(externalIP, ['sudo', 'mv', '/tmp/%s' % (f,), '/etc/flocker/%s' % (f,)])
            print "Pushed f to %s" % (externalIP,)
    print "Finished telling all nodes about the master."
<commit_msg>mention which file you are pushing<commit_after>from pipes import quote
import subprocess
import yaml

config = yaml.load(open("settings.yml"))

def runSSH(ip, command):
    command = 'ssh -i %s %s@%s %s' % (config["private_key_path"],
            config["remote_server_username"],
            ip, " ".join(map(quote, command)))
    return subprocess.check_output(command, shell=True)

def pushConfig(text, instances):
    f = open("master_address", "w")
    f.write(text)
    f.close()

    print "Written master address"
    for (externalIP, internalIP) in instances:
        runSSH(externalIP, ['sudo', 'mkdir', '-p', '/etc/flocker'])

        f = open("my_address", "w")
        f.write(externalIP)
        f.close()

        for f in ('master_address', 'my_address'):
            scp = ("scp -i %(private_key_path)s %(filename)s "
                   "%(remote_server_username)s@%(external_ip)s:/tmp/%(filename)s") % dict(
                        private_key_path=config["private_key_path"],
                        remote_server_username=config["remote_server_username"],
                        external_ip=externalIP, filename=f)
            subprocess.check_output(scp, shell=True)
            runSSH(externalIP, ['sudo', 'mv', '/tmp/%s' % (f,), '/etc/flocker/%s' % (f,)])
            print "Pushed", f, "to", externalIP
    print "Finished telling all nodes about the master."
<|endoftext|>