<commit_before>#!/usr/bin/env python
from __future__ import print_function

import sys, os, datetime, time, shutil, tempfile, subprocess, random
from optparse import OptionParser
from ._backup import *

info = "'rethinkdb index-rebuild' recreates outdated secondary indexes in a cluster.\n" + \
       "  This should be used after upgrading to a newer version of rethinkdb.  There\n" + \
       "  will be a notification in the web UI if any secondary indexes are out-of-date."
usage = "rethinkdb index-rebuild [-c HOST:PORT] [-a AUTH_KEY] [-n NUM] [-r (DB | DB.TABLE)]..."

# Prefix used for indexes that are being rebuilt
temp_index_prefix = '$reql_temp_index$_'

def print_restore_help():
    print(info)
    print(usage)
    print("")
    print("  FILE                             the archive file to restore data from")
    print("  -h [ --help ]                    print this help")
    print("  -c [ --connect ] HOST:PORT       host and client port of a rethinkdb node to connect")
    print("                                   to (defaults to localhost:28015)")
    print("  -a [ --auth ] AUTH_KEY           authorization key for rethinkdb clients")
    print("  -r [ --rebuild ] (DB | DB.TABLE) the databases or tables to rebuild indexes on")
    print("                                   (defaults to all databases and tables)")
    print("  -n NUM                           the number of concurrent indexes to rebuild")
    print("                                   (defaults to 1)")
    print("")
    print("EXAMPLES:")
    print("rethinkdb index-rebuild -c mnemosyne:39500")
    print("  rebuild all outdated secondary indexes from the cluster on the host 'mnemosyne',")
    print("  one at a time")
    print("")
    print("rethinkdb index-rebuild -r test -r production.users -n 5")
    print("  rebuild all outdated secondary indexes from a local cluster on all tables in the")
    print("  'test' database as well as the 'production.users' table, five at a time")

def parse_options():
    parser = OptionParser(add_help_option=False, usage=usage)
    parser.add_option("-c", "--connect", dest="host", metavar="HOST:PORT", default="localhost:28015", type="string")
    parser.add_option("-a", "--auth", dest="auth_key", metavar="KEY", default="", type="string")
    parser.add_option("-r", "--rebuild", dest="tables", metavar="DB | DB.TABLE", default=[], action="append", type="string")
    parser.add_option("-n", dest="concurrent", metavar="NUM", default=1, type="int")
    parser.add_option("--debug", dest="debug", default=False, action="store_true")
    parser.add_option("-h", "--help", dest="help", default=False, action="store_true")
    (options, args) = parser.parse_args()

    if options.help:
        print_restore_help()
        exit(0)

    # Check validity of arguments
    if len(args) != 0:
        raise RuntimeError("Error: No positional arguments supported")

    res = { }

    # Verify valid host:port --connect option
    (res["host"], res["port"]) = parse_connect_option(options.host)

    # Verify valid --import options
    res["tables"] = parse_db_table_options(options.tables)

    res["auth_key"] = options.auth_key
    res["concurrent"] = options.concurrent
    res["debug"] = options.debug
    return res

def print_progress(ratio):
    total_width = 40
    done_width = int(ratio * total_width)
    equals = "=" * done_width
    spaces = " " * (total_width - done_width)
    percent = int(100 * ratio)
    print("\r[%s%s] %3d%%" % (equals, spaces, percent), end='')
    sys.stdout.flush()

def new_connection(conn_store, options):
    # Check if the connection is up
    try:
        r.expr(0).run(conn_store[0])
    except:
        conn_store[0] = r.connect(options["host"], options["port"])
    return conn_store[0]

def get_table_outdated_indexes(conn, db, table):
    return r.db(db).table(table).index_status().filter(lambda i: i['outdated'])['index'].run(conn)

def get_outdated_indexes(progress, conn, db_tables):
    res = [ ]
    if len(db_tables) == 0:
        dbs = r.db_list().run(conn)
        db_tables = [(db, None) for db in dbs]

    for db_table in db_tables:
        if db_table[1] is not None:
            table_list = [db_table[1]]
        else:
            table_list = r.db(db_table[0]).table_list().run(conn)
        for table in table_list:
            outdated_indexes = get_table_outdated_indexes(conn, db_table[0], table)

            for index in outdated_indexes:
                res.append({'db': db_table[0], 'table': table, 'name': index})
    return res

def drop_outdated_temp_indexes(progress, conn, indexes):
    indexes_to_drop = [i for i in indexes if i['name'].find(temp_index_prefix) == 0]
    for index in indexes_to_drop:
        r.db(index['db']).table(index['table']).index_drop(index['name']).run(conn)
        indexes.remove(index)

def create_temp_index(progress, conn, index):
    # If this index is already being rebuilt, don't try to recreate it
    extant_indexes = r.db(index['db']).table(index['table']).index_status().map(lambda i: i['index']).run(conn)
    if index['temp_name'] not in extant_indexes:
        index_fn = r.db(index['db']).table(index['table']).index_status(index['name']).nth(0)['function']
        r.db(index['db']).table(index['table']).index_create(index['temp_name'], index_fn).run(conn)

def get_index_progress(progress, conn, index):
    status = r.db(index['db']).table(index['table']).index_status(index['temp_name']).nth(0).run(conn)
    index['function'] = status['function']
    if status['ready']:
        return None
    else:
        processed = float(status.get('blocks_processed', 0))
        total = float(status.get('blocks_total', 1))
        if total != 0:
            return processed / total
        else:
            return 0.0

def rename_index(progress, conn, index):
    r.db(index['db']).table(index['table']).index_rename(index['temp_name'], index['name'], overwrite=True).run(conn)

def check_index_renamed(progress, conn, index):
    status = r.db(index['db']).table(index['table']).index_status(index['name']).nth(0).run(conn)
    if status['outdated'] or status['ready'] != True or status['function'] != index['function']:
        raise RuntimeError("Error: failed to rename `%s.%s` temporary index for `%s`" % \
                           (index['db'], index['table'], index['name']))

def rebuild_indexes(options):
    conn_store = [r.connect(options['host'], options['port'])]
    conn_fn = lambda: new_connection(conn_store, options)

    indexes_to_build = rdb_call_wrapper(conn_fn, "get outdated indexes", get_outdated_indexes, options["tables"])
    indexes_in_progress = [ ]

    # Drop any outdated indexes with the temp_index_prefix
    rdb_call_wrapper(conn_fn, "drop temporary outdated indexes", drop_outdated_temp_indexes, indexes_to_build)

    random.shuffle(indexes_to_build)
    total_indexes = len(indexes_to_build)
    indexes_completed = 0
    progress_ratio = 0.0
    highest_progress = 0.0

    print("Rebuilding %d indexes: %s" % (total_indexes,
                                         ", ".join(map(lambda i: "`%s.%s:%s`" % (i['db'], i['table'], i['name']),
                                                       indexes_to_build))))

    while len(indexes_to_build) > 0 or len(indexes_in_progress) > 0:
        # Make sure we're running the right number of concurrent index rebuilds
        while len(indexes_to_build) > 0 and len(indexes_in_progress) < options["concurrent"]:
            index = indexes_to_build.pop()
            indexes_in_progress.append(index)
            index['temp_name'] = temp_index_prefix + index['name']
            index['progress'] = 0
            index['ready'] = False

            try:
                rdb_call_wrapper(conn_fn, "create `%s.%s` index `%s`" % (index['db'], index['table'], index['name']),
                                 create_temp_index, index)
            except RuntimeError as ex:
                # This may be caused by a suprious failure (see github issue #2904), ignore if so
                if ex.message != "ReQL error during 'create `%s.%s` index `%s`': Index `%s` already exists on table `%s.%s`." % \
                                 (index['db'], index['table'], index['name'], index['temp_name'], index['db'], index['table']):
                    raise

        # Report progress
        highest_progress = max(highest_progress, progress_ratio)
        print_progress(highest_progress)

        # Check the status of indexes in progress
        progress_ratio = 0.0
        for index in indexes_in_progress:
            index_progress = rdb_call_wrapper(conn_fn, "progress `%s.%s` index `%s`" % (index['db'], index['table'], index['name']),
                                              get_index_progress, index)
            if index_progress is None:
                index['ready'] = True
                try:
                    rdb_call_wrapper(conn_fn, "rename `%s.%s` index `%s`" % (index['db'], index['table'], index['name']),
                                     rename_index, index)
                except r.RqlRuntimeError as ex:
                    # This may be caused by a spurious failure (see github issue #2904), check if it actually succeeded
                    if ex.message != "ReQL error during 'rename `%s.%s` index `%s`': Index `%s` does not exist on table `%s.%s`." % \
                                     (index['db'], index['table'], index['name'], index['temp_name'], index['db'], index['table']):
                        raise
                    rdb_call_wrapper(conn_fn, "check rename `%s.%s` index `%s`" % (index['db'], index['table'], index['name']),
                                     check_index_renamed, index)
            else:
                progress_ratio += index_progress / total_indexes

        indexes_in_progress = [index for index in indexes_in_progress if not index['ready']]
        indexes_completed = total_indexes - len(indexes_to_build) - len(indexes_in_progress)
        progress_ratio += float(indexes_completed) / total_indexes

        if len(indexes_in_progress) == options['concurrent'] or \
           (len(indexes_in_progress) > 0 and len(indexes_to_build) == 0):
            # Short sleep to keep from killing the CPU
            time.sleep(0.1)

    # Make sure the progress bar says we're done and get past the progress bar line
    print_progress(1.0)
    print("")

def main():
    try:
        options = parse_options()
    except RuntimeError as ex:
        print("Usage: %s" % usage, file=sys.stderr)
        print(ex, file=sys.stderr)
        return 1

    try:
        start_time = time.time()
        rebuild_indexes(options)
    except RuntimeError as ex:
        print(ex, file=sys.stderr)
        return 1
    print("  Done (%d seconds)" % (time.time() - start_time))
    return 0

if __name__ == "__main__":
    exit(main())
<commit_msg>making index-rebuild handle auth-key<commit_after>#!/usr/bin/env python
from __future__ import print_function

import sys, os, datetime, time, shutil, tempfile, subprocess, random
from optparse import OptionParser
from ._backup import *

info = "'rethinkdb index-rebuild' recreates outdated secondary indexes in a cluster.\n" + \
       "  This should be used after upgrading to a newer version of rethinkdb.  There\n" + \
       "  will be a notification in the web UI if any secondary indexes are out-of-date."
usage = "rethinkdb index-rebuild [-c HOST:PORT] [-a AUTH_KEY] [-n NUM] [-r (DB | DB.TABLE)]..."

# Prefix used for indexes that are being rebuilt
temp_index_prefix = '$reql_temp_index$_'

def print_restore_help():
    print(info)
    print(usage)
    print("")
    print("  FILE                             the archive file to restore data from")
    print("  -h [ --help ]                    print this help")
    print("  -c [ --connect ] HOST:PORT       host and client port of a rethinkdb node to connect")
    print("                                   to (defaults to localhost:28015)")
    print("  -a [ --auth ] AUTH_KEY           authorization key for rethinkdb clients")
    print("  -r [ --rebuild ] (DB | DB.TABLE) the databases or tables to rebuild indexes on")
    print("                                   (defaults to all databases and tables)")
    print("  -n NUM                           the number of concurrent indexes to rebuild")
    print("                                   (defaults to 1)")
    print("")
    print("EXAMPLES:")
    print("rethinkdb index-rebuild -c mnemosyne:39500")
    print("  rebuild all outdated secondary indexes from the cluster on the host 'mnemosyne',")
    print("  one at a time")
    print("")
    print("rethinkdb index-rebuild -r test -r production.users -n 5")
    print("  rebuild all outdated secondary indexes from a local cluster on all tables in the")
    print("  'test' database as well as the 'production.users' table, five at a time")

def parse_options():
    parser = OptionParser(add_help_option=False, usage=usage)
    parser.add_option("-c", "--connect", dest="host", metavar="HOST:PORT", default="localhost:28015", type="string")
    parser.add_option("-a", "--auth", dest="auth_key", metavar="KEY", default="", type="string")
    parser.add_option("-r", "--rebuild", dest="tables", metavar="DB | DB.TABLE", default=[], action="append", type="string")
    parser.add_option("-n", dest="concurrent", metavar="NUM", default=1, type="int")
    parser.add_option("--debug", dest="debug", default=False, action="store_true")
    parser.add_option("-h", "--help", dest="help", default=False, action="store_true")
    (options, args) = parser.parse_args()

    if options.help:
        print_restore_help()
        exit(0)

    # Check validity of arguments
    if len(args) != 0:
        raise RuntimeError("Error: No positional arguments supported")

    res = { }

    # Verify valid host:port --connect option
    (res["host"], res["port"]) = parse_connect_option(options.host)

    # Verify valid --import options
    res["tables"] = parse_db_table_options(options.tables)

    res["auth_key"] = options.auth_key
    res["concurrent"] = options.concurrent
    res["debug"] = options.debug
    return res

def print_progress(ratio):
    total_width = 40
    done_width = int(ratio * total_width)
    equals = "=" * done_width
    spaces = " " * (total_width - done_width)
    percent = int(100 * ratio)
    print("\r[%s%s] %3d%%" % (equals, spaces, percent), end='')
    sys.stdout.flush()

def do_connect(options):
    try:
        return r.connect(options['host'], options['port'], auth_key=options['auth_key'])
    except (r.RqlError, r.RqlDriverError) as ex:
        raise RuntimeError("Error when connecting: %s" % ex.message)

def new_connection(conn_store, options):
    # Check if the connection is up
    try:
        r.expr(0).run(conn_store[0])
    except:
        conn_store[0] = do_connect(options)
    return conn_store[0]

def get_table_outdated_indexes(conn, db, table):
    return r.db(db).table(table).index_status().filter(lambda i: i['outdated'])['index'].run(conn)

def get_outdated_indexes(progress, conn, db_tables):
    res = [ ]
    if len(db_tables) == 0:
        dbs = r.db_list().run(conn)
        db_tables = [(db, None) for db in dbs]

    for db_table in db_tables:
        if db_table[1] is not None:
            table_list = [db_table[1]]
        else:
            table_list = r.db(db_table[0]).table_list().run(conn)
        for table in table_list:
            outdated_indexes = get_table_outdated_indexes(conn, db_table[0], table)

            for index in outdated_indexes:
                res.append({'db': db_table[0], 'table': table, 'name': index})
    return res

def drop_outdated_temp_indexes(progress, conn, indexes):
    indexes_to_drop = [i for i in indexes if i['name'].find(temp_index_prefix) == 0]
    for index in indexes_to_drop:
        r.db(index['db']).table(index['table']).index_drop(index['name']).run(conn)
        indexes.remove(index)

def create_temp_index(progress, conn, index):
    # If this index is already being rebuilt, don't try to recreate it
    extant_indexes = r.db(index['db']).table(index['table']).index_status().map(lambda i: i['index']).run(conn)
    if index['temp_name'] not in extant_indexes:
        index_fn = r.db(index['db']).table(index['table']).index_status(index['name']).nth(0)['function']
        r.db(index['db']).table(index['table']).index_create(index['temp_name'], index_fn).run(conn)

def get_index_progress(progress, conn, index):
    status = r.db(index['db']).table(index['table']).index_status(index['temp_name']).nth(0).run(conn)
    index['function'] = status['function']
    if status['ready']:
        return None
    else:
        processed = float(status.get('blocks_processed', 0))
        total = float(status.get('blocks_total', 1))
        if total != 0:
            return processed / total
        else:
            return 0.0

def rename_index(progress, conn, index):
    r.db(index['db']).table(index['table']).index_rename(index['temp_name'], index['name'], overwrite=True).run(conn)

def check_index_renamed(progress, conn, index):
    status = r.db(index['db']).table(index['table']).index_status(index['name']).nth(0).run(conn)
    if status['outdated'] or status['ready'] != True or status['function'] != index['function']:
        raise RuntimeError("Error: failed to rename `%s.%s` temporary index for `%s`" % \
                           (index['db'], index['table'], index['name']))

def rebuild_indexes(options):
    conn_store = [do_connect(options)]
    conn_fn = lambda: new_connection(conn_store, options)

    indexes_to_build = rdb_call_wrapper(conn_fn, "get outdated indexes", get_outdated_indexes, options["tables"])
    indexes_in_progress = [ ]

    # Drop any outdated indexes with the temp_index_prefix
    rdb_call_wrapper(conn_fn, "drop temporary outdated indexes", drop_outdated_temp_indexes, indexes_to_build)

    random.shuffle(indexes_to_build)
    total_indexes = len(indexes_to_build)
    indexes_completed = 0
    progress_ratio = 0.0
    highest_progress = 0.0

    print("Rebuilding %d indexes: %s" % (total_indexes,
                                         ", ".join(map(lambda i: "`%s.%s:%s`" % (i['db'], i['table'], i['name']),
                                                       indexes_to_build))))

    while len(indexes_to_build) > 0 or len(indexes_in_progress) > 0:
        # Make sure we're running the right number of concurrent index rebuilds
        while len(indexes_to_build) > 0 and len(indexes_in_progress) < options["concurrent"]:
            index = indexes_to_build.pop()
            indexes_in_progress.append(index)
            index['temp_name'] = temp_index_prefix + index['name']
            index['progress'] = 0
            index['ready'] = False

            try:
                rdb_call_wrapper(conn_fn, "create `%s.%s` index `%s`" % (index['db'], index['table'], index['name']),
                                 create_temp_index, index)
            except RuntimeError as ex:
                # This may be caused by a suprious failure (see github issue #2904), ignore if so
                if ex.message != "ReQL error during 'create `%s.%s` index `%s`': Index `%s` already exists on table `%s.%s`." % \
                                 (index['db'], index['table'], index['name'], index['temp_name'], index['db'], index['table']):
                    raise

        # Report progress
        highest_progress = max(highest_progress, progress_ratio)
        print_progress(highest_progress)

        # Check the status of indexes in progress
        progress_ratio = 0.0
        for index in indexes_in_progress:
            index_progress = rdb_call_wrapper(conn_fn, "progress `%s.%s` index `%s`" % (index['db'], index['table'], index['name']),
                                              get_index_progress, index)
            if index_progress is None:
                index['ready'] = True
                try:
                    rdb_call_wrapper(conn_fn, "rename `%s.%s` index `%s`" % (index['db'], index['table'], index['name']),
                                     rename_index, index)
                except r.RqlRuntimeError as ex:
                    # This may be caused by a spurious failure (see github issue #2904), check if it actually succeeded
                    if ex.message != "ReQL error during 'rename `%s.%s` index `%s`': Index `%s` does not exist on table `%s.%s`." % \
                                     (index['db'], index['table'], index['name'], index['temp_name'], index['db'], index['table']):
                        raise
                    rdb_call_wrapper(conn_fn, "check rename `%s.%s` index `%s`" % (index['db'], index['table'], index['name']),
                                     check_index_renamed, index)
            else:
                progress_ratio += index_progress / total_indexes

        indexes_in_progress = [index for index in indexes_in_progress if not index['ready']]
        indexes_completed = total_indexes - len(indexes_to_build) - len(indexes_in_progress)
        progress_ratio += float(indexes_completed) / total_indexes

        if len(indexes_in_progress) == options['concurrent'] or \
           (len(indexes_in_progress) > 0 and len(indexes_to_build) == 0):
            # Short sleep to keep from killing the CPU
            time.sleep(0.1)

    # Make sure the progress bar says we're done and get past the progress bar line
    print_progress(1.0)
    print("")

def main():
    try:
        options = parse_options()
    except RuntimeError as ex:
        print("Usage: %s" % usage, file=sys.stderr)
        print(ex, file=sys.stderr)
        return 1

    try:
        start_time = time.time()
        rebuild_indexes(options)
    except RuntimeError as ex:
        print(ex, file=sys.stderr)
        return 1
    print("  Done (%d seconds)" % (time.time() - start_time))
    return 0

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


<commit_before>from astropy.extern.six.moves.urllib import parse as urlparse
from astropy.extern.six.moves import map, zip
from astropy.extern.six import StringIO
from astropy.extern import six

from collections import defaultdict

from astropy.io import ascii
from astropy.table import Table
from astropy import units as u
from bs4 import BeautifulSoup

from ..query import BaseQuery
from ..utils import commons
from ..utils import prepend_docstr_noreturns
from ..utils import async_to_sync
from . import ATOMIC_LINE_LIST_URL, TIMEOUT

__all__ = ['AtomicLineList', 'AtomicLineListClass']

@async_to_sync
class AtomicLineListClass(BaseQuery):
    FORM_URL = ATOMIC_LINE_LIST_URL()
    TIMEOUT = TIMEOUT()

    def __init__(self):
        BaseQuery.__init__(self)
        self._default_form_values = None

    def query_object(self, wavelength_range=None, wavelength_type=None,
                     wavelength_accuracy=None, element_spectrum=None):
        """
        Parameters
        ----------
        wavelength_range : pair of `astropy.units.Unit` values

        wavelength_type : str
            Either 'Air' or 'Vacuum'

        wavelength_accuracy :

        element_spectrum :

        Returns
        -------
        result : `~astropy.table.Table`
            The result of the query as a `~astropy.table.Table` object.

        """
        response = self.query_object_async(
            wavelength_range, wavelength_type, wavelength_accuracy, element_spectrum)
        table = self._parse_result(response)
        return table

    @prepend_docstr_noreturns(query_object.__doc__)
    def query_object_async(self, wavelength_range=None, wavelength_type='',
                           wavelength_accuracy=None, element_spectrum=None):
        """
        Returns
        -------
        response : `requests.Response`
            The HTTP response returned from the service.
        """
        if self._default_form_values is None:
            response = commons.send_request(self.FORM_URL, {}, self.TIMEOUT, 'GET')
            bs = BeautifulSoup(response.text)
            form = bs.find('form')
            self._default_form_values = self._get_default_form_values(form)
        default_values = self._default_form_values
        wltype = (wavelength_type or default_values.get('air', '')).lower()
        if wltype in ('air', 'vacuum'):
            air = wltype.capitalize()
        else:
            raise ValueError('parameter wavelength_type must be either "air" or "vacuum".')
        wlrange = wavelength_range or []
        if len(wlrange) not in (0, 2):
            raise ValueError(
                'length of `wavelength_range` must be 2 or 0, but is: {}'.format(
                    len(wlrange)))
        # convert wavelengths in incoming wavelength range to Angstroms
        wlrange_in_angstroms = (wl.to(u.Angstrom, equivalencies=u.spectral()).value for wl in wlrange)
        input = {
            'wavl': '-'.join(map(str, wlrange_in_angstroms)),
            'wave': 'Angstrom',
            'air': air,
            'wacc': wavelength_accuracy,
            'elmion': element_spectrum,
            # TODO: add support for more parameters
        }
        response = self._submit_form(input)
        return response

    def _parse_result(self, response):
        data = StringIO(BeautifulSoup(response.text).find('pre').text.strip())
        # `header` is e.g. "u'-LAMBDA-VAC-ANG-|-SPECTRUM--|TT|--------TERM---------|---J-J---|----LEVEL ENERGY--CM-1----'"
        header = data.readline().strip().strip('|')
        # `colnames` is then "[u'LAMBDA VAC ANG', u'SPECTRUM', u'TT', u'TERM', u'J J', u'LEVEL ENERGY  CM 1']"
        colnames = [colname.strip('-').replace('-', ' ') for colname in header.split('|') if colname.strip()]
        indices = [i for i, c in enumerate(header) if c == '|']
        input = []
        for line in data:
            row = []
            for start, end in zip([0]+indices, indices+[None]):
                value = line[start:end].strip()
                if value:
                    row.append(value)
            if row:
                input.append('\t'.join(row))
        if input:
            return ascii.read(input, data_start=0, delimiter='\t', names=colnames)
        else:
            # return an empty table if the query yielded no results
            return Table()

    def _submit_form(self, input=None):
        """Fill out the form of the SkyView site and submit it with the
        values given in `input` (a dictionary where the keys are the form
        element's names and the values are their respective values).

        """
        if input is None:
            input = {}
        response = commons.send_request(self.FORM_URL, {}, self.TIMEOUT, 'GET')
        bs = BeautifulSoup(response.text)
        form = bs.find('form')
        # cache the default values to save HTTP traffic
        if self._default_form_values is None:
            self._default_form_values = self._get_default_form_values(form)
        # only overwrite payload's values if the `input` value is not None
        # to avoid overwriting of the form's default values
        payload = self._default_form_values.copy()
        for k, v in six.iteritems(input):
            if v is not None:
                payload[k] = v
        url = urlparse.urljoin(self.FORM_URL, form.get('action'))
        response = commons.send_request(url, payload, self.TIMEOUT)
        return response

    def _get_default_form_values(self, form):
        """Return the already selected values of a given form (a BeautifulSoup
        form node) as a dict.

        """
        res = defaultdict(list)
        for elem in form.find_all(['input', 'select']):
            key = elem.get('name')
            value = None
            # ignore the submit and reset buttons
            if elem.get('type') in ['submit', 'reset']:
                continue
            # check boxes: enabled boxes have the value "on" if not specificed
            # otherwise. Found out by debugging, perhaps not documented.
            if elem.get('type') == 'checkbox' and elem.get('checked') in ["", "checked"]:
                value = elem.get('value', 'on')
            # radio buttons and simple input fields
            if elem.get('type') == 'radio' and\
                    elem.get('checked') in ["", "checked"] or\
                    elem.get('type') in [None, 'text']:
                value = elem.get('value')
            # dropdown menu, multi-section possible
            if elem.name == 'select':
                for option in elem.find_all('option'):
                    if option.get('selected') == '':
                        value = option.get('value', option.text.strip())
            if value and value not in [None, u'None', u'null']:
                res[key].append(value)
        # avoid values with size 1 lists
        d = dict(res)
        for k, v in six.iteritems(d):
            if len(v) == 1:
                d[k] = v[0]
        return d

AtomicLineList = AtomicLineListClass()
<commit_msg>atomic: added more comments to explain _parse_result method<commit_after>from astropy.extern.six.moves.urllib import parse as urlparse
from astropy.extern.six.moves import map, zip
from astropy.extern.six import StringIO
from astropy.extern import six

from collections import defaultdict

from astropy.io import ascii
from astropy.table import Table
from astropy import units as u
from bs4 import BeautifulSoup

from ..query import BaseQuery
from ..utils import commons
from ..utils import prepend_docstr_noreturns
from ..utils import async_to_sync
from . import ATOMIC_LINE_LIST_URL, TIMEOUT

__all__ = ['AtomicLineList', 'AtomicLineListClass']

@async_to_sync
class AtomicLineListClass(BaseQuery):
    FORM_URL = ATOMIC_LINE_LIST_URL()
    TIMEOUT = TIMEOUT()

    def __init__(self):
        BaseQuery.__init__(self)
        self._default_form_values = None

    def query_object(self, wavelength_range=None, wavelength_type=None,
                     wavelength_accuracy=None, element_spectrum=None):
        """
        Parameters
        ----------
        wavelength_range : pair of `astropy.units.Unit` values

        wavelength_type : str
            Either 'Air' or 'Vacuum'

        wavelength_accuracy :

        element_spectrum :

        Returns
        -------
        result : `~astropy.table.Table`
            The result of the query as a `~astropy.table.Table` object.

        """
        response = self.query_object_async(
            wavelength_range, wavelength_type, wavelength_accuracy, element_spectrum)
        table = self._parse_result(response)
        return table

    @prepend_docstr_noreturns(query_object.__doc__)
    def query_object_async(self, wavelength_range=None, wavelength_type='',
                           wavelength_accuracy=None, element_spectrum=None):
        """
        Returns
        -------
        response : `requests.Response`
            The HTTP response returned from the service.
        """
        if self._default_form_values is None:
            response = commons.send_request(self.FORM_URL, {}, self.TIMEOUT, 'GET')
            bs = BeautifulSoup(response.text)
            form = bs.find('form')
            self._default_form_values = self._get_default_form_values(form)
        default_values = self._default_form_values
        wltype = (wavelength_type or default_values.get('air', '')).lower()
        if wltype in ('air', 'vacuum'):
            air = wltype.capitalize()
        else:
            raise ValueError('parameter wavelength_type must be either "air" or "vacuum".')
        wlrange = wavelength_range or []
        if len(wlrange) not in (0, 2):
            raise ValueError(
                'length of `wavelength_range` must be 2 or 0, but is: {}'.format(
                    len(wlrange)))
        # convert wavelengths in incoming wavelength range to Angstroms
        wlrange_in_angstroms = (wl.to(u.Angstrom, equivalencies=u.spectral()).value for wl in wlrange)
        input = {
            'wavl': '-'.join(map(str, wlrange_in_angstroms)),
            'wave': 'Angstrom',
            'air': air,
            'wacc': wavelength_accuracy,
            'elmion': element_spectrum,
            # TODO: add support for more parameters
        }
        response = self._submit_form(input)
        return response

    def _parse_result(self, response):
        data = StringIO(BeautifulSoup(response.text).find('pre').text.strip())
        # `header` is e.g. "u'-LAMBDA-VAC-ANG-|-SPECTRUM--|TT|--------TERM---------|---J-J---|----LEVEL ENERGY--CM-1----'"
        header = data.readline().strip().strip('|')
        # `colnames` is then "[u'LAMBDA VAC ANG', u'SPECTRUM', u'TT', u'TERM', u'J J', u'LEVEL ENERGY  CM 1']"
        colnames = [colname.strip('-').replace('-', ' ') for colname in header.split('|') if colname.strip()]
        indices = [i for i, c in enumerate(header) if c == '|']
        input = []
        for line in data:
            # example for `line`:
            # u'      1.010799    Zn XXX     E1          1-10          1/2-*            0.00 - 98933890.00\n'
            row = []
            for start, end in zip([0]+indices, indices+[None]):
                # `value` will hold all cell values in the line, so
                # `u'1.010799'`, `u'Zn XXX'` etc.
                value = line[start:end].strip()
                if value:
                    row.append(value)
            if row:
                input.append('\t'.join(row))
        if input:
            return ascii.read(input, data_start=0, delimiter='\t', names=colnames)
        else:
            # return an empty table if the query yielded no results
            return Table()

    def _submit_form(self, input=None):
        """Fill out the form of the SkyView site and submit it with the
        values given in `input` (a dictionary where the keys are the form
        element's names and the values are their respective values).

        """
        if input is None:
            input = {}
        response = commons.send_request(self.FORM_URL, {}, self.TIMEOUT, 'GET')
        bs = BeautifulSoup(response.text)
        form = bs.find('form')
        # cache the default values to save HTTP traffic
        if self._default_form_values is None:
            self._default_form_values = self._get_default_form_values(form)
        # only overwrite payload's values if the `input` value is not None
        # to avoid overwriting of the form's default values
        payload = self._default_form_values.copy()
        for k, v in six.iteritems(input):
            if v is not None:
                payload[k] = v
        url = urlparse.urljoin(self.FORM_URL, form.get('action'))
        response = commons.send_request(url, payload, self.TIMEOUT)
        return response

    def _get_default_form_values(self, form):
        """Return the already selected values of a given form (a BeautifulSoup
        form node) as a dict.

        """
        res = defaultdict(list)
        for elem in form.find_all(['input', 'select']):
            key = elem.get('name')
            value = None
            # ignore the submit and reset buttons
            if elem.get('type') in ['submit', 'reset']:
                continue
            # check boxes: enabled boxes have the value "on" if not specificed
            # otherwise. Found out by debugging, perhaps not documented.
            if elem.get('type') == 'checkbox' and elem.get('checked') in ["", "checked"]:
                value = elem.get('value', 'on')
            # radio buttons and simple input fields
            if elem.get('type') == 'radio' and\
                    elem.get('checked') in ["", "checked"] or\
                    elem.get('type') in [None, 'text']:
                value = elem.get('value')
            # dropdown menu, multi-section possible
            if elem.name == 'select':
                for option in elem.find_all('option'):
                    if option.get('selected') == '':
                        value = option.get('value', option.text.strip())
            if value and value not in [None, u'None', u'null']:
                res[key].append(value)
        # avoid values with size 1 lists
        d = dict(res)
        for k, v in six.iteritems(d):
            if len(v) == 1:
                d[k] = v[0]
        return d

AtomicLineList = AtomicLineListClass()
<|endoftext|>


<commit_before>"""
Django settings for root project.

Generated by 'django-admin startproject' using Django 1.10.1.

For more information on this file, see
https://docs.djangoproject.com/en/1.10/topics/settings/

For the full list of settings and their values, see
https://docs.djangoproject.com/en/1.10/ref/settings/
"""

import os

# Build paths inside the project like this: os.path.join(BASE_DIR, ...)
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# Quick-start development settings - unsuitable for production
# See https://docs.djangoproject.com/en/1.10/howto/deployment/checklist/

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = 'xjju44jdmn9e%ys03$w^1x-w*2$n3(@@qk9b=54b34c3&el1h9'

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = True

ALLOWED_HOSTS = []

# Application definition

INSTALLED_APPS = [
    'todolist.apps.TodolistConfig',
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
]

MIDDLEWARE = [
    'django.middleware.security.SecurityMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
]

ROOT_URLCONF = 'root.urls'

TEMPLATES = [
    {
        'BACKEND': 'django.template.backends.django.DjangoTemplates',
        'DIRS': [],
        'APP_DIRS': True,
        'OPTIONS': {
            'context_processors': [
                'django.template.context_processors.debug',
                'django.template.context_processors.request',
                'django.contrib.auth.context_processors.auth',
                'django.contrib.messages.context_processors.messages',
            ],
        },
    },
]

WSGI_APPLICATION = 'root.wsgi.application'

# Database
# https://docs.djangoproject.com/en/1.10/ref/settings/#databases

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
    }
}

# Password validation
# https://docs.djangoproject.com/en/1.10/ref/settings/#auth-password-validators

AUTH_PASSWORD_VALIDATORS = [
    {
        'NAME': 'django.contrib.auth.password_validation.UserAttributeSimilarityValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.MinimumLengthValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.CommonPasswordValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.NumericPasswordValidator',
    },
]

# Internationalization
# https://docs.djangoproject.com/en/1.10/topics/i18n/

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'UTC'

USE_I18N = True

USE_L10N = True

USE_TZ = True

# Static files (CSS, JavaScript, Images)
# https://docs.djangoproject.com/en/1.10/howto/static-files/

STATIC_URL = '/static/'

LOGIN_REDIRECT_URL = '/hello'
<commit_msg>settings login redirect<commit_after>"""
Django settings for root project.

Generated by 'django-admin startproject' using Django 1.10.1.

For more information on this file, see
https://docs.djangoproject.com/en/1.10/topics/settings/

For the full list of settings and their values, see
https://docs.djangoproject.com/en/1.10/ref/settings/
"""

import os

# Build paths inside the project like this: os.path.join(BASE_DIR, ...)
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# Quick-start development settings - unsuitable for production
# See https://docs.djangoproject.com/en/1.10/howto/deployment/checklist/

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = 'xjju44jdmn9e%ys03$w^1x-w*2$n3(@@qk9b=54b34c3&el1h9'

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = True

ALLOWED_HOSTS = []

# Application definition

INSTALLED_APPS = [
    'todolist.apps.TodolistConfig',
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
]

MIDDLEWARE = [
    'django.middleware.security.SecurityMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
]

ROOT_URLCONF = 'root.urls'

TEMPLATES = [
    {
        'BACKEND': 'django.template.backends.django.DjangoTemplates',
        'DIRS': [],
        'APP_DIRS': True,
        'OPTIONS': {
            'context_processors': [
                'django.template.context_processors.debug',
                'django.template.context_processors.request',
                'django.contrib.auth.context_processors.auth',
                'django.contrib.messages.context_processors.messages',
            ],
        },
    },
]

WSGI_APPLICATION = 'root.wsgi.application'

# Database
# https://docs.djangoproject.com/en/1.10/ref/settings/#databases

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
    }
}

# Password validation
# https://docs.djangoproject.com/en/1.10/ref/settings/#auth-password-validators

AUTH_PASSWORD_VALIDATORS = [
    {
        'NAME': 'django.contrib.auth.password_validation.UserAttributeSimilarityValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.MinimumLengthValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.CommonPasswordValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.NumericPasswordValidator',
    },
]

# Internationalization
# https://docs.djangoproject.com/en/1.10/topics/i18n/

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'UTC'

USE_I18N = True

USE_L10N = True

USE_TZ = True

# Static files (CSS, JavaScript, Images)
# https://docs.djangoproject.com/en/1.10/howto/static-files/

STATIC_URL = '/static/'

LOGIN_REDIRECT_URL = '/todo'
<|endoftext|>


<commit_before>__author__ = 'marion'

# add the shared settings file to namespace
import sys
from os.path import dirname, abspath
sys.path.insert(0, dirname(dirname(abspath(__file__))))
from settings import DATA_VIZ_METRICS
from settings import MUSE_PORTS
from settings import SPACEBREW_DATA_VIZ_NAME
from router.spacebrew_router import SpacebrewRouter

from spacebrew_utils import calculate_spacebrew_name
import time

start = time.time()

# todo: remove these lines line
subscriber_name = SPACEBREW_DATA_VIZ_NAME
#publisher_ip = '127.0.0.1'
publisher_ip = '104.236.243.233'
#subscriber_ip = '50.185.173.15'
#subscriber_ip = '108.74.162.96'
#subscriber_ip = '10.0.0.1'
subscriber_ip = '208.66.31.59'
spacebrew_server_ip = '208.66.31.59'

# spacebrew router
sp_router = SpacebrewRouter(server=spacebrew_server_ip)
for path in DATA_VIZ_METRICS:

    publisher_metric_name = calculate_spacebrew_name(path)
    for muse_port in MUSE_PORTS:

        subscriber_metric_name = '%s-%s' % (publisher_metric_name, muse_port)

        # route data FROM the muses
        publisher_name = 'muse-%s' % muse_port

        r = sp_router.link(publisher_metric_name, subscriber_metric_name, publisher_name, subscriber_name,
                            publisher_ip, subscriber_ip)
        print r

        time.sleep(0.1)

end = time.time()

print 'routing took %s s' % (end - start)

<commit_msg>updatd muse router<commit_after>__author__ = 'marion'

# add the shared settings file to namespace
import sys
from os.path import dirname, abspath
sys.path.insert(0, dirname(dirname(abspath(__file__))))
from settings import DATA_VIZ_METRICS
from settings import MUSE_PORTS
from settings import SPACEBREW_DATA_VIZ_NAME
from router.spacebrew_router import SpacebrewRouter

from spacebrew_utils import calculate_spacebrew_name
import time

start = time.time()

# todo: remove these lines line
subscriber_name = SPACEBREW_DATA_VIZ_NAME
#publisher_ip = '127.0.0.1'
publisher_ip = '107.170.205.177'
#subscriber_ip = '50.185.173.15'
#subscriber_ip = '108.74.162.96'
subscriber_ip = '10.0.0.245'
#subscriber_ip = '208.66.31.59'
spacebrew_server_ip = '107.170.205.177'

# spacebrew router
sp_router = SpacebrewRouter(server=spacebrew_server_ip)
for path in DATA_VIZ_METRICS:

    publisher_metric_name = calculate_spacebrew_name(path)
    for muse_port in MUSE_PORTS:

        subscriber_metric_name = '%s-muse-%s' % (publisher_metric_name, muse_port)

        # route data FROM the muses
        publisher_name = 'muse-%s' % muse_port

        r = sp_router.link(publisher_metric_name, subscriber_metric_name, publisher_name, subscriber_name,
                            publisher_ip, subscriber_ip)
        print r

        time.sleep(0.1)

end = time.time()

print 'routing took %s s' % (end - start)

<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-

from __future__ import absolute_import, division, unicode_literals, print_function

from distutils.version import LooseVersion
import sys
import inspect
import os
import subprocess

import six

from .. import environment
from ..console import log
from .. import util

WIN = (os.name == "nt")

class Virtualenv(environment.Environment):
    """
    Manage an environment using virtualenv.
    """
    tool_name = "virtualenv"

    def __init__(self, conf, python, requirements):
        """
        Parameters
        ----------
        conf : Config instance

        python : str
            Version of Python.  Must be of the form "MAJOR.MINOR".

        executable : str
            Path to Python executable.

        requirements : dict
            Dictionary mapping a PyPI package name to a version
            identifier string.
        """
        executable = Virtualenv._find_python(python)
        if executable is None:
            raise environment.EnvironmentUnavailable("No executable found for python {0}".format(python))

        self._executable = executable
        self._python = python
        self._requirements = requirements
        super(Virtualenv, self).__init__(conf, python, requirements)

        try:
            import virtualenv
        except ImportError:
            raise util.UserError(
                "virtualenv must be installed to run asv with the given "
                "Python executable")

        # Can't use `virtualenv.__file__` here, because that will refer to a
        # .pyc file which can't be used on another version of Python
        self._virtualenv_path = os.path.abspath(
            inspect.getsourcefile(virtualenv))

    @staticmethod
    def _find_python(python):
        """Find Python executable for the given Python version"""
        # Find Python executable on path
        try:
            return util.which('python{0}'.format(python))
        except IOError:
            pass

        # Maybe the current one is correct
        if '{0[0]}.{0[1]}'.format(sys.version_info) == python:
            return sys.executable

        return None

    @classmethod
    def matches(self, python):
        try:
            import virtualenv
        except ImportError:
            return False
        else:
            if LooseVersion(virtualenv.__version__) == LooseVersion('1.11.0'):
                log.warn(
                    "asv is not compatible with virtualenv 1.11 due to a bug in "
                    "setuptools.")
            if LooseVersion(virtualenv.__version__) < LooseVersion('1.10'):
                log.warn(
                    "If using virtualenv, it much be at least version 1.10")

        executable = Virtualenv._find_python(python)
        return executable is not None

    def _setup(self):
        """
        Setup the environment on disk using virtualenv.
        Then, all of the requirements are installed into
        it using `pip install`.
        """
        log.info("Creating virtualenv for {0}".format(self.name))
        util.check_call([
            self._executable,
            self._virtualenv_path,
            '--no-site-packages',
            self._path])

        log.info("Installing requirements for {0}".format(self.name))
        self._install_requirements()

    def _install_requirements(self):
        if sys.version_info[:2] == (3, 2):
            self.run_executable('pip', ['install', '-v', 'wheel<0.29.0', 'pip<8'])
        else:
            self.run_executable('pip', ['install', '-v', 'wheel', 'pip>=8'])

        if self._requirements:
            args = ['install', '-v', '--upgrade']
            for key, val in six.iteritems(self._requirements):
                pkg = key
                if key.startswith('pip+'):
                    pkg = key[4:]

                if val:
                    args.append("{0}=={1}".format(pkg, val))
                else:
                    args.append(pkg)
            self.run_executable('pip', args)

    def install(self, package):
        log.info("Installing into {0}".format(self.name))
        self.run_executable('pip', ['install', package])

    def uninstall(self, package):
        log.info("Uninstalling from {0}".format(self.name))
        self.run_executable('pip', ['uninstall', '-y', package],
                             valid_return_codes=None)

    def run(self, args, **kwargs):
        log.debug("Running '{0}' in {1}".format(' '.join(args), self.name))
        return self.run_executable('python', args, **kwargs)
<commit_msg>pass ``install_timeout`` parameter to run_executable in ``Virtualenv`` plugin ``install`` and ``_install_requirements`` calls (GH391)<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-

from __future__ import absolute_import, division, unicode_literals, print_function

from distutils.version import LooseVersion
import sys
import inspect
import os
import subprocess

import six

from .. import environment
from ..console import log
from .. import util

WIN = (os.name == "nt")

class Virtualenv(environment.Environment):
    """
    Manage an environment using virtualenv.
    """
    tool_name = "virtualenv"

    def __init__(self, conf, python, requirements):
        """
        Parameters
        ----------
        conf : Config instance

        python : str
            Version of Python.  Must be of the form "MAJOR.MINOR".

        executable : str
            Path to Python executable.

        requirements : dict
            Dictionary mapping a PyPI package name to a version
            identifier string.
        """
        executable = Virtualenv._find_python(python)
        if executable is None:
            raise environment.EnvironmentUnavailable(
                "No executable found for python {0}".format(python))

        self._executable = executable
        self._python = python
        self._requirements = requirements
        super(Virtualenv, self).__init__(conf, python, requirements)

        try:
            import virtualenv
        except ImportError:
            raise util.UserError(
                "virtualenv must be installed to run asv with the given "
                "Python executable")

        # Can't use `virtualenv.__file__` here, because that will refer to a
        # .pyc file which can't be used on another version of Python
        self._virtualenv_path = os.path.abspath(
            inspect.getsourcefile(virtualenv))

    @staticmethod
    def _find_python(python):
        """Find Python executable for the given Python version"""
        # Find Python executable on path
        try:
            return util.which('python{0}'.format(python))
        except IOError:
            pass

        # Maybe the current one is correct
        if '{0[0]}.{0[1]}'.format(sys.version_info) == python:
            return sys.executable

        return None

    @classmethod
    def matches(self, python):
        try:
            import virtualenv
        except ImportError:
            return False
        else:
            if LooseVersion(virtualenv.__version__) == LooseVersion('1.11.0'):
                log.warn(
                    "asv is not compatible with virtualenv 1.11 due to a bug in "
                    "setuptools.")
            if LooseVersion(virtualenv.__version__) < LooseVersion('1.10'):
                log.warn(
                    "If using virtualenv, it much be at least version 1.10")

        executable = Virtualenv._find_python(python)
        return executable is not None

    def _setup(self):
        """
        Setup the environment on disk using virtualenv.
        Then, all of the requirements are installed into
        it using `pip install`.
        """
        log.info("Creating virtualenv for {0}".format(self.name))
        util.check_call([
            self._executable,
            self._virtualenv_path,
            '--no-site-packages',
            self._path])

        log.info("Installing requirements for {0}".format(self.name))
        self._install_requirements()

    def _install_requirements(self):
        if sys.version_info[:2] == (3, 2):
            self.run_executable(
                'pip', ['install', '-v', 'wheel<0.29.0', 'pip<8'])
        else:
            self.run_executable('pip', ['install', '-v', 'wheel', 'pip>=8'])

        if self._requirements:
            args = ['install', '-v', '--upgrade']
            for key, val in six.iteritems(self._requirements):
                pkg = key
                if key.startswith('pip+'):
                    pkg = key[4:]

                if val:
                    args.append("{0}=={1}".format(pkg, val))
                else:
                    args.append(pkg)
            self.run_executable('pip', args, timeout=self._install_timeout)

    def install(self, package):
        log.info("Installing into {0}".format(self.name))
        self.run_executable('pip', ['install', package],
                            timeout=self._install_timeout)

    def uninstall(self, package):
        log.info("Uninstalling from {0}".format(self.name))
        self.run_executable('pip', ['uninstall', '-y', package],
                            valid_return_codes=None)

    def run(self, args, **kwargs):
        log.debug("Running '{0}' in {1}".format(' '.join(args), self.name))
        return self.run_executable('python', args, **kwargs)
<|endoftext|>


<commit_before>"""
Incorporated from django-storages, copyright all of those listed in:
http://code.welldev.org/django-storages/src/tip/AUTHORS
"""
import urllib
import os
import mimetypes
import re

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from django.conf import settings
from django.core.files.base import File
from django.core.files.storage import Storage
from django.core.exceptions import ImproperlyConfigured

try:
    from boto.s3.connection import S3Connection
    from boto.exception import S3ResponseError
    from boto.s3.key import Key
except ImportError:
    raise ImproperlyConfigured(
        "Could not load boto's S3 bindings. Please install boto."
    )

AWS_REGIONS = [
    'eu-west-1',
    'us-east-1',
    'us-west-1',
    'us-west-2',
    'sa-east-1',
    'ap-northeast-1',
    'ap-southeast-1',
]

REGION_RE = re.compile(r's3-(.+).amazonaws.com')

ACCESS_KEY_NAME = getattr(settings, 'AWS_ACCESS_KEY_ID', None)
SECRET_KEY_NAME = getattr(settings, 'AWS_SECRET_ACCESS_KEY', None)
HEADERS = getattr(settings, 'AWS_HEADERS', {})
STORAGE_BUCKET_NAME = getattr(settings, 'AWS_STORAGE_BUCKET_NAME', None)
STORAGE_BUCKET_CNAME = getattr(settings, 'AWS_STORAGE_BUCKET_CNAME', None)
AWS_REGION = getattr(settings, 'AWS_REGION', 'us-east-1')
AUTO_CREATE_BUCKET = getattr(settings, 'AWS_AUTO_CREATE_BUCKET', True)
DEFAULT_ACL = getattr(settings, 'AWS_DEFAULT_ACL', 'public-read')
QUERYSTRING_AUTH = getattr(settings, 'AWS_QUERYSTRING_AUTH', True)
QUERYSTRING_EXPIRE = getattr(settings, 'AWS_QUERYSTRING_EXPIRE', 3600)
IS_GZIPPED = getattr(settings, 'AWS_IS_GZIPPED', False)
GZIP_CONTENT_TYPES = getattr(settings, 'GZIP_CONTENT_TYPES', (
    'text/css',
    'application/javascript',
    'application/x-javascript'
))

if IS_GZIPPED:
    from gzip import GzipFile

class S3BotoStorage(Storage):
    """Amazon Simple Storage Service using Boto"""

    def __init__(self, bucket=STORAGE_BUCKET_NAME,
                       bucket_cname=STORAGE_BUCKET_CNAME,
                       region=AWS_REGION, access_key=None,
                       secret_key=None, acl=DEFAULT_ACL,
                       headers=HEADERS, gzip=IS_GZIPPED,
                       gzip_content_types=GZIP_CONTENT_TYPES,
                       querystring_auth=QUERYSTRING_AUTH,
                       force_no_ssl=False):
        self.bucket_name = bucket
        self.bucket_cname = bucket_cname
        self.host = self._get_host(region)
        self.acl = acl
        self.headers = headers
        self.gzip = gzip
        self.gzip_content_types = gzip_content_types
        self.querystring_auth = querystring_auth
        self.force_no_ssl = force_no_ssl
        # This is called as chunks are uploaded to S3. Useful for getting
        # around limitations in eventlet for things like gunicorn.
        self.s3_callback_during_upload = None

        if not access_key and not secret_key:
            access_key, secret_key = self._get_access_keys()

        self.connection = S3Connection(
            access_key, secret_key, host=self.host,
        )

    @property
    def bucket(self):
        if not hasattr(self, '_bucket'):
            self._bucket = self._get_or_create_bucket(self.bucket_name)
        return self._bucket

    def _get_access_keys(self):
        access_key = ACCESS_KEY_NAME
        secret_key = SECRET_KEY_NAME
        if (access_key or secret_key) and (not access_key or not secret_key):
            access_key = os.environ.get(ACCESS_KEY_NAME)
            secret_key = os.environ.get(SECRET_KEY_NAME)

        if access_key and secret_key:
            # Both were provided, so use them
            return access_key, secret_key

        return None, None

    def _get_host(self, region):
        """
        Returns correctly formatted host. Accepted formats:

            * simple region name, eg 'us-west-1' (see list in AWS_REGIONS)
            * full host name, eg 's3-us-west-1.amazonaws.com'.
        """
        if 'us-east-1' in region:
            return 's3.amazonaws.com'
        elif region in AWS_REGIONS:
            return 's3-%s.amazonaws.com' % region
        elif region and not REGION_RE.findall(region):
            raise ImproperlyConfigured('AWS_REGION improperly configured!')
        # can be full host or empty string, default region
        return  region

    def _get_or_create_bucket(self, name):
        """Retrieves a bucket if it exists, otherwise creates it."""
        try:
            return self.connection.get_bucket(name)
        except S3ResponseError, e:
            if AUTO_CREATE_BUCKET:
                return self.connection.create_bucket(name)
            raise ImproperlyConfigured, ("Bucket specified by "
            "AWS_STORAGE_BUCKET_NAME does not exist. Buckets can be "
            "automatically created by setting AWS_AUTO_CREATE_BUCKET=True")

    def _clean_name(self, name):
        # Useful for windows' paths
        return os.path.normpath(name).replace('\\', '/')

    def _compress_content(self, content):
        """Gzip a given string."""
        zbuf = StringIO()
        zfile = GzipFile(mode='wb', compresslevel=6, fileobj=zbuf)
        zfile.write(content.read())
        zfile.close()
        content.file = zbuf
        return content

    def _open(self, name, mode='rb'):
        name = self._clean_name(name)
        return S3BotoStorageFile(name, mode, self)

    def _save(self, name, content):
        name = self._clean_name(name)

        if callable(self.headers):
            headers = self.headers(name, content)
        else:
            headers = self.headers

        if hasattr(content.file, 'content_type'):
            content_type = content.file.content_type
        else:
            content_type = mimetypes.guess_type(name)[0] or "application/x-octet-stream"

        if self.gzip and content_type in self.gzip_content_types:
            content = self._compress_content(content)
            headers.update({'Content-Encoding': 'gzip'})

        headers.update({
            'Content-Type': content_type,
            'Content-Length' : len(content),
        })

        content.name = name
        k = self.bucket.get_key(name)
        if not k:
            k = self.bucket.new_key(name)
        # The callback seen here is particularly important for async WSGI
        # servers. This allows us to call back to eventlet or whatever
        # async support library we're using periodically to prevent timeouts.
        k.set_contents_from_file(content, headers=headers, policy=self.acl,
                                 cb=self.s3_callback_during_upload,
                                 num_cb=-1)
        return name

    def delete(self, name):
        name = self._clean_name(name)
        self.bucket.delete_key(name)

    def exists(self, name):
        name = self._clean_name(name)
        k = Key(self.bucket, name)
        return k.exists()

    def listdir(self, name):
        name = self._clean_name(name)
        return [l.name for l in self.bucket.list() if not len(name) or l.name[:len(name)] == name]

    def size(self, name):
        name = self._clean_name(name)
        return self.bucket.get_key(name).size

    def url(self, name):
        name = self._clean_name(name)
        if self.bucket.get_key(name) is None:
            return ''
        return self.bucket.get_key(name).generate_url(QUERYSTRING_EXPIRE,
                                                      method='GET',
                                                      query_auth=self.querystring_auth,
                                                      force_http=self.force_no_ssl)

    def url_as_attachment(self, name, filename=None):
        name = self._clean_name(name)

        if filename:
            disposition = 'attachment; filename="%s"' % filename
        else:
            disposition = 'attachment;'

        response_headers = {
            'response-content-disposition': disposition,
        }

        return self.connection.generate_url(QUERYSTRING_EXPIRE, 'GET',
                                            bucket=self.bucket.name, key=name,
                                            query_auth=True,
                                            force_http=self.force_no_ssl,
                                            response_headers=response_headers)

    def get_available_name(self, name):
        """ Overwrite existing file with the same name. """
        name = self._clean_name(name)
        return name

class S3BotoStorage_AllPublic(S3BotoStorage):
    """
    Same as S3BotoStorage, but defaults to uploading everything with a
    public acl. This has two primary beenfits:

    1) Non-encrypted requests just make a lot better sense for certain things
       like profile images. Much faster, no need to generate S3 auth keys.
    2) Since we don't have to hit S3 for auth keys, this backend is much
       faster than S3BotoStorage, as it makes no attempt to validate whether
       keys exist.

    WARNING: This backend makes absolutely no attempt to verify whether the
    given key exists on self.url(). This is much faster, but be aware.
    """
    def __init__(self, *args, **kwargs):
        super(S3BotoStorage_AllPublic, self).__init__(acl='public-read',
                                                      querystring_auth=False,
                                                      force_no_ssl=True,
                                                      *args, **kwargs)

    def url(self, name):
        """
        Since we assume all public storage with no authorization keys, we can
        just simply dump out a URL rather than having to query S3 for new keys.
        """
        name = urllib.quote_plus(self._clean_name(name), safe='/')

        if self.bucket_cname:
            return "http://%s/%s" % (self.bucket_cname, name)
        elif self.host:
            return "http://%s/%s/%s" % (self.host, self.bucket_name, name)
        # No host ? Then it's the default region
        return "http://s3.amazonaws.com/%s/%s" % (self.bucket_name, name)

class S3BotoStorageFile(File):
    def __init__(self, name, mode, storage):
        self._storage = storage
        self.name = name
        self._mode = mode
        self.key = storage.bucket.get_key(name)
        self._is_dirty = False
        self.file = StringIO()

    @property
    def size(self):
        if not self.key:
            raise IOError('No such S3 key: %s' % self.name)
        return self.key.size

    def read(self, *args, **kwargs):
        self.file = StringIO()
        self._is_dirty = False
        if not self.key:
            raise IOError('No such S3 key: %s' % self.name)
        self.key.get_contents_to_file(self.file)
        return self.file.getvalue()

    def write(self, content):
        if 'w' not in self._mode:
            raise AttributeError("File was opened for read-only access.")
        self.file = StringIO(content)
        self._is_dirty = True

    def close(self):
        if self._is_dirty:
            if not self.key:
                self.key = self._storage.bucket.new_key(key_name=self.name)
            self.key.set_contents_from_string(self.file.getvalue(), headers=self._storage.headers, policy=self.storage.acl)
        self.key.close()
<commit_msg>Use protocol-less urls for loading images from S3.<commit_after>"""
Incorporated from django-storages, copyright all of those listed in:
http://code.welldev.org/django-storages/src/tip/AUTHORS
"""
import urllib
import os
import mimetypes
import re

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from django.conf import settings
from django.core.files.base import File
from django.core.files.storage import Storage
from django.core.exceptions import ImproperlyConfigured

try:
    from boto.s3.connection import S3Connection
    from boto.exception import S3ResponseError
    from boto.s3.key import Key
except ImportError:
    raise ImproperlyConfigured(
        "Could not load boto's S3 bindings. Please install boto."
    )

AWS_REGIONS = [
    'eu-west-1',
    'us-east-1',
    'us-west-1',
    'us-west-2',
    'sa-east-1',
    'ap-northeast-1',
    'ap-southeast-1',
]

REGION_RE = re.compile(r's3-(.+).amazonaws.com')

ACCESS_KEY_NAME = getattr(settings, 'AWS_ACCESS_KEY_ID', None)
SECRET_KEY_NAME = getattr(settings, 'AWS_SECRET_ACCESS_KEY', None)
HEADERS = getattr(settings, 'AWS_HEADERS', {})
STORAGE_BUCKET_NAME = getattr(settings, 'AWS_STORAGE_BUCKET_NAME', None)
STORAGE_BUCKET_CNAME = getattr(settings, 'AWS_STORAGE_BUCKET_CNAME', None)
AWS_REGION = getattr(settings, 'AWS_REGION', 'us-east-1')
AUTO_CREATE_BUCKET = getattr(settings, 'AWS_AUTO_CREATE_BUCKET', True)
DEFAULT_ACL = getattr(settings, 'AWS_DEFAULT_ACL', 'public-read')
QUERYSTRING_AUTH = getattr(settings, 'AWS_QUERYSTRING_AUTH', True)
QUERYSTRING_EXPIRE = getattr(settings, 'AWS_QUERYSTRING_EXPIRE', 3600)
IS_GZIPPED = getattr(settings, 'AWS_IS_GZIPPED', False)
GZIP_CONTENT_TYPES = getattr(settings, 'GZIP_CONTENT_TYPES', (
    'text/css',
    'application/javascript',
    'application/x-javascript'
))

if IS_GZIPPED:
    from gzip import GzipFile

class S3BotoStorage(Storage):
    """Amazon Simple Storage Service using Boto"""

    def __init__(self, bucket=STORAGE_BUCKET_NAME,
                       bucket_cname=STORAGE_BUCKET_CNAME,
                       region=AWS_REGION, access_key=None,
                       secret_key=None, acl=DEFAULT_ACL,
                       headers=HEADERS, gzip=IS_GZIPPED,
                       gzip_content_types=GZIP_CONTENT_TYPES,
                       querystring_auth=QUERYSTRING_AUTH,
                       force_no_ssl=False):
        self.bucket_name = bucket
        self.bucket_cname = bucket_cname
        self.host = self._get_host(region)
        self.acl = acl
        self.headers = headers
        self.gzip = gzip
        self.gzip_content_types = gzip_content_types
        self.querystring_auth = querystring_auth
        self.force_no_ssl = force_no_ssl
        # This is called as chunks are uploaded to S3. Useful for getting
        # around limitations in eventlet for things like gunicorn.
        self.s3_callback_during_upload = None

        if not access_key and not secret_key:
            access_key, secret_key = self._get_access_keys()

        self.connection = S3Connection(
            access_key, secret_key, host=self.host,
        )

    @property
    def bucket(self):
        if not hasattr(self, '_bucket'):
            self._bucket = self._get_or_create_bucket(self.bucket_name)
        return self._bucket

    def _get_access_keys(self):
        access_key = ACCESS_KEY_NAME
        secret_key = SECRET_KEY_NAME
        if (access_key or secret_key) and (not access_key or not secret_key):
            access_key = os.environ.get(ACCESS_KEY_NAME)
            secret_key = os.environ.get(SECRET_KEY_NAME)

        if access_key and secret_key:
            # Both were provided, so use them
            return access_key, secret_key

        return None, None

    def _get_host(self, region):
        """
        Returns correctly formatted host. Accepted formats:

            * simple region name, eg 'us-west-1' (see list in AWS_REGIONS)
            * full host name, eg 's3-us-west-1.amazonaws.com'.
        """
        if 'us-east-1' in region:
            return 's3.amazonaws.com'
        elif region in AWS_REGIONS:
            return 's3-%s.amazonaws.com' % region
        elif region and not REGION_RE.findall(region):
            raise ImproperlyConfigured('AWS_REGION improperly configured!')
        # can be full host or empty string, default region
        return  region

    def _get_or_create_bucket(self, name):
        """Retrieves a bucket if it exists, otherwise creates it."""
        try:
            return self.connection.get_bucket(name)
        except S3ResponseError, e:
            if AUTO_CREATE_BUCKET:
                return self.connection.create_bucket(name)
            raise ImproperlyConfigured, ("Bucket specified by "
            "AWS_STORAGE_BUCKET_NAME does not exist. Buckets can be "
            "automatically created by setting AWS_AUTO_CREATE_BUCKET=True")

    def _clean_name(self, name):
        # Useful for windows' paths
        return os.path.normpath(name).replace('\\', '/')

    def _compress_content(self, content):
        """Gzip a given string."""
        zbuf = StringIO()
        zfile = GzipFile(mode='wb', compresslevel=6, fileobj=zbuf)
        zfile.write(content.read())
        zfile.close()
        content.file = zbuf
        return content

    def _open(self, name, mode='rb'):
        name = self._clean_name(name)
        return S3BotoStorageFile(name, mode, self)

    def _save(self, name, content):
        name = self._clean_name(name)

        if callable(self.headers):
            headers = self.headers(name, content)
        else:
            headers = self.headers

        if hasattr(content.file, 'content_type'):
            content_type = content.file.content_type
        else:
            content_type = mimetypes.guess_type(name)[0] or "application/x-octet-stream"

        if self.gzip and content_type in self.gzip_content_types:
            content = self._compress_content(content)
            headers.update({'Content-Encoding': 'gzip'})

        headers.update({
            'Content-Type': content_type,
            'Content-Length' : len(content),
        })

        content.name = name
        k = self.bucket.get_key(name)
        if not k:
            k = self.bucket.new_key(name)
        # The callback seen here is particularly important for async WSGI
        # servers. This allows us to call back to eventlet or whatever
        # async support library we're using periodically to prevent timeouts.
        k.set_contents_from_file(content, headers=headers, policy=self.acl,
                                 cb=self.s3_callback_during_upload,
                                 num_cb=-1)
        return name

    def delete(self, name):
        name = self._clean_name(name)
        self.bucket.delete_key(name)

    def exists(self, name):
        name = self._clean_name(name)
        k = Key(self.bucket, name)
        return k.exists()

    def listdir(self, name):
        name = self._clean_name(name)
        return [l.name for l in self.bucket.list() if not len(name) or l.name[:len(name)] == name]

    def size(self, name):
        name = self._clean_name(name)
        return self.bucket.get_key(name).size

    def url(self, name):
        name = self._clean_name(name)
        if self.bucket.get_key(name) is None:
            return ''
        return self.bucket.get_key(name).generate_url(QUERYSTRING_EXPIRE,
                                                      method='GET',
                                                      query_auth=self.querystring_auth,
                                                      force_http=self.force_no_ssl)

    def url_as_attachment(self, name, filename=None):
        name = self._clean_name(name)

        if filename:
            disposition = 'attachment; filename="%s"' % filename
        else:
            disposition = 'attachment;'

        response_headers = {
            'response-content-disposition': disposition,
        }

        return self.connection.generate_url(QUERYSTRING_EXPIRE, 'GET',
                                            bucket=self.bucket.name, key=name,
                                            query_auth=True,
                                            force_http=self.force_no_ssl,
                                            response_headers=response_headers)

    def get_available_name(self, name):
        """ Overwrite existing file with the same name. """
        name = self._clean_name(name)
        return name

class S3BotoStorage_AllPublic(S3BotoStorage):
    """
    Same as S3BotoStorage, but defaults to uploading everything with a
    public acl. This has two primary beenfits:

    1) Non-encrypted requests just make a lot better sense for certain things
       like profile images. Much faster, no need to generate S3 auth keys.
    2) Since we don't have to hit S3 for auth keys, this backend is much
       faster than S3BotoStorage, as it makes no attempt to validate whether
       keys exist.

    WARNING: This backend makes absolutely no attempt to verify whether the
    given key exists on self.url(). This is much faster, but be aware.
    """
    def __init__(self, *args, **kwargs):
        super(S3BotoStorage_AllPublic, self).__init__(acl='public-read',
                                                      querystring_auth=False,
                                                      force_no_ssl=True,
                                                      *args, **kwargs)

    def url(self, name):
        """
        Since we assume all public storage with no authorization keys, we can
        just simply dump out a URL rather than having to query S3 for new keys.
        """
        name = urllib.quote_plus(self._clean_name(name), safe='/')

        if self.bucket_cname:
            return "//%s/%s" % (self.bucket_cname, name)
        elif self.host:
            return "//%s/%s/%s" % (self.host, self.bucket_name, name)
        # No host ? Then it's the default region
        return "//s3.amazonaws.com/%s/%s" % (self.bucket_name, name)

class S3BotoStorageFile(File):
    def __init__(self, name, mode, storage):
        self._storage = storage
        self.name = name
        self._mode = mode
        self.key = storage.bucket.get_key(name)
        self._is_dirty = False
        self.file = StringIO()

    @property
    def size(self):
        if not self.key:
            raise IOError('No such S3 key: %s' % self.name)
        return self.key.size

    def read(self, *args, **kwargs):
        self.file = StringIO()
        self._is_dirty = False
        if not self.key:
            raise IOError('No such S3 key: %s' % self.name)
        self.key.get_contents_to_file(self.file)
        return self.file.getvalue()

    def write(self, content):
        if 'w' not in self._mode:
            raise AttributeError("File was opened for read-only access.")
        self.file = StringIO(content)
        self._is_dirty = True

    def close(self):
        if self._is_dirty:
            if not self.key:
                self.key = self._storage.bucket.new_key(key_name=self.name)
            self.key.set_contents_from_string(self.file.getvalue(), headers=self._storage.headers, policy=self.storage.acl)
        self.key.close()
<|endoftext|>


<commit_before># This script can be used to process imSim simulations that were produced with atmospheric (but not
# optics) PSF, to study the statistical properties of the PSFs.

import numpy as np
import matplotlib.pyplot as plt
import os, sys

# First, get access to the power spectrum estimator.  It's not yet in GalSim or installed on my
# system, so have to add its dir to path.  This should work on other people's machines if they are
# sitting in the directory containing this file (process_imsim.py) in the great3-private repo.
pse_path = os.path.abspath('../../metrics')
sys.path.append(pse_path)
import pse

# Helper function: read in a single file
def process_file(filename, verbose):
    # Read in the data
    data = np.loadtxt(filename).transpose()

    # Separate into the fields we want
    x = 0.2*data[1] # arcsec (convert from 0.2" LSST pixels)
    y = 0.2*data[2] # arcsec (convert from 0.2" LSST pixels)
    e1 = data[3]
    e2 = data[4]
    fwhm = 3600.*data[6] # to get arcsec

    # Process FWHM to get fluctuations around mean.
    dfwhm = fwhm/np.mean(fwhm)-1.0
    # Remove the large-scale correlations due to atmospheric dispersion, which are roughly like a
    # constant non-zero mean ellipticity.  For the PS, that will look like something peaked at k=0,
    # which we do not represent with our lensing engine.  Hence we remove it here, and when
    # simulating atmospheric PSFs, we can simply draw a random value of e (from within a range given
    # by the sims) to add after the fact.
    de1 = e1-np.mean(e1)
    de2 = e2-np.mean(e2)

    # Get some basic statistics:
    if verbose:
        print "Some statistics of file ",filename
        print "Number of stars in input grid (expected 8820):",data.shape[1]
        print "Min, max, mean x (arcsec):",np.min(x), np.max(x), np.mean(x)
        print "Min, max, mean y (arcsec):",np.min(y), np.max(y), np.mean(y)
        print "Mean FWHM and std dev of fluctuations around mean:",np.mean(fwhm), np.std(dfwhm)
        print "Mean e1, std dev of e1:",np.mean(e1), np.std(e1)
        print "Mean e2, std dev of e2:",np.mean(e2), np.std(e2)
        print "Showing grid:"
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(x, y, marker='+')
        ax.set_xlabel('x (arcsec)')
        ax.set_ylabel('y (arcsec)')
        plt.xlim((0,1000))
        plt.ylim((0,1000))
        ax.set_title('Locations of stars')
        plt.show()

    # Now define expected grid locations
    x_grid_min = 200.
    y_grid_min = 48.
    dgrid = 360. # 10 degrees / 100
    ngrid = 21 # because this grid is only covering 2 degrees
    x_grid_index = (x-x_grid_min)/dgrid
    round_x_grid_index = np.round(x_grid_index)
    y_grid_index = (y-y_grid_min)/dgrid
    round_y_grid_index = np.round(y_grid_index)
    x_use = round_x_grid_index[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.int)
    y_use = round_y_grid_index[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.int)
    de1_use = de1[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.int)
    de2_use = de2[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.int)
    dfwhm_use = dfwhm[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.int)
    if verbose:
        print "After grid selection:"
        print "Number of grid points:",len(x_use)
        print "Min,max x,y:",np.min(x_use), np.max(x_use), np.min(y_use), np.max(y_use)
        #fig = plt.figure()
        #ax = fig.add_subplot(111)
        #ax.plot(x_use, y_use, marker='+')
        #ax.set_xlabel('x index')
        #ax.set_ylabel('y index')
        #plt.show()

    # Make grids of the interesting parameters.
    de1_grid = np.zeros((ngrid, ngrid))
    de2_grid = np.zeros((ngrid, ngrid))
    dfwhm_grid = np.zeros((ngrid, ngrid))
    de1_grid[x_use, y_use] = de1_use
    de2_grid[x_use, y_use] = de2_use
    dfwhm_grid[x_use, y_use] = dfwhm_use

# Main function: for now, just specify a single file directly.
infile = '/Users/rmandelb/great3/chihway-sims/catalogs/1min/focalplane_47.txt'
verbose = 1
process_file(infile, verbose)

# Later, will need to write code that crawls through all directories and processes everything.
<commit_msg>more work on processing imsim inputs, with weird results (#10)<commit_after># This script can be used to process imSim simulations that were produced with atmospheric (but not
# optics) PSF, to study the statistical properties of the PSFs.

import numpy as np
import matplotlib.pyplot as plt
import os, sys
import scipy.optimize as opti

# First, get access to the power spectrum estimator.  It's not yet in GalSim or installed on my
# system, so have to add its dir to path.  This should work on other people's machines if they are
# sitting in the directory containing this file (process_imsim.py) in the great3-private repo.
pse_path = os.path.abspath('../../metrics')
sys.path.append(pse_path)
import pse

# Helper function: read in a single file (`filename').  `verbose' determines how much information
# gets outputted to stdout, 'do_plot' determines whether or not to display plots, and `n_ell'
# determines how many logarithmic ell bins to use for the PS estimation.
def process_file(filename, n_ell, verbose=True, do_plot=False):
    # Read in the data
    data = np.loadtxt(filename).transpose()

    # Separate into the fields we want
    x = 0.2*data[1] # arcsec (convert from 0.2" LSST pixels)
    y = 0.2*data[2] # arcsec (convert from 0.2" LSST pixels)
    e1 = data[3]
    e2 = data[4]
    fwhm = 3600.*data[6] # to get arcsec

    # Process FWHM to get fluctuations around mean.
    dfwhm = fwhm/np.mean(fwhm)-1.0
    # Remove the large-scale correlations due to atmospheric dispersion, which are roughly like a
    # constant non-zero mean ellipticity.  For the PS, that will look like something peaked at k=0,
    # which we do not represent with our lensing engine.  Hence we remove it here, and when
    # simulating atmospheric PSFs, we can simply draw a random value of e (from within a range given
    # by the sims) to add after the fact.
    de1 = e1-np.mean(e1)
    de2 = e2-np.mean(e2)

    # Get some basic statistics:
    if verbose:
        print "Some statistics of file ",filename
        print "Number of stars in input grid (expected 8820):",data.shape[1]
        print "Min, max, mean x (arcsec):",np.min(x), np.max(x), np.mean(x)
        print "Min, max, mean y (arcsec):",np.min(y), np.max(y), np.mean(y)
        print "Mean FWHM and std dev of fluctuations around mean:",np.mean(fwhm), np.std(dfwhm)
        print "Mean e1, std dev of e1:",np.mean(e1), np.std(e1)
        print "Mean e2, std dev of e2:",np.mean(e2), np.std(e2)
        print "Showing grid:"
    if do_plot:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(x, y, marker='+', linestyle='None')
        ax.set_xlabel('x (arcsec)')
        ax.set_ylabel('y (arcsec)')
        plt.xlim((0,1000))
        plt.ylim((0,1000))
        ax.set_title('Locations of stars')
        plt.show()

    # Now define expected grid locations
    x_grid_min = 200.
    y_grid_min = 48.
    dgrid = 360. # in arcsec: 10 degrees / 100
    ngrid = 21 # because this grid is only covering 2 degrees
    x_grid_index = (x-x_grid_min)/dgrid
    round_x_grid_index = np.round(x_grid_index)
    y_grid_index = (y-y_grid_min)/dgrid
    round_y_grid_index = np.round(y_grid_index)
    x_use = round_x_grid_index[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.int)
    y_use = round_y_grid_index[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.int)
    de1_use = de1[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.float32)
    de2_use = de2[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.float32)
    dfwhm_use = dfwhm[(np.abs(round_x_grid_index-x_grid_index)<0.05) & (np.abs(round_y_grid_index-y_grid_index)<0.05)].astype(np.float32)
    if verbose:
        print "After grid selection:"
        print "Number of grid points:",len(x_use)
        print "Min,max x,y:",np.min(x_use), np.max(x_use), np.min(y_use), np.max(y_use)
    if do_plot:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(x_use, y_use, marker='+', linestyle='None')
        ax.set_xlabel('x index')
        ax.set_ylabel('y index')
        plt.show()

    # Make grids of the interesting parameters.
    de1_grid = np.zeros((ngrid, ngrid))
    de2_grid = np.zeros((ngrid, ngrid))
    dfwhm_grid = np.zeros((ngrid, ngrid))
    de1_grid[x_use, y_use] = de1_use
    de2_grid[x_use, y_use] = de2_use
    dfwhm_grid[x_use, y_use] = dfwhm_use

    # PS estimation based on ellipticities
    theta = ngrid*dgrid/3600. # total size of grid in degrees, required for PSE
    atmos_e_pse = pse.PowerSpectrumEstimator(ngrid, theta, n_ell)
    ell, cee, cbb, ceb = atmos_e_pse.estimate(de1_grid, de2_grid)

    # Check best-fit power-law on large k (small scales).
    subell = ell[(ell>=500) & (ell<= 2500)]
    sublgell = np.log10(ell[(ell>=500) & (ell<= 2500)])-np.log10(1000.)
    sublgcee = np.log10(subell**2 * cee[(ell>=500) & (ell<= 2500)]/(2.*np.pi))
    sublgcbb = np.log10(subell**2 * cbb[(ell>=500) & (ell<= 2500)]/(2.*np.pi))
    A = np.vander(sublgell, 2)
    (coeffs_ee, residuals, rank, sing_vals) = np.linalg.lstsq(A, sublgcee)
    (coeffs_bb, residuals, rank, sing_vals) = np.linalg.lstsq(A, sublgcbb)
    if verbose:
        print "ell^2 * P_EE fit: amplitude at ell=1000 and power-law index are",10**coeffs_ee[1], coeffs_ee[0]
        print "ell^2 * P_EE fit: amplitude at ell=1000 and power-law index are",10**coeffs_bb[1], coeffs_bb[0]
    f_ee = 10**(coeffs_ee[0]*sublgell + coeffs_ee[1])
    f_bb = 10**(coeffs_bb[0]*sublgell + coeffs_bb[1])
    f_kol = 10**((-5./3.)*sublgell + 0.5*(coeffs_ee[1]+coeffs_bb[1]))
    if do_plot:
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(ell, ell*(ell+1)*cee/(2.*np.pi), marker='+', linestyle='None',
                label='Observed EE power')
        ax.plot(ell, ell*(ell+1)*cbb/(2.*np.pi), marker='+', linestyle='None',
                label='Observed BB power')
        ax.plot(ell, ell*(ell+1)*ceb/(2.*np.pi), marker='+', linestyle='None',
                label='Observed EB power')
        ax.plot(subell, f_ee, label='Best-fit power law (EE)')
        ax.plot(subell, f_bb, label='Best-fit power law (BB)')
        ax.plot(subell, f_kol, label='Kolmogorov expectation')
        lim = (1.e-6, 1.e-3)
        plt.ylim(lim)
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.set_xlabel('ell')
        ax.set_ylabel('ell*(ell+1)*C_ell/2pi')
        ax.set_title('Power spectrum from atmospheric PSF shapes')
        plt.legend(loc=4)
        min_ell = 2.*np.pi/(theta*np.pi/180.)
        max_ell = np.sqrt(2.)*ngrid*min_ell/2.
        ax.plot(np.array((min_ell,min_ell)), np.array(lim), color='black')
        ax.plot(np.array((max_ell,max_ell)), np.array(lim), color='black')
        plt.show()

    # Compare with atmospheric power spectra from papers, etc.

    # Reconstruct "kappa" field based on e1, e2, and compare with dfwhm.  Check statistical
    # properties, etc.

    # Save / output / return information
    # Average over all the many sub-grids?

# Main function: for now, just specify a single file directly.
infile = '/Users/rmandelb/great3/chihway-sims/catalogs/3min/focalplane_33.txt'
n_ell = 8
do_plot = True
verbose = True
process_file(infile, n_ell, verbose=verbose, do_plot=do_plot)

# Later, will need to write code that crawls through all directories and processes everything,
# stores and postprocesses results.

<|endoftext|>