<commit_before>from __future__ import print_function

import sys
from re import sub
import HTMLParser

from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.sql.functions import col, udf, array_contains, explode, desc
from pyspark.sql.types import StringType

def guessFormat(path, default="json"):
    if path.endswith(".json"):
        return ("json", {})
    elif path.endswith(".parquet"):
        return ("parquet", {})
    elif path.endswith(".csv"):
        return ("com.databricks.spark.csv", {'header': 'true'})
    else:
        return (default, {})

def formatURL(url, corpus, id, pages, regions):
    if corpus == 'ca':
        r = regions[0]
        return "%s/print/image_600x600_from_%d%%2C%d_to_%d%%2C%d/" \
            % (url, r.x/3, r.y/3, (r.x + r.w)/3, (r.y + r.h)/3)
    elif corpus == 'onb':
        return "%s&seite=%s" % (sub("&", "&", url), pages[0])
    elif corpus == 'trove':
        return "http://trove.nla.gov.au/ndp/del/article/%s" % sub("^trove/", "", id)
    else:
        return url

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: pretty-cluster.py <metadata> <input> <output> [<query>]", file=sys.stderr)
        exit(-1)
    sc = SparkContext(appName="Prettyprint Clusters")
    sqlContext = SQLContext(sc)

    outpath = sys.argv[3]
    (outputFormat, outputOptions) = guessFormat(outpath, "json")

    ## Should do more field renaming in meta to avoid clashing with fields in raw.
    meta = sqlContext.read.json(sys.argv[1])\
           .withColumnRenamed('lang', 'series_lang')\
           .withColumnRenamed('id', 'series').dropDuplicates(['series'])

    df = sqlContext.read.load(sys.argv[2])

    if len(sys.argv) >= 5:
        df.registerTempTable("clusters")
        print(sys.argv[4])
        raw = df.join(df.filter(sys.argv[4]).select("cluster").distinct(), 'cluster')
    else:
        raw = df

    h = HTMLParser.HTMLParser()
    removeTags = udf(lambda s: h.unescape(sub("</?[A-Za-z][^>]*>", "", s)), StringType())
    constructURL = udf(lambda url, corpus, id, pages, regions: formatURL(url, corpus, id, pages, regions),
                       StringType())

    raw\
        .withColumnRenamed("title", "heading")\
        .withColumn("url", constructURL(raw.url, raw.corpus, raw.id, raw.pages, raw.regions))\
        .drop("locs")\
        .drop("pages")\
        .drop("regions")\
        .withColumn("text", removeTags(raw.text))\
        .join(meta, 'series', 'left_outer')\
        .orderBy(desc("size"), "cluster", "date", "id", "begin")\
        .write.format(outputFormat).options(**outputOptions).save(outpath)

    sc.stop()

<commit_msg>Metadata should already be keyed by series.<commit_after>from __future__ import print_function

import sys
from re import sub
import HTMLParser

from pyspark import SparkContext
from pyspark.sql import SQLContext
from pyspark.sql.functions import col, udf, array_contains, explode, desc
from pyspark.sql.types import StringType

def guessFormat(path, default="json"):
    if path.endswith(".json"):
        return ("json", {})
    elif path.endswith(".parquet"):
        return ("parquet", {})
    elif path.endswith(".csv"):
        return ("com.databricks.spark.csv", {'header': 'true'})
    else:
        return (default, {})

## This should be made obsolete by including links in document records.
def formatURL(url, corpus, id, pages, regions):
    if corpus == 'ca':
        r = regions[0]
        return "%s/print/image_600x600_from_%d%%2C%d_to_%d%%2C%d/" \
            % (url, r.x/3, r.y/3, (r.x + r.w)/3, (r.y + r.h)/3)
    elif corpus == 'onb':
        return "%s&seite=%s" % (sub("&", "&", url), pages[0])
    elif corpus == 'trove':
        return "http://trove.nla.gov.au/ndp/del/article/%s" % sub("^trove/", "", id)
    else:
        return url

if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: pretty-cluster.py <metadata> <input> <output> [<query>]", file=sys.stderr)
        exit(-1)
    sc = SparkContext(appName="Prettyprint Clusters")
    sqlContext = SQLContext(sc)

    outpath = sys.argv[3]
    (outputFormat, outputOptions) = guessFormat(outpath, "json")

    ## Should do more field renaming in meta to avoid clashing with fields in raw.
    meta = sqlContext.read.json(sys.argv[1])\
           .withColumnRenamed('lang', 'series_lang')\
           .dropDuplicates(['series'])

    df = sqlContext.read.load(sys.argv[2])

    if len(sys.argv) >= 5:
        df.registerTempTable("clusters")
        print(sys.argv[4])
        raw = df.join(df.filter(sys.argv[4]).select("cluster").distinct(), 'cluster')
    else:
        raw = df

    h = HTMLParser.HTMLParser()
    removeTags = udf(lambda s: h.unescape(sub("</?[A-Za-z][^>]*>", "", s)), StringType())
    constructURL = udf(lambda url, corpus, id, pages, regions: formatURL(url, corpus, id, pages, regions),
                       StringType())

    raw\
        .withColumnRenamed("title", "heading")\
        .withColumn("url", constructURL(raw.url, raw.corpus, raw.id, raw.pages, raw.regions))\
        .drop("locs")\
        .drop("pages")\
        .drop("regions")\
        .withColumn("text", removeTags(raw.text))\
        .join(meta, 'series', 'left_outer')\
        .orderBy(desc("size"), "cluster", "date", "id", "begin")\
        .write.format(outputFormat).options(**outputOptions).save(outpath)

    sc.stop()

<|endoftext|>


<commit_before>import json

class D3MapBuilder:
  def __init__(self, db):
    self._db = db

  def build(self):
    output = dict()

    nodes = self._db.get_nodes()

    output['nodes'] = [{'group': x.group, 'name': x.name, 'id': x.id,
                        'macs': ', '.join(x.macs)
                       } for x in nodes if x.online]
    output['links'] = [{'source': x.pair[0], 'target': x.pair[1],
                        'distance': x.distance,
                        'strength': x.strength,
                        'quality': x.quality,
                        'id': "-".join(nodes[i].id for i in x.pair)
                       } for x in self._db.get_links()]

    return json.dumps(output)

<commit_msg>add geo field to nodes<commit_after>import json

class D3MapBuilder:
  def __init__(self, db):
    self._db = db

  def build(self):
    output = dict()

    nodes = self._db.get_nodes()

    output['nodes'] = [{'group': x.group, 'name': x.name, 'id': x.id,
                        'macs': ', '.join(x.macs),
                        'geo': x.gps.split(" ") if x.gps else None
                       } for x in nodes if x.online]
    output['links'] = [{'source': x.pair[0], 'target': x.pair[1],
                        'distance': x.distance,
                        'strength': x.strength,
                        'quality': x.quality,
                        'id': "-".join(nodes[i].id for i in x.pair)
                       } for x in self._db.get_links()]

    return json.dumps(output)

<|endoftext|>


<commit_before>"""
A number of functions that help the Amcat3 selection page to retrieve selections from the database
"""

from amcat.models import article

def getQuerySet(projects=None, articlesets=None, mediums=None, startDate=None, endDate=None, articleids=None, **kargs):
    queryset = article.Article.objects
    if articlesets:
        queryset = queryset.filter(articlesets__in=articlesets)
    elif projects:
        queryset = queryset.filter(project__in=projects)
    else:
        raise ValueError("Either projects or article sets needs to be specified")
    if mediums:
        queryset = queryset.filter(medium__in=mediums)
    if startDate:
        queryset = queryset.filter(date__gte=startDate)
    if endDate:
        queryset = queryset.filter(date__lte=endDate)
    if articleids:
        queryset = queryset.filter(id__in=articleids)
    return queryset
<commit_msg>Fixed issue #199<commit_after>"""
A number of functions that help the Amcat3 selection page to retrieve selections from the database
"""

from amcat.models import article

def getQuerySet(projects=None, articlesets=None, mediums=None, startDate=None, endDate=None, articleids=None, **kargs):
    queryset = article.Article.objects
    if articlesets:
        queryset = queryset.filter(articlesets__in=articlesets)
    elif projects:
        queryset = queryset.filter(articlesets__project__in=projects)
    else:
        raise ValueError("Either projects or article sets needs to be specified")

    if mediums:
        queryset = queryset.filter(medium__in=mediums)
    if startDate:
        queryset = queryset.filter(date__gte=startDate)
    if endDate:
        queryset = queryset.filter(date__lte=endDate)
    if articleids:
        queryset = queryset.filter(id__in=articleids)
    return queryset
<|endoftext|>


<commit_before>import functools
import logging
import multiprocessing
import signal
import socket
import sys

from oslo.config import cfg

from akanda.rug import event
from akanda.rug import health
from akanda.rug import notifications
from akanda.rug import scheduler
from akanda.rug import populate
from akanda.rug import worker
from akanda.rug.api import quantum as quantum_api

LOG = logging.getLogger(__name__)

def shuffle_notifications(notification_queue, sched):
    """Copy messages from the notification queue into the scheduler.
    """
    while True:
        try:
            target, message = notification_queue.get()
            sched.handle_message(target, message)
        # FIXME(rods): if a signal arrive during an IO operation an
        # IOError is raised. We catch the exceptions in meantime
        # waiting for a better solution.
        except IOError:
            pass
        except KeyboardInterrupt:
            sched.stop()
            break

def register_and_load_opts(argv):
    cfg.CONF.register_opts([
        cfg.StrOpt('host',
                   default=socket.getfqdn(),
                   help="The hostname Akanda is running on"),

        # FIXME(dhellmann): Use a separate group for these auth params
        cfg.StrOpt('admin_user'),
        cfg.StrOpt('admin_password', secret=True),
        cfg.StrOpt('admin_tenant_name'),
        cfg.StrOpt('auth_url'),
        cfg.StrOpt('auth_strategy', default='keystone'),
        cfg.StrOpt('auth_region'),

        cfg.StrOpt('management_network_id'),
        cfg.StrOpt('external_network_id'),
        cfg.StrOpt('management_subnet_id'),
        cfg.StrOpt('router_image_uuid'),

        cfg.StrOpt('management_prefix', default='fdca:3ba5:a17a:acda::/64'),
        cfg.IntOpt('akanda_mgt_service_port', default=5000),
        cfg.IntOpt('router_instance_flavor', default=1),

        # needed for plugging locally into management network
        cfg.StrOpt('interface_driver'),
        cfg.StrOpt('ovs_integration_bridge', default='br-int'),
        cfg.BoolOpt('ovs_use_veth', default=False),
        cfg.IntOpt('network_device_mtu'),
    ])

    AGENT_OPTIONS = [
        cfg.StrOpt('root_helper', default='sudo'),
    ]

    cfg.CONF.register_opts(AGENT_OPTIONS, 'AGENT')

    # FIXME: Convert these to regular options, not command line options.
    cfg.CONF.register_cli_opts([
        cfg.IntOpt('health-check-period',
                   default=60,
                   help='seconds between health checks'),
        cfg.IntOpt('num-worker-processes',
                   short='w',
                   default=16,
                   help='the number of worker processes to run'),
        cfg.IntOpt('num-worker-threads',
                   short='t',
                   default=4,
                   help='the number of worker threads to run per process'),

        # FIXME(dhellmann): set up a group for these messaging params
        cfg.StrOpt('amqp-url',
                   default='amqp://guest:secrete@localhost:5672/',
                   help='connection for AMQP server'),
        cfg.StrOpt('incoming-notifications-exchange',
                   default='neutron',
                   help='name of the exchange where we receive notifications'),
        cfg.StrOpt('outgoing-notifications-exchange',
                   default='neutron',
                   help='name of the exchange where we send notifications'),
        cfg.StrOpt('rpc-exchange',
                   default='l3_agent_fanout',
                   help='name of the exchange where we receive RPC calls'),
    ])
    cfg.CONF(argv, project='akanda')

def main(argv=sys.argv[1:]):
    register_and_load_opts(argv)
    logging.basicConfig(
        level=logging.DEBUG,
        format=':'.join('%(' + n + ')s'
                        for n in ['processName',
                                  'threadName',
                                  'name',
                                  'levelname',
                                  'message']),
    )

    # Purge the mgt tap interface on startup
    quantum = quantum_api.Quantum(cfg.CONF)
    #TODO(mark): develop better way restore after machine reboot
    #quantum.purge_management_interface()

    # bring the mgt tap interface up
    quantum.ensure_local_service_port()

    # Set up the queue to move messages between the eventlet-based
    # listening process and the scheduler.
    notification_queue = multiprocessing.Queue()

    # Listen for notifications.
    notification_proc = multiprocessing.Process(
        target=notifications.listen,
        kwargs={
            'host_id': cfg.CONF.host,
            'amqp_url': cfg.CONF.amqp_url,
            'notifications_exchange_name':
            cfg.CONF.incoming_notifications_exchange,
            'rpc_exchange_name': cfg.CONF.rpc_exchange,
            'notification_queue': notification_queue,
        },
        name='notification-listener',
    )
    notification_proc.start()

    # Set up the notifications publisher
    publisher = notifications.Publisher(
        cfg.CONF.amqp_url,
        exchange_name=cfg.CONF.outgoing_notifications_exchange,
        topic='notifications.info',
    )

    # Set up a factory to make Workers that know how many threads to
    # run.
    worker_factory = functools.partial(
        worker.Worker,
        num_threads=cfg.CONF.num_worker_threads,
        notifier=publisher,
    )

    # Set up the scheduler that knows how to manage the routers and
    # dispatch messages.
    sched = scheduler.Scheduler(
        num_workers=cfg.CONF.num_worker_processes,
        worker_factory=worker_factory,
    )

    def debug_workers(signum, stack):
        sched.handle_message(
            'debug',
            event.Event('debug', '', event.POLL, {'verbose': 1})
        )

    signal.signal(signal.SIGUSR1, debug_workers)

    # Prepopulate the workers with existing routers on startup
    populate.pre_populate_workers(sched)

    # Set up the periodic health check
    health.start_inspector(cfg.CONF.health_check_period, sched)

    # Block the main process, copying messages from the notification
    # listener to the scheduler
    shuffle_notifications(notification_queue, sched)

    # Terminate the listening process
    notification_proc.terminate()

    LOG.info('exiting')
<commit_msg>Add timestamps to logs<commit_after>import functools
import logging
import multiprocessing
import signal
import socket
import sys

from oslo.config import cfg

from akanda.rug import event
from akanda.rug import health
from akanda.rug import notifications
from akanda.rug import scheduler
from akanda.rug import populate
from akanda.rug import worker
from akanda.rug.api import quantum as quantum_api

LOG = logging.getLogger(__name__)

def shuffle_notifications(notification_queue, sched):
    """Copy messages from the notification queue into the scheduler.
    """
    while True:
        try:
            target, message = notification_queue.get()
            sched.handle_message(target, message)
        # FIXME(rods): if a signal arrive during an IO operation an
        # IOError is raised. We catch the exceptions in meantime
        # waiting for a better solution.
        except IOError:
            pass
        except KeyboardInterrupt:
            sched.stop()
            break

def register_and_load_opts(argv):
    cfg.CONF.register_opts([
        cfg.StrOpt('host',
                   default=socket.getfqdn(),
                   help="The hostname Akanda is running on"),

        # FIXME(dhellmann): Use a separate group for these auth params
        cfg.StrOpt('admin_user'),
        cfg.StrOpt('admin_password', secret=True),
        cfg.StrOpt('admin_tenant_name'),
        cfg.StrOpt('auth_url'),
        cfg.StrOpt('auth_strategy', default='keystone'),
        cfg.StrOpt('auth_region'),

        cfg.StrOpt('management_network_id'),
        cfg.StrOpt('external_network_id'),
        cfg.StrOpt('management_subnet_id'),
        cfg.StrOpt('router_image_uuid'),

        cfg.StrOpt('management_prefix', default='fdca:3ba5:a17a:acda::/64'),
        cfg.IntOpt('akanda_mgt_service_port', default=5000),
        cfg.IntOpt('router_instance_flavor', default=1),

        # needed for plugging locally into management network
        cfg.StrOpt('interface_driver'),
        cfg.StrOpt('ovs_integration_bridge', default='br-int'),
        cfg.BoolOpt('ovs_use_veth', default=False),
        cfg.IntOpt('network_device_mtu'),
    ])

    AGENT_OPTIONS = [
        cfg.StrOpt('root_helper', default='sudo'),
    ]

    cfg.CONF.register_opts(AGENT_OPTIONS, 'AGENT')

    # FIXME: Convert these to regular options, not command line options.
    cfg.CONF.register_cli_opts([
        cfg.IntOpt('health-check-period',
                   default=60,
                   help='seconds between health checks'),
        cfg.IntOpt('num-worker-processes',
                   short='w',
                   default=16,
                   help='the number of worker processes to run'),
        cfg.IntOpt('num-worker-threads',
                   short='t',
                   default=4,
                   help='the number of worker threads to run per process'),

        # FIXME(dhellmann): set up a group for these messaging params
        cfg.StrOpt('amqp-url',
                   default='amqp://guest:secrete@localhost:5672/',
                   help='connection for AMQP server'),
        cfg.StrOpt('incoming-notifications-exchange',
                   default='neutron',
                   help='name of the exchange where we receive notifications'),
        cfg.StrOpt('outgoing-notifications-exchange',
                   default='neutron',
                   help='name of the exchange where we send notifications'),
        cfg.StrOpt('rpc-exchange',
                   default='l3_agent_fanout',
                   help='name of the exchange where we receive RPC calls'),
    ])
    cfg.CONF(argv, project='akanda')

def main(argv=sys.argv[1:]):
    register_and_load_opts(argv)
    logging.basicConfig(
        level=logging.DEBUG,
        format=':'.join('%(' + n + ')s'
                        for n in ['asctime',
                                  'processName',
                                  'threadName',
                                  'name',
                                  'levelname',
                                  'message']),
    )

    # Purge the mgt tap interface on startup
    quantum = quantum_api.Quantum(cfg.CONF)
    #TODO(mark): develop better way restore after machine reboot
    #quantum.purge_management_interface()

    # bring the mgt tap interface up
    quantum.ensure_local_service_port()

    # Set up the queue to move messages between the eventlet-based
    # listening process and the scheduler.
    notification_queue = multiprocessing.Queue()

    # Listen for notifications.
    notification_proc = multiprocessing.Process(
        target=notifications.listen,
        kwargs={
            'host_id': cfg.CONF.host,
            'amqp_url': cfg.CONF.amqp_url,
            'notifications_exchange_name':
            cfg.CONF.incoming_notifications_exchange,
            'rpc_exchange_name': cfg.CONF.rpc_exchange,
            'notification_queue': notification_queue,
        },
        name='notification-listener',
    )
    notification_proc.start()

    # Set up the notifications publisher
    publisher = notifications.Publisher(
        cfg.CONF.amqp_url,
        exchange_name=cfg.CONF.outgoing_notifications_exchange,
        topic='notifications.info',
    )

    # Set up a factory to make Workers that know how many threads to
    # run.
    worker_factory = functools.partial(
        worker.Worker,
        num_threads=cfg.CONF.num_worker_threads,
        notifier=publisher,
    )

    # Set up the scheduler that knows how to manage the routers and
    # dispatch messages.
    sched = scheduler.Scheduler(
        num_workers=cfg.CONF.num_worker_processes,
        worker_factory=worker_factory,
    )

    def debug_workers(signum, stack):
        sched.handle_message(
            'debug',
            event.Event('debug', '', event.POLL, {'verbose': 1})
        )

    signal.signal(signal.SIGUSR1, debug_workers)

    # Prepopulate the workers with existing routers on startup
    populate.pre_populate_workers(sched)

    # Set up the periodic health check
    health.start_inspector(cfg.CONF.health_check_period, sched)

    # Block the main process, copying messages from the notification
    # listener to the scheduler
    shuffle_notifications(notification_queue, sched)

    # Terminate the listening process
    notification_proc.terminate()

    LOG.info('exiting')
<|endoftext|>


<commit_before># Convert twoPowerCurves(Fig 5.15 (b)) to python/JAX
# Author: Garvit9000c

import jax.numpy as jnp
import matplotlib.pyplot as plt

pi=jnp.pi

#Domain
x=jnp.arange(-1*pi,pi,0.01)

#functions
y1=-1*jnp.cos(x)+0.03
y2=-1*jnp.cos(2*x/3)

plt.ylim(-2,2.5)

#Ploting Curve
plt.plot(x,y1,'b') #Curve_B
plt.plot(x,y2,'r') #Curve_A

#Small line indicating position of alpha
plt.hlines(y=-1.025, xmin=-0.3, xmax=0.3, linewidth=1, color='k')

#Axis Arrows
plt.arrow(0,-2.5,0,4.5,color='black',width = 0.03)       #y-axis
plt.arrow(-1*pi,-1.5,2*pi,0,color='black',width = 0.03)  #x-axis

#labels
plt.text(-0.3, 2.2, '1 - β', fontsize=15) # 1 - β
plt.text(-0.3, 0.5, '1', fontsize=15)     # 1
plt.text(-0.3, -1.3, 'α', fontsize=15)    # α
plt.text(-0.34, -1.8, '$θ_0$', fontsize=12)  # θo
plt.text(3.1, -1.8, 'θ', fontsize=15)     # θ
plt.text(2.5, -0.2, 'A', fontsize=15)     # A
plt.text(2.5, 1.1, 'B', fontsize=15)      # B

plt.savefig('../figures/twoPowerCurves.pdf', dpi=300)
plt.show()
<commit_msg>Update twoPowerCurves.py<commit_after># Convert twoPowerCurves(Fig 5.15 (b)) to python/JAX
# Author: Garvit9000c

import numpy as np
import matplotlib.pyplot as plt

pi=np.pi

#Domain
x=np.arange(-1*pi,pi,0.01)

#functions
y1=-1*np.cos(x)+0.03
y2=-1*np.cos(2*x/3)

plt.ylim(-2,2.5)

#Ploting Curve
plt.plot(x,y1,'b') #Curve_B
plt.plot(x,y2,'r') #Curve_A

#Small line indicating position of alpha
plt.hlines(y=-1.025, xmin=-0.3, xmax=0.3, linewidth=1, color='k')

#Axis Arrows
plt.arrow(0,-2.5,0,4.5,color='black',width = 0.03)       #y-axis
plt.arrow(-1*pi,-1.5,2*pi,0,color='black',width = 0.03)  #x-axis

#labels
plt.text(-0.3, 2.2, r'$1 - \beta$', fontsize=15) # 1 - β
plt.text(-0.3, 0.5, '1', fontsize=15)     # 1
plt.text(-0.3, -1.3, r'$\alpha$', fontsize=15)    # α
plt.text(-0.34, -1.8, '$θ_0$', fontsize=12)  # θo
plt.text(3.1, -1.8, 'θ', fontsize=15)     # θ
plt.text(2.5, -0.2, 'A', fontsize=15)     # A
plt.text(2.5, 1.1, 'B', fontsize=15)      # B

plt.savefig('../figures/twoPowerCurves.pdf', dpi=300)
plt.show()
<|endoftext|>


<commit_before>import os
import requests

# NOTE: this URL should be updated once the reach_export branch is merged
base_url = ('https://raw.githubusercontent.com/bgyori/famplex/reach_export/'
            'export/')
famplex_groundings = 'famplex_groundings.tsv'
famplex_overrides = 'famplex_grounding_overrides.tsv'

def get_real_lines(fname):
    with open(fname, 'r') as fh:
        lines = [l.strip() for l in fh.readlines() if l.strip() and
                 not l.startswith('#')]
    lines = [l.split('\t') for l in lines]
    return lines

def get_other_strings(fname):
    kb_files = [row[0] + '.tsv' for row in get_real_lines(fname)]
    strings = []
    for kb_fname in kb_files:
        # We skip famplex here to avoid redundancy
        if kb_fname == 'famplex.tsv':
            continue
        strings += [row[0] for row in
                    get_real_lines(os.path.join(kb_dir, kb_fname))]
    return set(strings)

def get_overrides(overrides_rows, other_strings, famplex_only=True):
    overrides = []
    for txt, db_id, db_ns, type in overrides_rows:
        # In famplex_only mode, skip and non-fplx groundings
        if famplex_only and db_ns != 'fplx':
            continue
        # If this is not an actual override, skip it
        if txt not in other_strings:
            continue
        overrides.append((txt, db_id, db_ns, type))
    return overrides

def extend_overrides(override_fname, overrides):
    with open(override_fname, 'r') as fh:
        lines = [l.strip() for l in fh.readlines()]
    try:
        idx = lines.index('# FamPlex overrides')
    except ValueError:
        idx = None
    lines_out = lines[:idx+1] if idx else lines
    # Add comment block if it is not already there
    lines_out += [] if idx else ['#', '# FamPlex overrides', '#']
    lines_out += ['\t'.join(l) for l in overrides]
    with open(override_fname, 'w') as fh:
        for line in lines_out:
            fh.write('%s\n' % line)

if __name__ == '__main__':
    # Basic positioning of folders
    here = os.path.dirname(os.path.abspath(__file__))
    kb_dir = os.path.join(here, os.pardir, 'src', 'main', 'resources', 'org',
                          'clulab', 'reach', 'kb')
    groundings_fname = os.path.join(kb_dir, 'famplex.tsv')
    override_fname = os.path.join(kb_dir, 'NER-Grounding-Override.tsv')

    # Download and write to groundings file
    with open(groundings_fname, 'w') as fh:
        fh.write(requests.get(base_url + famplex_groundings).text)

    # Now get all the "other" strings so we can figure out what to add
    # to the overrides file
    other_strings = get_other_strings(os.path.join(here, os.pardir,
                                      'ner_kb.config'))

    overrides_rows = \
        [row.split('\t') for row in
         requests.get(base_url + famplex_overrides).text.split('\n')]
    overrides = get_overrides(overrides_rows, other_strings, famplex_only=True)

    extend_overrides(override_fname, overrides)
<commit_msg>Simplify update script for FamPlex groundings/overrides<commit_after>import os
import requests

base_url = ('https://raw.githubusercontent.com/sorgerlab/famplex/master/'
            'export/')
famplex_groundings = 'famplex_groundings.tsv'

def get_real_lines(fname):
    with open(fname, 'r') as fh:
        lines = [l.strip() for l in fh.readlines() if l.strip() and
                 not l.startswith('#')]
    lines = [l.split('\t') for l in lines]
    return lines

def get_other_strings(fname):
    kb_files = [row[0] + '.tsv' for row in get_real_lines(fname)]
    strings = []
    for kb_fname in kb_files:
        # We skip famplex here to avoid redundancy
        if kb_fname == 'famplex.tsv':
            continue
        strings += [row[0] for row in
                    get_real_lines(os.path.join(kb_dir, kb_fname))]
    return set(strings)

def get_overrides(overrides_rows, other_strings, famplex_only=True):
    overrides = []
    for txt, db_id, db_ns, type in overrides_rows:
        # In famplex_only mode, skip and non-fplx groundings
        if famplex_only and db_ns != 'fplx':
            continue
        # If this is not an actual override, skip it
        if txt not in other_strings:
            continue
        overrides.append((txt, db_id, db_ns, type))
    return overrides

def extend_overrides(override_fname, overrides):
    with open(override_fname, 'r') as fh:
        lines = [l.strip() for l in fh.readlines()]
    try:
        idx = lines.index('# FamPlex overrides')
    except ValueError:
        idx = None
    lines_out = lines[:idx+2] if idx else lines
    # Add comment block if it is not already there
    lines_out += [] if idx else ['#', '# FamPlex overrides', '#']
    lines_out += ['\t'.join(l) for l in overrides]
    with open(override_fname, 'w') as fh:
        for line in lines_out:
            fh.write('%s\n' % line)

if __name__ == '__main__':
    # Basic positioning of folders
    here = os.path.dirname(os.path.abspath(__file__))
    kb_dir = os.path.join(here, os.pardir, 'src', 'main', 'resources', 'org',
                          'clulab', 'reach', 'kb')
    groundings_fname = os.path.join(kb_dir, 'famplex.tsv')
    override_fname = os.path.join(kb_dir, 'NER-Grounding-Override.tsv')

    groundings_rows = [line.strip().split('\t') for line in
                       requests.get(base_url +
                                    famplex_groundings).text.split('\n')]

    # Download and write to groundings file
    with open(groundings_fname, 'w') as fh:
        for line in groundings_rows:
            if line[2] == 'fplx':
                fh.write('%s\t%s\n' % (line[0], line[1]))

    # Now get all the "other" strings so we can figure out what to add
    # to the overrides file
    other_strings = get_other_strings(os.path.join(here, os.pardir,
                                      'ner_kb.config'))

    overrides = get_overrides(groundings_rows, other_strings,
                              famplex_only=True)

    extend_overrides(override_fname, overrides)
<|endoftext|>


<commit_before>import sys
import os
import pwd

try:
  import builtins as __builtin__
except ImportError:
  import __builtin__

from hashlib import sha256
from os.path import abspath, basename, dirname
from time import time
from twisted.internet import defer
from twisted.python.util import initgroups
from twisted.scripts.twistd import runApp

# BytesIO is needed on py3 as StringIO does not operate on byte input anymore
# We could use BytesIO on py2 as well but it is slower than StringIO
if sys.version_info >= (3, 0):
  from io import BytesIO as StringIO
else:
  try:
    from cStringIO import StringIO
  except ImportError:
    from StringIO import StringIO

try:
  import cPickle as pickle
  USING_CPICKLE = True
except ImportError:
  import pickle
  USING_CPICKLE = False

def dropprivs(user):
  uid, gid = pwd.getpwnam(user)[2:4]
  initgroups(uid, gid)
  os.setregid(gid, gid)
  os.setreuid(uid, uid)
  return (uid, gid)

def run_twistd_plugin(filename):
    from carbon.conf import get_parser
    from twisted.scripts.twistd import ServerOptions

    bin_dir = dirname(abspath(filename))
    root_dir = dirname(bin_dir)
    os.environ.setdefault('GRAPHITE_ROOT', root_dir)

    program = basename(filename).split('.')[0]

    # First, parse command line options as the legacy carbon scripts used to
    # do.
    parser = get_parser(program)
    (options, args) = parser.parse_args()

    if not args:
      parser.print_usage()
      return

    # This isn't as evil as you might think
    __builtin__.instance = options.instance
    __builtin__.program = program

    # Then forward applicable options to either twistd or to the plugin itself.
    twistd_options = ["--no_save"]

    # If no reactor was selected yet, try to use the epoll reactor if
    # available.
    try:
        from twisted.internet import epollreactor  # noqa: F401
        twistd_options.append("--reactor=epoll")
    except ImportError:
        pass

    if options.debug or options.nodaemon:
        twistd_options.extend(["--nodaemon"])
    if options.profile:
        twistd_options.extend(["--profile", options.profile])
    if options.profiler:
        twistd_options.extend(["--profiler", options.profiler])
    if options.pidfile:
        twistd_options.extend(["--pidfile", options.pidfile])
    if options.umask:
        twistd_options.extend(["--umask", options.umask])
    if options.syslog:
        twistd_options.append("--syslog")

    # Now for the plugin-specific options.
    twistd_options.append(program)

    if options.debug:
        twistd_options.append("--debug")

    for option_name, option_value in vars(options).items():
        if (option_value is not None and option_name not in (
                "debug", "profile", "profiler", "pidfile", "umask", "nodaemon", "syslog")):
            twistd_options.extend(["--%s" % option_name.replace("_", "-"), option_value])

    # Finally, append extra args so that twistd has a chance to process them.
    twistd_options.extend(args)

    config = ServerOptions()
    config.parseOptions(twistd_options)

    runApp(config)

def parseDestination(dest_string):
    s = dest_string.strip()
    bidx = s.rfind(']:')    # find closing bracket and following colon.
    cidx = s.find(':')
    if s.startswith('[') and bidx is not None:
        server = s[1:bidx]
        port = s[bidx + 2:]
    elif cidx is not None:
        server = s[:cidx]
        port = s[cidx + 1:]
    else:
        raise ValueError("Invalid destination string \"%s\"" % dest_string)

    if ':' in port:
        port, _, instance = port.partition(':')
    else:
        instance = None

    return server, int(port), instance

def parseDestinations(destination_strings):
    return [parseDestination(dest_string) for dest_string in destination_strings]

# Yes this is duplicated in whisper. Yes, duplication is bad.
# But the code is needed in both places and we do not want to create
# a dependency on whisper especiaily as carbon moves toward being a more
# generic storage service that can use various backends.
UnitMultipliers = {
  's': 1,
  'm': 60,
  'h': 60 * 60,
  'd': 60 * 60 * 24,
  'w': 60 * 60 * 24 * 7,
  'y': 60 * 60 * 24 * 365,
}

def getUnitString(s):
  if s not in UnitMultipliers:
    raise ValueError("Invalid unit '%s'" % s)
  return s

def parseRetentionDef(retentionDef):
  import re
  (precision, points) = retentionDef.strip().split(':')

  if precision.isdigit():
    precision = int(precision) * UnitMultipliers[getUnitString('s')]
  else:
    precision_re = re.compile(r'^(\d+)([a-z]+)$')
    match = precision_re.match(precision)
    if match:
      precision = int(match.group(1)) * UnitMultipliers[getUnitString(match.group(2))]
    else:
      raise ValueError("Invalid precision specification '%s'" % precision)

  if points.isdigit():
    points = int(points)
  else:
    points_re = re.compile(r'^(\d+)([a-z]+)$')
    match = points_re.match(points)
    if match:
      points = int(match.group(1)) * UnitMultipliers[getUnitString(match.group(2))] / precision
    else:
      raise ValueError("Invalid retention specification '%s'" % points)

  return (precision, points)

# This whole song & dance is due to pickle being insecure
# yet performance critical for carbon. We leave the insecure
# mode (which is faster) as an option (USE_INSECURE_UNPICKLER).
# The SafeUnpickler classes were largely derived from
# http://nadiana.com/python-pickle-insecure
if USING_CPICKLE:
  class SafeUnpickler(object):
    PICKLE_SAFE = {
      'copy_reg': set(['_reconstructor']),
      '__builtin__': set(['object']),
    }

    @classmethod
    def find_class(cls, module, name):
      if module not in cls.PICKLE_SAFE:
        raise pickle.UnpicklingError('Attempting to unpickle unsafe module %s' % module)
      __import__(module)
      mod = sys.modules[module]
      if name not in cls.PICKLE_SAFE[module]:
        raise pickle.UnpicklingError('Attempting to unpickle unsafe class %s' % name)
      return getattr(mod, name)

    @classmethod
    def loads(cls, pickle_string):
      pickle_obj = pickle.Unpickler(StringIO(pickle_string))
      pickle_obj.find_global = cls.find_class
      return pickle_obj.load()

else:
  class SafeUnpickler(pickle.Unpickler):
    PICKLE_SAFE = {
      'copy_reg': set(['_reconstructor']),
      '__builtin__': set(['object']),
    }

    def find_class(self, module, name):
      if module not in self.PICKLE_SAFE:
        raise pickle.UnpicklingError('Attempting to unpickle unsafe module %s' % module)
      __import__(module)
      mod = sys.modules[module]
      if name not in self.PICKLE_SAFE[module]:
        raise pickle.UnpicklingError('Attempting to unpickle unsafe class %s' % name)
      return getattr(mod, name)

    @classmethod
    def loads(cls, pickle_string):
      if sys.version_info >= (3, 0):
        return cls(StringIO(pickle_string), encoding='utf-8').load()
      else:
        return cls(StringIO(pickle_string)).load()

def get_unpickler(insecure=False):
  if insecure:
    return pickle
  else:
    return SafeUnpickler

class TokenBucket(object):
  '''This is a basic tokenbucket rate limiter implementation for use in
  enforcing various configurable rate limits'''
  def __init__(self, capacity, fill_rate, reactor):
    '''Capacity is the total number of tokens the bucket can hold, fill rate is
    the rate in tokens (or fractional tokens) to be added to the bucket per
    second.'''
    self.capacity = float(capacity)
    self._tokens = float(capacity)
    self.fill_rate = float(fill_rate)
    self.timestamp = time()
    self.reactor = reactor

  def drain(self, cost, blocking=False):
    '''Given a number of tokens (or fractions) drain will return True and
    drain the number of tokens from the bucket if the capacity allows,
    otherwise we return false and leave the contents of the bucket.'''
    if not blocking:
      if cost <= self.tokens:
        self._tokens -= cost
        return True

      return False

    d = defer.Deferred()

    if cost <= self.tokens:
      self._tokens -= cost
      d.callback(True)
      return d

    tokens_needed = cost - self._tokens
    seconds_per_token = 1 / self.fill_rate
    seconds_left = seconds_per_token * tokens_needed
    time_to_sleep = self.timestamp + seconds_left - time()
    if time_to_sleep <= 0:
      self._tokens -= cost
      d.callback(True)
      return d

    def drain_and_return(val):
      self._tokens -= cost
      return d.callback(val)

    self.reactor.callLater(time_to_sleep, drain_and_return, True)
    return d

  def setCapacityAndFillRate(self, new_capacity, new_fill_rate):
    delta = float(new_capacity) - self.capacity
    self.capacity = float(new_capacity)
    self.fill_rate = float(new_fill_rate)
    self._tokens = delta + self._tokens

  @property
  def tokens(self):
    '''The tokens property will return the current number of tokens in the
    bucket.'''
    if self._tokens < self.capacity:
      now = time()
      delta = self.fill_rate * (now - self.timestamp)
      self._tokens = min(self.capacity, self._tokens + delta)
      self.timestamp = now
    return self._tokens

class PluginRegistrar(type):
  """Clever subclass detection hack that makes plugin loading trivial.
  To use this, define an abstract base class for plugin implementations
  that defines the plugin API. Give that base class a __metaclass__ of
  PluginRegistrar, and define a 'plugins = {}' class member. Subclasses
  defining a 'plugin_name' member will then appear in the plugins dict.
  """
  def __init__(classObj, name, bases, members):
    super(PluginRegistrar, classObj).__init__(name, bases, members)
    if hasattr(classObj, 'plugin_name'):
      classObj.plugins[classObj.plugin_name] = classObj

class TaggedSeries(object):
  @classmethod
  def parse(cls, path):
    tags = {}

    # if path is in openmetrics format: metric{tag="value",...}
    # https://github.com/RichiH/OpenMetrics
    if path[-2:] == '"}' and '{' in path:
      (metric, rawtags) = path[0:-1].split('{', 2)
      if not metric:
        raise Exception('Cannot parse path %s, no metric found' % path)

      while len(rawtags) > 0:
        m = re.match(r'([^=]+)="((?:[\\]["\\]|[^"\\])+)"(:?,|$)', rawtags)
        if not m:
          raise Exception('Cannot parse path %s, invalid segment %s' % (path, rawtags))

        tags[m.group(1)] = m.group(2).replace(r'\"', '"').replace(r'\\', '\\')
        rawtags = rawtags[len(m.group(0)):]

      tags['name'] = metric
      return cls(metric, tags)

    # path is a carbon path with optional tags: metric;tag=value;...
    segments = path.split(';')

    metric = segments[0]
    if not metric:
      raise Exception('Cannot parse path %s, no metric found' % path)

    for segment in segments[1:]:
      tag = segment.split('=', 1)
      if len(tag) != 2 or not tag[0]:
        raise Exception('Cannot parse path %s, invalid segment %s' % (path, segment))

      tags[tag[0]] = tag[1]

    tags['name'] = metric
    return cls(metric, tags)

  @staticmethod
  def format(tags):
    return tags.get('name', '') + ''.join(sorted([
      ';%s=%s' % (tag, value)
      for tag, value in tags.items()
      if tag != 'name'
    ]))

  @staticmethod
  def encode(metric, sep='.'):
    # if metric is tagged, encode it for storage in whisper etc
    if ';' in metric:
      metric_hash = sha256(metric.encode('utf8')).hexdigest()
      return sep.join(['_tagged', metric_hash[0:3], metric_hash[3:6], metric.replace('.', '-')])

    # metric isn't tagged, just replace dots with the separator and trim any leading separator
    return metric.replace('.', sep).lstrip(sep)

  def __init__(self, metric, tags, series_id=None):
    self.metric = metric
    self.tags = tags
    self.series_id = id

  @property
  def path(self):
    return self.__class__.format(self.tags)
<commit_msg>add missing import<commit_after>import sys
import os
import pwd
import re

try:
  import builtins as __builtin__
except ImportError:
  import __builtin__

from hashlib import sha256
from os.path import abspath, basename, dirname
from time import time
from twisted.internet import defer
from twisted.python.util import initgroups
from twisted.scripts.twistd import runApp

# BytesIO is needed on py3 as StringIO does not operate on byte input anymore
# We could use BytesIO on py2 as well but it is slower than StringIO
if sys.version_info >= (3, 0):
  from io import BytesIO as StringIO
else:
  try:
    from cStringIO import StringIO
  except ImportError:
    from StringIO import StringIO

try:
  import cPickle as pickle
  USING_CPICKLE = True
except ImportError:
  import pickle
  USING_CPICKLE = False

def dropprivs(user):
  uid, gid = pwd.getpwnam(user)[2:4]
  initgroups(uid, gid)
  os.setregid(gid, gid)
  os.setreuid(uid, uid)
  return (uid, gid)

def run_twistd_plugin(filename):
    from carbon.conf import get_parser
    from twisted.scripts.twistd import ServerOptions

    bin_dir = dirname(abspath(filename))
    root_dir = dirname(bin_dir)
    os.environ.setdefault('GRAPHITE_ROOT', root_dir)

    program = basename(filename).split('.')[0]

    # First, parse command line options as the legacy carbon scripts used to
    # do.
    parser = get_parser(program)
    (options, args) = parser.parse_args()

    if not args:
      parser.print_usage()
      return

    # This isn't as evil as you might think
    __builtin__.instance = options.instance
    __builtin__.program = program

    # Then forward applicable options to either twistd or to the plugin itself.
    twistd_options = ["--no_save"]

    # If no reactor was selected yet, try to use the epoll reactor if
    # available.
    try:
        from twisted.internet import epollreactor  # noqa: F401
        twistd_options.append("--reactor=epoll")
    except ImportError:
        pass

    if options.debug or options.nodaemon:
        twistd_options.extend(["--nodaemon"])
    if options.profile:
        twistd_options.extend(["--profile", options.profile])
    if options.profiler:
        twistd_options.extend(["--profiler", options.profiler])
    if options.pidfile:
        twistd_options.extend(["--pidfile", options.pidfile])
    if options.umask:
        twistd_options.extend(["--umask", options.umask])
    if options.syslog:
        twistd_options.append("--syslog")

    # Now for the plugin-specific options.
    twistd_options.append(program)

    if options.debug:
        twistd_options.append("--debug")

    for option_name, option_value in vars(options).items():
        if (option_value is not None and option_name not in (
                "debug", "profile", "profiler", "pidfile", "umask", "nodaemon", "syslog")):
            twistd_options.extend(["--%s" % option_name.replace("_", "-"), option_value])

    # Finally, append extra args so that twistd has a chance to process them.
    twistd_options.extend(args)

    config = ServerOptions()
    config.parseOptions(twistd_options)

    runApp(config)

def parseDestination(dest_string):
    s = dest_string.strip()
    bidx = s.rfind(']:')    # find closing bracket and following colon.
    cidx = s.find(':')
    if s.startswith('[') and bidx is not None:
        server = s[1:bidx]
        port = s[bidx + 2:]
    elif cidx is not None:
        server = s[:cidx]
        port = s[cidx + 1:]
    else:
        raise ValueError("Invalid destination string \"%s\"" % dest_string)

    if ':' in port:
        port, _, instance = port.partition(':')
    else:
        instance = None

    return server, int(port), instance

def parseDestinations(destination_strings):
    return [parseDestination(dest_string) for dest_string in destination_strings]

# Yes this is duplicated in whisper. Yes, duplication is bad.
# But the code is needed in both places and we do not want to create
# a dependency on whisper especiaily as carbon moves toward being a more
# generic storage service that can use various backends.
UnitMultipliers = {
  's': 1,
  'm': 60,
  'h': 60 * 60,
  'd': 60 * 60 * 24,
  'w': 60 * 60 * 24 * 7,
  'y': 60 * 60 * 24 * 365,
}

def getUnitString(s):
  if s not in UnitMultipliers:
    raise ValueError("Invalid unit '%s'" % s)
  return s

def parseRetentionDef(retentionDef):
  import re
  (precision, points) = retentionDef.strip().split(':')

  if precision.isdigit():
    precision = int(precision) * UnitMultipliers[getUnitString('s')]
  else:
    precision_re = re.compile(r'^(\d+)([a-z]+)$')
    match = precision_re.match(precision)
    if match:
      precision = int(match.group(1)) * UnitMultipliers[getUnitString(match.group(2))]
    else:
      raise ValueError("Invalid precision specification '%s'" % precision)

  if points.isdigit():
    points = int(points)
  else:
    points_re = re.compile(r'^(\d+)([a-z]+)$')
    match = points_re.match(points)
    if match:
      points = int(match.group(1)) * UnitMultipliers[getUnitString(match.group(2))] / precision
    else:
      raise ValueError("Invalid retention specification '%s'" % points)

  return (precision, points)

# This whole song & dance is due to pickle being insecure
# yet performance critical for carbon. We leave the insecure
# mode (which is faster) as an option (USE_INSECURE_UNPICKLER).
# The SafeUnpickler classes were largely derived from
# http://nadiana.com/python-pickle-insecure
if USING_CPICKLE:
  class SafeUnpickler(object):
    PICKLE_SAFE = {
      'copy_reg': set(['_reconstructor']),
      '__builtin__': set(['object']),
    }

    @classmethod
    def find_class(cls, module, name):
      if module not in cls.PICKLE_SAFE:
        raise pickle.UnpicklingError('Attempting to unpickle unsafe module %s' % module)
      __import__(module)
      mod = sys.modules[module]
      if name not in cls.PICKLE_SAFE[module]:
        raise pickle.UnpicklingError('Attempting to unpickle unsafe class %s' % name)
      return getattr(mod, name)

    @classmethod
    def loads(cls, pickle_string):
      pickle_obj = pickle.Unpickler(StringIO(pickle_string))
      pickle_obj.find_global = cls.find_class
      return pickle_obj.load()

else:
  class SafeUnpickler(pickle.Unpickler):
    PICKLE_SAFE = {
      'copy_reg': set(['_reconstructor']),
      '__builtin__': set(['object']),
    }

    def find_class(self, module, name):
      if module not in self.PICKLE_SAFE:
        raise pickle.UnpicklingError('Attempting to unpickle unsafe module %s' % module)
      __import__(module)
      mod = sys.modules[module]
      if name not in self.PICKLE_SAFE[module]:
        raise pickle.UnpicklingError('Attempting to unpickle unsafe class %s' % name)
      return getattr(mod, name)

    @classmethod
    def loads(cls, pickle_string):
      if sys.version_info >= (3, 0):
        return cls(StringIO(pickle_string), encoding='utf-8').load()
      else:
        return cls(StringIO(pickle_string)).load()

def get_unpickler(insecure=False):
  if insecure:
    return pickle
  else:
    return SafeUnpickler

class TokenBucket(object):
  '''This is a basic tokenbucket rate limiter implementation for use in
  enforcing various configurable rate limits'''
  def __init__(self, capacity, fill_rate, reactor):
    '''Capacity is the total number of tokens the bucket can hold, fill rate is
    the rate in tokens (or fractional tokens) to be added to the bucket per
    second.'''
    self.capacity = float(capacity)
    self._tokens = float(capacity)
    self.fill_rate = float(fill_rate)
    self.timestamp = time()
    self.reactor = reactor

  def drain(self, cost, blocking=False):
    '''Given a number of tokens (or fractions) drain will return True and
    drain the number of tokens from the bucket if the capacity allows,
    otherwise we return false and leave the contents of the bucket.'''
    if not blocking:
      if cost <= self.tokens:
        self._tokens -= cost
        return True

      return False

    d = defer.Deferred()

    if cost <= self.tokens:
      self._tokens -= cost
      d.callback(True)
      return d

    tokens_needed = cost - self._tokens
    seconds_per_token = 1 / self.fill_rate
    seconds_left = seconds_per_token * tokens_needed
    time_to_sleep = self.timestamp + seconds_left - time()
    if time_to_sleep <= 0:
      self._tokens -= cost
      d.callback(True)
      return d

    def drain_and_return(val):
      self._tokens -= cost
      return d.callback(val)

    self.reactor.callLater(time_to_sleep, drain_and_return, True)
    return d

  def setCapacityAndFillRate(self, new_capacity, new_fill_rate):
    delta = float(new_capacity) - self.capacity
    self.capacity = float(new_capacity)
    self.fill_rate = float(new_fill_rate)
    self._tokens = delta + self._tokens

  @property
  def tokens(self):
    '''The tokens property will return the current number of tokens in the
    bucket.'''
    if self._tokens < self.capacity:
      now = time()
      delta = self.fill_rate * (now - self.timestamp)
      self._tokens = min(self.capacity, self._tokens + delta)
      self.timestamp = now
    return self._tokens

class PluginRegistrar(type):
  """Clever subclass detection hack that makes plugin loading trivial.
  To use this, define an abstract base class for plugin implementations
  that defines the plugin API. Give that base class a __metaclass__ of
  PluginRegistrar, and define a 'plugins = {}' class member. Subclasses
  defining a 'plugin_name' member will then appear in the plugins dict.
  """
  def __init__(classObj, name, bases, members):
    super(PluginRegistrar, classObj).__init__(name, bases, members)
    if hasattr(classObj, 'plugin_name'):
      classObj.plugins[classObj.plugin_name] = classObj

class TaggedSeries(object):
  @classmethod
  def parse(cls, path):
    tags = {}

    # if path is in openmetrics format: metric{tag="value",...}
    # https://github.com/RichiH/OpenMetrics
    if path[-2:] == '"}' and '{' in path:
      (metric, rawtags) = path[0:-1].split('{', 2)
      if not metric:
        raise Exception('Cannot parse path %s, no metric found' % path)

      while len(rawtags) > 0:
        m = re.match(r'([^=]+)="((?:[\\]["\\]|[^"\\])+)"(:?,|$)', rawtags)
        if not m:
          raise Exception('Cannot parse path %s, invalid segment %s' % (path, rawtags))

        tags[m.group(1)] = m.group(2).replace(r'\"', '"').replace(r'\\', '\\')
        rawtags = rawtags[len(m.group(0)):]

      tags['name'] = metric
      return cls(metric, tags)

    # path is a carbon path with optional tags: metric;tag=value;...
    segments = path.split(';')

    metric = segments[0]
    if not metric:
      raise Exception('Cannot parse path %s, no metric found' % path)

    for segment in segments[1:]:
      tag = segment.split('=', 1)
      if len(tag) != 2 or not tag[0]:
        raise Exception('Cannot parse path %s, invalid segment %s' % (path, segment))

      tags[tag[0]] = tag[1]

    tags['name'] = metric
    return cls(metric, tags)

  @staticmethod
  def format(tags):
    return tags.get('name', '') + ''.join(sorted([
      ';%s=%s' % (tag, value)
      for tag, value in tags.items()
      if tag != 'name'
    ]))

  @staticmethod
  def encode(metric, sep='.'):
    # if metric is tagged, encode it for storage in whisper etc
    if ';' in metric:
      metric_hash = sha256(metric.encode('utf8')).hexdigest()
      return sep.join(['_tagged', metric_hash[0:3], metric_hash[3:6], metric.replace('.', '-')])

    # metric isn't tagged, just replace dots with the separator and trim any leading separator
    return metric.replace('.', sep).lstrip(sep)

  def __init__(self, metric, tags, series_id=None):
    self.metric = metric
    self.tags = tags
    self.series_id = id

  @property
  def path(self):
    return self.__class__.format(self.tags)
<|endoftext|>


<commit_before>#
# Copyright 2016 Quantopian, Inc.
#
# 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 pandas as pd
import numpy as np
import warnings
from functools import wraps
from IPython.display import display

class NonMatchingTimezoneError(Exception):
    pass

def non_unique_bin_edges_error(func):
    """
    Give user a more informative error in case it is not possible
    to properly calculate quantiles on the input dataframe (factor)
    """
    def dec(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except ValueError as e:
            if 'Bin edges must be unique' in str(e):
                print("""
It's not possible to compute quantiles for the input provided.
This usually happens when the input contains too many identical
values so that they span more than one quantile. A typical example
is a factor with discrete values.
A workaround is to use 'bins' option instead of 'quantiles', or
to provide custom ranges to 'quantiles' option. 'bins' chooses the
buckets to be evenly spaced according to the values themselves, while
'quantiles' forces the buckets to have the same number of records.
Please see utils.get_clean_factor_and_forward_returns documentation.
                      """)
            raise
    return dec

@non_unique_bin_edges_error
def quantize_factor(factor_data, quantiles=5, bins=None, by_group=False):
    """
    Computes period wise factor quantiles.

    Parameters
    ----------
    factor_data : pd.DataFrame - MultiIndex
        A MultiIndex DataFrame indexed by date (level 0) and asset (level 1),
        containing the values for a single alpha factor, forward returns for each period,
        The factor quantile/bin that factor value belongs too, and (optionally) the group the
        asset belongs to.
    quantiles : int or sequence[float]
        Number of equal-sized quantile buckets to use in factor bucketing.
        Alternately sequence of quantiles, allowing non-equal-sized buckets
        e.g. [0, .10, .5, .90, 1.] or [.05, .5, .95]
        Only one of 'quantiles' or 'bins' can be not-None
    bins : int or sequence[float]
        Number of equal-width (valuewise) bins to use in factor bucketing.
        Alternately sequence of bin edges allowing for non-uniform bin width
        e.g. [-4, -2, -0.5, 0, 10]
        Only one of 'quantiles' or 'bins' can be not-None
    by_group : bool
        If True, compute quantile buckets separately for each group.

    Returns
    -------
    factor_quantile : pd.Series
        Factor quantiles indexed by date and asset.
    """

    def quantile_calc(x, _quantiles, _bins):
        if _quantiles is not None and _bins is None:
            return pd.qcut(x, _quantiles, labels=False) + 1
        elif _bins is not None and _quantiles is None:
            return pd.cut(x, _bins, labels=False) + 1
        raise ValueError('Either quantiles or bins should be provided')

    grouper = [factor_data.index.get_level_values('date')]
    if by_group:
        grouper.append('group')

    factor_quantile = factor_data.groupby(grouper)['factor'].apply(quantile_calc, quantiles, bins)
    factor_quantile.name = 'factor_quantile'

    return factor_quantile.dropna()

def compute_forward_returns(prices, periods=(1, 5, 10), filter_zscore=None):
    """
    Finds the N period forward returns (as percent change) for each asset provided.

    Parameters
    ----------
    prices : pd.DataFrame
        Pricing data to use in forward price calculation.
        Assets as columns, dates as index. Pricing data must
        span the factor analysis time period plus an additional buffer window
        that is greater than the maximum number of expected periods
        in the forward returns calculations.
    periods : sequence[int]
        periods to compute forward returns on.
    filter_zscore : int or float
        Sets forward returns greater than X standard deviations
        from the the mean to nan.
        Caution: this outlier filtering incorporates lookahead bias.

    Returns
    -------
    forward_returns : pd.DataFrame - MultiIndex
        Forward returns in indexed by date and asset.
        Separate column for each forward return window.
    """

    forward_returns = pd.DataFrame(index=pd.MultiIndex.from_product(
        [prices.index, prices.columns], names=['date', 'asset']))

    for period in periods:
        delta = prices.pct_change(period).shift(-period)

        if filter_zscore is not None:
            mask = abs(delta - delta.mean()) > (filter_zscore * delta.std())
            delta[mask] = np.nan

        forward_returns[period] = delta.stack()

    forward_returns.index = forward_returns.index.rename(['date', 'asset'])

    return forward_returns

def demean_forward_returns(factor_data, grouper=None):
    """
    Convert forward returns to returns relative to mean
    period wise all-universe or group returns.
    group-wise normalization incorporates the assumption of a
    group neutral portfolio constraint and thus allows allows the
    factor to be evaluated across groups.

    For example, if AAPL 5 period return is 0.1% and mean 5 period
    return for the Technology stocks in our universe was 0.5% in the
    same period, the group adjusted 5 period return for AAPL in this
    period is -0.4%.

    Parameters
    ----------
    factor_data : pd.DataFrame - MultiIndex
        Forward returns in indexed by date and asset.
        Separate column for each forward return window.
    grouper : list
        If True, demean according to group.

    Returns
    -------
    adjusted_forward_returns : pd.DataFrame - MultiIndex
        DataFrame of the same format as the input, but with each
        security's returns normalized by group.
    """

    factor_data = factor_data.copy()

    if not grouper:
        grouper = factor_data.index.get_level_values('date')

    cols = get_forward_returns_columns(factor_data.columns)
    factor_data[cols] = factor_data.groupby(grouper)[cols].transform(lambda x: x - x.mean())

    return factor_data

def print_table(table, name=None, fmt=None):
    """
    Pretty print a pandas DataFrame.

    Uses HTML output if running inside Jupyter Notebook, otherwise
    formatted text output.

    Parameters
    ----------
    table : pd.Series or pd.DataFrame
        Table to pretty-print.
    name : str, optional
        Table name to display in upper left corner.
    fmt : str, optional
        Formatter to use for displaying table elements.
        E.g. '{0:.2f}%' for displaying 100 as '100.00%'.
        Restores original setting after displaying.
    """
    if isinstance(table, pd.Series):
        table = pd.DataFrame(table)

    if isinstance(table, pd.DataFrame):
        table.columns.name = name

    prev_option = pd.get_option('display.float_format')
    if fmt is not None:
        pd.set_option('display.float_format', lambda x: fmt.format(x))

    display(table)

    if fmt is not None:
        pd.set_option('display.float_format', prev_option)

def get_clean_factor_and_forward_returns(factor,
                                         prices,
                                         groupby=None,
                                         by_group=False,
                                         quantiles=5,
                                         bins=None,
                                         periods=(1, 5, 10),
                                         filter_zscore=20,
                                         groupby_labels=None):
    """
    Formats the factor data, pricing data, and group mappings
    into a DataFrame that contains aligned MultiIndex
    indices of date and asset.

    Parameters
    ----------
    factor : pd.Series - MultiIndex
        A MultiIndex Series indexed by date (level 0) and asset (level 1), containing
        the values for a single alpha factor.
        ::
            -----------------------------------
                date    |    asset   |
            -----------------------------------
                        |   AAPL     |   0.5
                        -----------------------
                        |   BA       |  -1.1
                        -----------------------
            2014-01-01  |   CMG      |   1.7
                        -----------------------
                        |   DAL      |  -0.1
                        -----------------------
                        |   LULU     |   2.7
                        -----------------------

    prices : pd.DataFrame
        A wide form Pandas DataFrame indexed by date with assets
        in the columns. It is important to pass the
        correct pricing data in depending on what time of period your
        signal was generated so to avoid lookahead bias, or
        delayed calculations. Pricing data must span the factor
        analysis time period plus an additional buffer window
        that is greater than the maximum number of expected periods
        in the forward returns calculations.
        'Prices' must contain at least an entry for each date/asset
        combination in 'factor'. This entry must be the asset price
        at the time the asset factor value is computed and it will be
        considered the buy price for that asset at that date.
        'Prices' must also contain entries for dates following each
        date/asset combination in 'factor', as many more dates as the
        maximum value in 'periods'. The asset price after 'period'
        dates will be considered the sell price for that asset when
        computing 'period' forward returns.
    groupby : pd.Series - MultiIndex or dict
        Either A MultiIndex Series indexed by date and asset,
        containing the period wise group codes for each asset, or
        a dict of asset to group mappings. If a dict is passed,
        it is assumed that group mappings are unchanged for the
        entire time period of the passed factor data.
    by_group : bool
        If True, compute statistics separately for each group.
    quantiles : int or sequence[float]
        Number of equal-sized quantile buckets to use in factor bucketing.
        Alternately sequence of quantiles, allowing non-equal-sized buckets
        e.g. [0, .10, .5, .90, 1.] or [.05, .5, .95]
        Only one of 'quantiles' or 'bins' can be not-None
    bins : int or sequence[float]
        Number of equal-width (valuewise) bins to use in factor bucketing.
        Alternately sequence of bin edges allowing for non-uniform bin width
        e.g. [-4, -2, -0.5, 0, 10]
        Chooses the buckets to be evenly spaced according to the values themselves.
        Useful when the factor contains discrete values.
        Only one of 'quantiles' or 'bins' can be not-None
    periods : sequence[int]
        periods to compute forward returns on.
    filter_zscore : int or float
        Sets forward returns greater than X standard deviations
        from the the mean to nan.
        Caution: this outlier filtering incorporates lookahead bias.
    groupby_labels : dict
        A dictionary keyed by group code with values corresponding
        to the display name for each group.

    Returns
    -------
    merged_data : pd.DataFrame - MultiIndex
        A MultiIndex Series indexed by date (level 0) and asset (level 1),
        containing the values for a single alpha factor, forward returns for each period,
        The factor quantile/bin that factor value belongs too, and (optionally) the group the
        asset belongs to.
    """

    if factor.index.levels[0].tz != prices.index.tz:
        raise NonMatchingTimezoneError("The timezone of 'factor' is not the "
                                       "same as the timezone of 'prices'. See "
                                       "the pandas methods tz_localize and tz_convert.")

    merged_data = compute_forward_returns(prices, periods, filter_zscore)

    factor = factor.copy()
    factor.index = factor.index.rename(['date', 'asset'])
    merged_data['factor'] = factor

    if groupby is not None:
        if isinstance(groupby, dict):
            diff = set(factor.index.get_level_values(
                'asset')) - set(groupby.keys())
            if len(diff) > 0:
                raise KeyError(
                    "Assets {} not in group mapping".format(
                        list(diff)))

            ss = pd.Series(groupby)
            groupby = pd.Series(index=factor.index,
                                data=ss[factor.index.get_level_values(
                                    'asset')].values)

        if groupby_labels is not None:
            diff = set(groupby.values) - set(groupby_labels.keys())
            if len(diff) > 0:
                raise KeyError(
                    "groups {} not in passed group names".format(
                        list(diff)))

            sn = pd.Series(groupby_labels)
            groupby = pd.Series(index=factor.index,
                                data=sn[groupby.values].values)

        merged_data['group'] = groupby.astype('category')

    merged_data = merged_data.dropna()

    merged_data['factor_quantile'] = quantize_factor(merged_data,
                                                     quantiles,
                                                     bins,
                                                     by_group)

    merged_data = merged_data.dropna()

    return merged_data

def common_start_returns(factor,
                         prices,
                         before,
                         after,
                         cumulative=False,
                         mean_by_date=False,
                         demean=None):
    """
    A date and equity pair is extracted from each index row in the factor
    dataframe and for each of these pairs a return series is built starting
    from 'before' the date and ending 'after' the date specified in the pair.
    All those returns series are then aligned to a common index (-before to after)
    and returned as a single DataFrame

    Parameters
    ----------
    factor : pd.DataFrame
        DataFrame with at least date and equity as index, the columns are irrelevant
    prices : pd.DataFrame
        A wide form Pandas DataFrame indexed by date with assets
        in the columns. Pricing data should span the factor
        analysis time period plus/minus an additional buffer window
        corresponding to after/before period parameters.
    before:
        How many returns to load before factor date
    after:
        How many returns to load after factor date
    cumulative: bool, optional
        Return cumulative returns
    mean_by_date: bool, optional
        If True, compute mean returns for each date and return that
        instead of a return series for each asset
    demean: pd.DataFrame, optional
        DataFrame with at least date and equity as index, the columns are irrelevant
        For each date a list of equities is extracted from 'demean' index and used
        as universe to compute demeaned mean returns (long short portfolio)

    Returns
    -------
    aligned_returns : pd.DataFrame
        Dataframe containing returns series for each factor aligned to the same index:
        -before to after
    """

    if cumulative:
        returns = prices
    else:
        returns = prices.pct_change(axis=0)

    all_returns = []

    for timestamp, df in factor.groupby(level='date'):

        equities = df.index.get_level_values('asset')

        try:
            day_zero_index = returns.index.get_loc(timestamp)
        except KeyError:
            continue

        starting_index = max(day_zero_index - before, 0)
        ending_index = min(day_zero_index + after + 1,
                           len(returns.index))

        equities_slice = set(equities)
        if demean is not None:
            demean_equities = demean.loc[timestamp].index.get_level_values('asset')
            equities_slice |= set(demean_equities)

        series = returns.loc[returns.index[starting_index:ending_index],
                             equities_slice]
        series.index = range(starting_index - day_zero_index,
                             ending_index - day_zero_index)

        if cumulative:
            series = (series / series.loc[0, :]) - 1

        if demean is not None:
            mean   = series.loc[:, demean_equities].mean(axis=1)
            series = series.loc[:, equities]
            series = series.sub(mean, axis=0)

        if mean_by_date:
            series = series.mean(axis=1)

        all_returns.append(series)

    return pd.concat(all_returns, axis=1)

def rate_of_return(period_ret):
    """
    1-period Growth Rate: the average rate of 1-period returns
    """
    return period_ret.add(1).pow(1./period_ret.name).sub(1)

def std_conversion(period_std):
    """
    1-period standard deviation (or standard error) approximation

    Parameters
    ----------
    period_std: pd.DataFrame
        DataFrame containing standard deviation or standard error values
        with column headings representing the return period.

    Returns
    -------
    pd.DataFrame
        DataFrame in same format as input but with one-period
        standard deviation/error values.
    """
    period_len = period_std.name
    return period_std / np.sqrt(period_len)

def get_forward_returns_columns(columns):
    return columns[columns.astype('str').str.isdigit()]
<commit_msg>STY pep8 for utils.py<commit_after>#
# Copyright 2016 Quantopian, Inc.
#
# 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 pandas as pd
import numpy as np
from IPython.display import display

class NonMatchingTimezoneError(Exception):
    pass

def non_unique_bin_edges_error(func):
    """
    Give user a more informative error in case it is not possible
    to properly calculate quantiles on the input dataframe (factor)
    """
    def dec(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except ValueError as e:
            if 'Bin edges must be unique' in str(e):
                print("""
It's not possible to compute quantiles for the input provided.
This usually happens when the input contains too many identical
values so that they span more than one quantile. A typical example
is a factor with discrete values.
A workaround is to use 'bins' option instead of 'quantiles', or
to provide custom ranges to 'quantiles' option. 'bins' chooses the
buckets to be evenly spaced according to the values themselves, while
'quantiles' forces the buckets to have the same number of records.
Please see utils.get_clean_factor_and_forward_returns documentation.
                      """)
            raise
    return dec

@non_unique_bin_edges_error
def quantize_factor(factor_data, quantiles=5, bins=None, by_group=False):
    """
    Computes period wise factor quantiles.

    Parameters
    ----------
    factor_data : pd.DataFrame - MultiIndex
        A MultiIndex DataFrame indexed by date (level 0) and asset (level 1),
        containing the values for a single alpha factor, forward returns for
        each period, the factor quantile/bin that factor value belongs too, and
        (optionally) the group the asset belongs to.
    quantiles : int or sequence[float]
        Number of equal-sized quantile buckets to use in factor bucketing.
        Alternately sequence of quantiles, allowing non-equal-sized buckets
        e.g. [0, .10, .5, .90, 1.] or [.05, .5, .95]
        Only one of 'quantiles' or 'bins' can be not-None
    bins : int or sequence[float]
        Number of equal-width (valuewise) bins to use in factor bucketing.
        Alternately sequence of bin edges allowing for non-uniform bin width
        e.g. [-4, -2, -0.5, 0, 10]
        Only one of 'quantiles' or 'bins' can be not-None
    by_group : bool
        If True, compute quantile buckets separately for each group.

    Returns
    -------
    factor_quantile : pd.Series
        Factor quantiles indexed by date and asset.
    """

    def quantile_calc(x, _quantiles, _bins):
        if _quantiles is not None and _bins is None:
            return pd.qcut(x, _quantiles, labels=False) + 1
        elif _bins is not None and _quantiles is None:
            return pd.cut(x, _bins, labels=False) + 1
        raise ValueError('Either quantiles or bins should be provided')

    grouper = [factor_data.index.get_level_values('date')]
    if by_group:
        grouper.append('group')

    factor_quantile = factor_data.groupby(grouper)['factor'] \
        .apply(quantile_calc, quantiles, bins)
    factor_quantile.name = 'factor_quantile'

    return factor_quantile.dropna()

def compute_forward_returns(prices, periods=(1, 5, 10), filter_zscore=None):
    """
    Finds the N period forward returns (as percent change) for each asset
    provided.

    Parameters
    ----------
    prices : pd.DataFrame
        Pricing data to use in forward price calculation.
        Assets as columns, dates as index. Pricing data must
        span the factor analysis time period plus an additional buffer window
        that is greater than the maximum number of expected periods
        in the forward returns calculations.
    periods : sequence[int]
        periods to compute forward returns on.
    filter_zscore : int or float
        Sets forward returns greater than X standard deviations
        from the the mean to nan.
        Caution: this outlier filtering incorporates lookahead bias.

    Returns
    -------
    forward_returns : pd.DataFrame - MultiIndex
        Forward returns in indexed by date and asset.
        Separate column for each forward return window.
    """

    forward_returns = pd.DataFrame(index=pd.MultiIndex.from_product(
        [prices.index, prices.columns], names=['date', 'asset']))

    for period in periods:
        delta = prices.pct_change(period).shift(-period)

        if filter_zscore is not None:
            mask = abs(delta - delta.mean()) > (filter_zscore * delta.std())
            delta[mask] = np.nan

        forward_returns[period] = delta.stack()

    forward_returns.index = forward_returns.index.rename(['date', 'asset'])

    return forward_returns

def demean_forward_returns(factor_data, grouper=None):
    """
    Convert forward returns to returns relative to mean
    period wise all-universe or group returns.
    group-wise normalization incorporates the assumption of a
    group neutral portfolio constraint and thus allows allows the
    factor to be evaluated across groups.

    For example, if AAPL 5 period return is 0.1% and mean 5 period
    return for the Technology stocks in our universe was 0.5% in the
    same period, the group adjusted 5 period return for AAPL in this
    period is -0.4%.

    Parameters
    ----------
    factor_data : pd.DataFrame - MultiIndex
        Forward returns in indexed by date and asset.
        Separate column for each forward return window.
    grouper : list
        If True, demean according to group.

    Returns
    -------
    adjusted_forward_returns : pd.DataFrame - MultiIndex
        DataFrame of the same format as the input, but with each
        security's returns normalized by group.
    """

    factor_data = factor_data.copy()

    if not grouper:
        grouper = factor_data.index.get_level_values('date')

    cols = get_forward_returns_columns(factor_data.columns)
    factor_data[cols] = factor_data.groupby(grouper)[cols] \
        .transform(lambda x: x - x.mean())

    return factor_data

def print_table(table, name=None, fmt=None):
    """
    Pretty print a pandas DataFrame.

    Uses HTML output if running inside Jupyter Notebook, otherwise
    formatted text output.

    Parameters
    ----------
    table : pd.Series or pd.DataFrame
        Table to pretty-print.
    name : str, optional
        Table name to display in upper left corner.
    fmt : str, optional
        Formatter to use for displaying table elements.
        E.g. '{0:.2f}%' for displaying 100 as '100.00%'.
        Restores original setting after displaying.
    """
    if isinstance(table, pd.Series):
        table = pd.DataFrame(table)

    if isinstance(table, pd.DataFrame):
        table.columns.name = name

    prev_option = pd.get_option('display.float_format')
    if fmt is not None:
        pd.set_option('display.float_format', lambda x: fmt.format(x))

    display(table)

    if fmt is not None:
        pd.set_option('display.float_format', prev_option)

def get_clean_factor_and_forward_returns(factor,
                                         prices,
                                         groupby=None,
                                         by_group=False,
                                         quantiles=5,
                                         bins=None,
                                         periods=(1, 5, 10),
                                         filter_zscore=20,
                                         groupby_labels=None):
    """
    Formats the factor data, pricing data, and group mappings
    into a DataFrame that contains aligned MultiIndex
    indices of date and asset.

    Parameters
    ----------
    factor : pd.Series - MultiIndex
        A MultiIndex Series indexed by date (level 0) and asset (level 1),
        containing the values for a single alpha factor.
        ::
            -----------------------------------
                date    |    asset   |
            -----------------------------------
                        |   AAPL     |   0.5
                        -----------------------
                        |   BA       |  -1.1
                        -----------------------
            2014-01-01  |   CMG      |   1.7
                        -----------------------
                        |   DAL      |  -0.1
                        -----------------------
                        |   LULU     |   2.7
                        -----------------------

    prices : pd.DataFrame
        A wide form Pandas DataFrame indexed by date with assets
        in the columns. It is important to pass the
        correct pricing data in depending on what time of period your
        signal was generated so to avoid lookahead bias, or
        delayed calculations. Pricing data must span the factor
        analysis time period plus an additional buffer window
        that is greater than the maximum number of expected periods
        in the forward returns calculations.
        'Prices' must contain at least an entry for each date/asset
        combination in 'factor'. This entry must be the asset price
        at the time the asset factor value is computed and it will be
        considered the buy price for that asset at that date.
        'Prices' must also contain entries for dates following each
        date/asset combination in 'factor', as many more dates as the
        maximum value in 'periods'. The asset price after 'period'
        dates will be considered the sell price for that asset when
        computing 'period' forward returns.
    groupby : pd.Series - MultiIndex or dict
        Either A MultiIndex Series indexed by date and asset,
        containing the period wise group codes for each asset, or
        a dict of asset to group mappings. If a dict is passed,
        it is assumed that group mappings are unchanged for the
        entire time period of the passed factor data.
    by_group : bool
        If True, compute statistics separately for each group.
    quantiles : int or sequence[float]
        Number of equal-sized quantile buckets to use in factor bucketing.
        Alternately sequence of quantiles, allowing non-equal-sized buckets
        e.g. [0, .10, .5, .90, 1.] or [.05, .5, .95]
        Only one of 'quantiles' or 'bins' can be not-None
    bins : int or sequence[float]
        Number of equal-width (valuewise) bins to use in factor bucketing.
        Alternately sequence of bin edges allowing for non-uniform bin width
        e.g. [-4, -2, -0.5, 0, 10]
        Chooses the buckets to be evenly spaced according to the values
        themselves. Useful when the factor contains discrete values.
        Only one of 'quantiles' or 'bins' can be not-None
    periods : sequence[int]
        periods to compute forward returns on.
    filter_zscore : int or float
        Sets forward returns greater than X standard deviations
        from the the mean to nan.
        Caution: this outlier filtering incorporates lookahead bias.
    groupby_labels : dict
        A dictionary keyed by group code with values corresponding
        to the display name for each group.

    Returns
    -------
    merged_data : pd.DataFrame - MultiIndex
        A MultiIndex Series indexed by date (level 0) and asset (level 1),
        containing the values for a single alpha factor, forward returns for
        each period, the factor quantile/bin that factor value belongs to, and
        (optionally) the group the asset belongs to.
    """

    if factor.index.levels[0].tz != prices.index.tz:
        raise NonMatchingTimezoneError("The timezone of 'factor' is not the "
                                       "same as the timezone of 'prices'. See "
                                       "the pandas methods tz_localize and "
                                       "tz_convert.")

    merged_data = compute_forward_returns(prices, periods, filter_zscore)

    factor = factor.copy()
    factor.index = factor.index.rename(['date', 'asset'])
    merged_data['factor'] = factor

    if groupby is not None:
        if isinstance(groupby, dict):
            diff = set(factor.index.get_level_values(
                'asset')) - set(groupby.keys())
            if len(diff) > 0:
                raise KeyError(
                    "Assets {} not in group mapping".format(
                        list(diff)))

            ss = pd.Series(groupby)
            groupby = pd.Series(index=factor.index,
                                data=ss[factor.index.get_level_values(
                                    'asset')].values)

        if groupby_labels is not None:
            diff = set(groupby.values) - set(groupby_labels.keys())
            if len(diff) > 0:
                raise KeyError(
                    "groups {} not in passed group names".format(
                        list(diff)))

            sn = pd.Series(groupby_labels)
            groupby = pd.Series(index=factor.index,
                                data=sn[groupby.values].values)

        merged_data['group'] = groupby.astype('category')

    merged_data = merged_data.dropna()

    merged_data['factor_quantile'] = quantize_factor(merged_data,
                                                     quantiles,
                                                     bins,
                                                     by_group)

    merged_data = merged_data.dropna()

    return merged_data

def common_start_returns(factor,
                         prices,
                         before,
                         after,
                         cumulative=False,
                         mean_by_date=False,
                         demean=None):
    """
    A date and equity pair is extracted from each index row in the factor
    dataframe and for each of these pairs a return series is built starting
    from 'before' the date and ending 'after' the date specified in the pair.
    All those returns series are then aligned to a common index (-before to
    after) and returned as a single DataFrame

    Parameters
    ----------
    factor : pd.DataFrame
        DataFrame with at least date and equity as index, the columns are
        irrelevant
    prices : pd.DataFrame
        A wide form Pandas DataFrame indexed by date with assets
        in the columns. Pricing data should span the factor
        analysis time period plus/minus an additional buffer window
        corresponding to after/before period parameters.
    before:
        How many returns to load before factor date
    after:
        How many returns to load after factor date
    cumulative: bool, optional
        Return cumulative returns
    mean_by_date: bool, optional
        If True, compute mean returns for each date and return that
        instead of a return series for each asset
    demean: pd.DataFrame, optional
        DataFrame with at least date and equity as index, the columns are
        irrelevant. For each date a list of equities is extracted from 'demean'
        index and used as universe to compute demeaned mean returns (long short
        portfolio)

    Returns
    -------
    aligned_returns : pd.DataFrame
        Dataframe containing returns series for each factor aligned to the same
        index: -before to after
    """

    if cumulative:
        returns = prices
    else:
        returns = prices.pct_change(axis=0)

    all_returns = []

    for timestamp, df in factor.groupby(level='date'):

        equities = df.index.get_level_values('asset')

        try:
            day_zero_index = returns.index.get_loc(timestamp)
        except KeyError:
            continue

        starting_index = max(day_zero_index - before, 0)
        ending_index = min(day_zero_index + after + 1,
                           len(returns.index))

        equities_slice = set(equities)
        if demean is not None:
            demean_equities = demean.loc[timestamp] \
                .index.get_level_values('asset')
            equities_slice |= set(demean_equities)

        series = returns.loc[returns.index[starting_index:ending_index],
                             equities_slice]
        series.index = range(starting_index - day_zero_index,
                             ending_index - day_zero_index)

        if cumulative:
            series = (series / series.loc[0, :]) - 1

        if demean is not None:
            mean = series.loc[:, demean_equities].mean(axis=1)
            series = series.loc[:, equities]
            series = series.sub(mean, axis=0)

        if mean_by_date:
            series = series.mean(axis=1)

        all_returns.append(series)

    return pd.concat(all_returns, axis=1)

def rate_of_return(period_ret):
    """
    1-period Growth Rate: the average rate of 1-period returns
    """
    return period_ret.add(1).pow(1./period_ret.name).sub(1)

def std_conversion(period_std):
    """
    1-period standard deviation (or standard error) approximation

    Parameters
    ----------
    period_std: pd.DataFrame
        DataFrame containing standard deviation or standard error values
        with column headings representing the return period.

    Returns
    -------
    pd.DataFrame
        DataFrame in same format as input but with one-period
        standard deviation/error values.
    """
    period_len = period_std.name
    return period_std / np.sqrt(period_len)

def get_forward_returns_columns(columns):
    return columns[columns.astype('str').str.isdigit()]
<|endoftext|>


<commit_before>#!/usr/bin/env python2.7
# -*- mode: python; coding: utf-8 -*-

from __future__ import print_function, division, absolute_import

from memory_profiler import profile
import numpy as np
from astropy import constants as const
from astropy.io import fits
from pyuvdata import UVData

@profile
def read_uvfits():
    filename = '/Volumes/Data1/mwa_uvfits/1066571272.uvfits'

    # first test uvdata.read_uvfits. First read metadata then full data
    uv_obj = UVData()
    uv_obj.read_uvfits(filename, metadata_only=True)
    uv_obj.read_uvfits_data(filename)
    del(uv_obj)

    # now test details with astropy
    hdu_list = fits.open(filename, memmap=True)
    vis_hdu = hdu_list[0]

    # only read in times, then uvws, then visibilities
    time0_array = vis_hdu.data.par('date')
    uvw_array = (np.array(np.stack((vis_hdu.data.par('UU'),
                                    vis_hdu.data.par('VV'),
                                    vis_hdu.data.par('WW')))) * const.c.to('m/s').value).T

    if vis_hdu.header['NAXIS'] == 7:

        data_array = (vis_hdu.data.data[:, 0, 0, :, :, :, 0] +
                      1j * vis_hdu.data.data[:, 0, 0, :, :, :, 1])
    else:
        data_array = (vis_hdu.data.data[:, 0, 0, :, :, 0] +
                      1j * vis_hdu.data.data[:, 0, 0, :, :, 1])
        data_array = data_array[:, np.newaxis, :, :]

    # test for releasing resources
    del(time0_array)
    del(uvw_array)
    del(data_array)

    # release file handles
    del(vis_hdu)
    del(hdu_list)
    del(filename)

    return

if __name__ == '__main__':
    read_uvfits()
<commit_msg>add slicing test to memtest script<commit_after>#!/usr/bin/env python2.7
# -*- mode: python; coding: utf-8 -*-

from __future__ import print_function, division, absolute_import

from memory_profiler import profile
import numpy as np
from astropy import constants as const
from astropy.io import fits
from pyuvdata import UVData

@profile
def read_uvfits():
    filename = '/Volumes/Data1/mwa_uvfits/1066571272.uvfits'

    # first test uvdata.read_uvfits. First read metadata then full data
    uv_obj = UVData()
    uv_obj.read_uvfits(filename, metadata_only=True)
    uv_obj.read_uvfits_data(filename)
    del(uv_obj)

    # now test details with astropy
    hdu_list = fits.open(filename, memmap=True)
    vis_hdu = hdu_list[0]

    # only read in times, then uvws, then visibilities
    time0_array = vis_hdu.data.par('date')
    uvw_array = (np.array(np.stack((vis_hdu.data.par('UU'),
                                    vis_hdu.data.par('VV'),
                                    vis_hdu.data.par('WW')))) * const.c.to('m/s').value).T

    if vis_hdu.header['NAXIS'] == 7:

        data_array = (vis_hdu.data.data[:, 0, 0, :, :, :, 0] +
                      1j * vis_hdu.data.data[:, 0, 0, :, :, :, 1])
    else:
        data_array = (vis_hdu.data.data[:, 0, 0, :, :, 0] +
                      1j * vis_hdu.data.data[:, 0, 0, :, :, 1])
        data_array = data_array[:, np.newaxis, :, :]

    # test for releasing resources
    del(time0_array)
    del(uvw_array)
    del(data_array)

    # release file handles
    del(vis_hdu)
    del(hdu_list)

    # now test reading a slice of the data
    hdu_list = fits.open(filename, memmap=True)
    vis_hdu = hdu_list[0]
    Nfreqs = vis_hdu.header['NAXIS4']
    freq_index = int(Nfreqs / 2)

    if vis_hdu.header['NAXIS'] == 7:

        data_slice = (vis_hdu.data.data[:, 0, 0, :, 0:freq_index, :, 0] +
                      1j * vis_hdu.data.data[:, 0, 0, :, 0:freq_index, :, 1])
    else:
        data_slice = (vis_hdu.data.data[:, 0, 0, 0:freq_index, :, 0] +
                      1j * vis_hdu.data.data[:, 0, 0, 0:freq_index, :, 1])
        data_slice = data_slice[:, np.newaxis, :, :]

    del(data_slice)

    del(vis_hdu)
    del(hdu_list)
    del(filename)

    return

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