<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Functions for supernova light curve I/O"""

# This module is designed to be self-contained, with only a dependency on
# numpy

# Try to get OrderedDict from collections, then from astropy, then just use
# the regular built-in dictionary.
try:
    from collections import OrderedDict as odict
except ImportError:
    try:
        from astropy.utils import OrderedDict as odict
    except ImportError:
        odict = dict

HAS_NUMPY = True
try:
    import numpy as np
except ImportError:
    HAS_NUMPY = False

__all__ = ['readlc', 'writelc']

def _cast_str(s):
    try:
        return int(s)
    except:
        try:
            return float(s)
        except:
            return s.strip()

def dict_to_array(d):
    """Convert a dictionary of lists (of equal length) to a structured
    numpy.ndarray"""
    import numpy as np

    # first convert all lists to 1-d arrays, in order to let numpy
    # figure out the necessary size of the string arrays.
    for key in d:
        d[key] = np.array(d[key])

    # Determine dtype of output array.
    dtypelist = []
    for key in d:
        dtypelist.append((key, d[key].dtype))

    # Initialize ndarray and then fill it.
    firstkey = d.keys()[0]
    col_len = len(d[firstkey])
    result = np.empty(col_len, dtype=dtypelist)
    for key in d:
        result[key] = d[key]

    return result

# ------------------------------- Readers ------------------------------------
def _read_csv(f, **kwargs):

    delim = kwargs.get('delim', None)
    metachar = kwargs.get('metachar', '@')
    commentchar = kwargs.get('commentchar', '#')

    meta = odict()
    colnames = []
    cols = []
    readingdata = False
    for line in f:

        # strip leading & trailing whitespace, newline, and comments
        line = line.strip()
        pos = line.find(commentchar)
        if pos > -1:
            line = line[:pos]
        if len(line) == 0: continue

        if not readingdata:
            # Read metadata
            if line[0] == metachar:
                pos = line.find(' ')  # Find first space.
                if pos in [-1, 1]:  # Space must exist and key must exist.
                    raise ValueError('Incorrectly formatted metadata line: '+
                                     line)
                meta[line[1:pos]] = _cast_str(line[pos:])
                continue

            # Read header line
            for item in line.split(delim):
                colnames.append(item.strip())
                cols.append([])
            readingdata = True
            continue

        # Now we're reading data
        items = line.split(delim)
        for col, item in zip(cols, items):
            col.append(_cast_str(item))

    data = odict(zip(colnames, cols))
    return meta, data

# SALT2 conversions upon reading:
# Names are converted to lowercase, then the following lookup table is used.
SALT2KEY_TO_KEY = {'redshift': 'z',
                   'z_heliocentric': 'z_helio',
                   'filter': 'band'}
def _salt2_rename_keys(d):
    newd = odict()
    for key, val in d.iteritems():
        key = key.lower()
        if key in SALT2KEY_TO_KEY:
            key = SALT2KEY_TO_KEY[key]
        newd[key] = val
    return newd

def _read_salt2(f, **kwargs):
    """Read a new-style SALT2 file.

    Such a file has metadata on lines starting with '@' and column names
    on lines starting with '#' and containing a ':' after the column name.
    There is optionally a line containing '#end' before the start of data.
    """

    raw = kwargs.get('raw', False)

    meta = odict()
    colnames = []
    cols = []
    readingdata = False
    for line in f:

        # strip leading & trailing whitespace & newline
        line = line.strip()
        if len(line) == 0: continue

        if not readingdata:
            # Read metadata
            if line[0] == '@':
                pos = line.find(' ')  # Find first space.
                if pos in [-1, 1]:  # Space must exist and key must exist.
                    raise ValueError('Incorrectly formatted metadata line: '+
                                     line)
                meta[line[1:pos]] = _cast_str(line[pos:])
                continue

            # Read header line
            if line[0] == '#':
                pos = line.find(':')
                if pos in [-1, 1]: continue  # comment line
                colname = line[1:pos].strip()
                if colname == 'end': continue
                colnames.append(colname)
                cols.append([])
                continue

            # If the first non-whitespace character is not '@' or '#',
            # assume the line is the first data line.
            readingdata = True

        # strip comments
        pos = line.find('#')
        if pos > -1:
            line = line[:pos]
        if len(line) == 0: continue

        # Now we're reading data
        items = line.split()
        for col, item in zip(cols, items):
            col.append(_cast_str(item))

    data = odict(zip(colnames, cols))

    # convert keys, if requested
    if not raw:
        meta = _salt2_rename_keys(meta)
        data = _salt2_rename_keys(data)

    return meta, data

READERS = {'csv': _read_csv,
           'salt2': _read_salt2}

def readlc(fname, fmt='csv', array=True, **kwargs):
    """Read light curve data.

    Parameters
    ----------
    fname : str
        Filename
    fmt : {'csv', 'salt2', 'snana'}, optional
        Format of file. Default is 'csv'. 'salt2' is the new format available
        in snfit version >= 2.3.0.
    array : bool, optional
        If True (default), returned data is a numpy array, Otherwise, data
        is a dictionary of lists (representing columns).
    delim : str, optional
        **[csv only]** Used to split entries on a line. Default is `None`.
        Extra whitespace is ignored.
    metachar : str, optional
        **[csv only]** Lines whose first non-whitespace character is `metachar`
        are treated as metadata lines, where the key and value are split on
        the first whitespace. Default is '@'
    commentchar : str, optional
        **[csv only]** One-character string indicating a comment. Default is
        '#'.
    raw : bool, optional
        **[salt2 only]** By default, the SALT2 reader converts all column names
        and metadata keys to lowercase and renames a few of them. Set to True
        to override this. Default is False.

    Returns
    -------
    meta : dict
        A (possibly empty) dictionary of metadata in the file.
    data : `~numpy.ndarray` or dict
        Data.
    """

    if fmt not in READERS:
        raise ValueError("Reader not defined for format '{}'. Options: "
                         .format(fmt) + ", ".join(READERS.keys()))

    with open(fname, 'rb') as f:
        meta, data = READERS[fmt](f, **kwargs)

    if array:
        if HAS_NUMPY:
            data = dict_to_array(data)
        else:
            raise ValueError('numpy required for array output. Use '
                             'array=False for dictionary output.')
    return meta, data

# ------------------------------- Writers -----------------------------------#
def _write_csv(f, data, meta, **kwargs):

    delim = kwargs.get('delim', ' ')
    metachar = kwargs.get('metachar', '@')

    if meta is not None:
        for key, val in meta.iteritems():
            f.write('{}{}{}{}\n'.format(metachar, key, delim, str(val)))

    if hasattr(data, 'dtype'):
        keys = data.dtype.names
        length = len(data)
    else:
        keys = data.keys()
        length = len(data[keys[0]])

    f.write(delim.join(keys))
    f.write('\n')
    for i in range(length):
        f.write(delim.join([str(data[key][i]) for key in keys]))
        f.write('\n')

KEY_TO_SALT2KEY_META = {
    'Z': 'REDSHIFT',              # Not sure if this is used.
    'Z_HELIOCENTRIC': 'Z_HELIO',
    'MAGSYS': 'MagSys',
    'Z_SOURCE': 'z_source'}
KEY_TO_SALT2KEY_COLUMN = {
    'Mjd': 'Date',
    'Time': 'Date',
    'Flux': 'FluxPsf',
    'Fluxpsf': 'FluxPsf',
    'Fluxerr': 'FluxPsferr',
    'Fluxpsferr': 'FluxPsferr',
    'Airmass': 'AirMass',
    'Zp': 'ZP',
    'Zpsys': 'MagSys',
    'Magsys': 'MagSys',
    'Band': 'Filter'}

def _write_salt2(f, data, meta, **kwargs):

    raw = kwargs.get('raw', False)
    pedantic = kwargs.get('pedantic', True)

    if meta is not None:
        for key, val in meta.iteritems():
            if not raw:
                key = key.upper()
                key = KEY_TO_SALT2KEY_META.get(key, key)
            f.write('@{} {}\n'.format(key, str(val)))

    if hasattr(data, 'dtype'):
        keys = data.dtype.names
        length = len(data)
    else:
        keys = data.keys()
        length = len(data[keys[0]])

    # Write column names
    keys_as_written = []
    for key in keys:
        if not raw:
            key = key.capitalize()
            key = KEY_TO_SALT2KEY_COLUMN.get(key, key)
        f.write('#{} :\n'.format(key))
        keys_as_written.append(key)
    f.write('#end :\n')

    # Check that necessary fields exist
    if pedantic:
        if not ('Filter' in keys_as_written and 'MagSys' in keys_as_written):
            raise ValueError('photometry data missing required some fields '
                             ': Filter, MagSys')

    # Write the data itself
    for i in range(length):
        f.write(' '.join([str(data[key][i]) for key in keys]))
        f.write('\n')

KEY_TO_SNANAKEY_COLUMN = {
    'TIME': 'MJD',
    'DATE': 'MJD',
    'FILTER': 'FLT',
    'BAND': 'FLT',
    'FLUX': 'FLUXCAL',
    'FLUXERR': 'FLUXCALERR',
    'ZP': 'ZPT',
    'ZEROPOINT': 'ZPT'
    }
KEY_TO_SNANAKEY_META = {
    'DEC': 'DECL'
    }
SNANA_REQUIRED_META = ['RA', 'DECL', 'SURVEY', 'FILTERS', 'MWEBV']
SNANA_REQUIRED_COLUMN = ['MJD', 'FLT', 'FLUXCAL', 'FLUXCALERR', 'ZPT']
def _write_snana(f, data, meta, **kwargs):

    raw = kwargs.get('raw', False)
    pedantic = kwargs.get('pedantic', True)

    # Write metadata
    keys_as_written = []
    if meta is not None:
        for key, val in meta.iteritems():
            if not raw:
                key = key.upper()
                key = KEY_TO_SNANAKEY_META.get(key, key)
            f.write('{}: {}\n'.format(key, str(val)))
            keys_as_written.append(key)

    # Check that necessary metadata was written
    if pedantic:
        for key in SNANA_REQUIRED_META:
            if key not in keys_as_written:
                raise ValueError('Missing required metadata kw: ' + key)

    # Get column names and data length
    if hasattr(data, 'dtype'):
        keys = data.dtype.names
        length = len(data)
    else:
        keys = data.keys()
        length = len(data[keys[0]])

    # Convert column names
    keys_to_write = []
    for key in keys:
        if not raw:
            key = key.upper()
            key = KEY_TO_SNANAKEY_COLUMN.get(key, key)
        keys_to_write.append(key)

    # Check that necessary column names are included
    if pedantic:
        for key in SNANA_REQUIRED_COLUMN:
            if key not in keys_to_write:
                raise ValueError('Missing required column name: ' + key)

    # Write the header
    f.write('\n'
            '# ==========================================\n'
            '# TERSE LIGHT CURVE OUTPUT:\n'
            '#\n'
            'NOBS: {:d}\n'
            'NVAR: {:d}\n'
            'VARLIST: {}\n'
            .format(length, len(keys), ' '.join(keys_to_write)))

    # Write data
    for i in range(length):
        f.write('OBS: ')
        f.write(' '.join([str(data[key][i]) for key in keys]))
        f.write('\n')

WRITERS = {'csv': _write_csv,
           'salt2': _write_salt2,
           'snana': _write_snana}

def writelc(data, fname, meta=None, fmt='csv', **kwargs):
    """Write light curve data.

    Parameters
    ----------
    data : `~numpy.ndarray` or dict
        Data.
    fname : str
        Filename.
    meta : dict, optional
        A (possibly empty) dictionary of metadata. Default is None.
    fmt : {'csv', 'salt2', 'snana'}, optional
        Format of file. Default is 'csv'. 'salt2' is the new format available
        in snfit version >= 2.3.0.
    delim : str, optional
        **[csv only]** Character used to separate entries on a line.
        Default is ' '.
    metachar : str, optional
        **[csv only]** Metadata designator. Default is '@'.
    raw : bool, optional
        **[salt2, snana]** By default, the SALT2 and SNANA writers rename
        some metadata keys and column names in order to comply with what
        snfit and SNANA expect. Set to True to override this.
        Default is False.
    pedantic : bool, optional
        **[salt2, snana]** If True, check that output column names and header
        keys comply with expected formatting, and raise a ValueError if not.
        It is probably a good idea to set to False when raw is True.
        Default is True.
    """

    if fmt not in WRITERS:
        raise ValueError("Writer not defined for format '{}'. Options: "
                         .format(fmt) + ", ".join(WRITERS.keys()))
    with open(fname, 'w') as f:
        WRITERS[fmt](f, data, meta, **kwargs)
<commit_msg>Add support for fits, json lightcurve files<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Functions for supernova light curve I/O"""

# This module is designed to be self-contained, with only a dependency on
# numpy. Since OrderedDict is only in 2.7+, try to get `OrderedDict` from
# collections, then from astropy, then just use a regular `dict`.
from warnings import warn
import json
import sys
import re
import numpy as np
try:
    from collections import OrderedDict as odict
except ImportError:
    try:
        from astropy.utils import OrderedDict as odict
    except ImportError:
        odict = dict

__all__ = ['read_lc', 'write_lc']

# Helper functions ========================================================== #
def _cast_str(s):
    try:
        return int(s)
    except:
        try:
            return float(s)
        except:
            return s.strip()

def dict_to_array(d):
    """Convert a dictionary of lists (or single values) to a structured
    numpy.ndarray."""

    # Convert all lists/values to 1-d arrays, in order to let numpy
    # figure out the necessary size of the string arrays.
    new_d = odict()
    for key in d:
        new_d[key] = np.atleast_1d(d[key])

    # Determine dtype of output array.
    dtype = [(key, arr.dtype) for key, arr in new_d.iteritems()]

    # Initialize ndarray and then fill it.
    col_len = max([len(v) for v in new_d.values()])
    result = np.empty(col_len, dtype=dtype)
    for key in new_d:
        result[key] = new_d[key]

    return result

# Reader: csv =============================================================== #
def _read_csv(f, **kwargs):

    delim = kwargs.get('delim', None)
    metachar = kwargs.get('metachar', '@')
    commentchar = kwargs.get('commentchar', '#')

    meta = odict()
    colnames = []
    cols = []
    readingdata = False
    for line in f:

        # strip leading & trailing whitespace, newline, and comments
        line = line.strip()
        pos = line.find(commentchar)
        if pos > -1:
            line = line[:pos]
        if len(line) == 0: continue

        if not readingdata:
            # Read metadata
            if line[0] == metachar:
                pos = line.find(' ')  # Find first space.
                if pos in [-1, 1]:  # Space must exist and key must exist.
                    raise ValueError('Incorrectly formatted metadata line: '+
                                     line)
                meta[line[1:pos]] = _cast_str(line[pos:])
                continue

            # Read header line
            for item in line.split(delim):
                colnames.append(item.strip())
                cols.append([])
            readingdata = True
            continue

        # Now we're reading data
        items = line.split(delim)
        for col, item in zip(cols, items):
            col.append(_cast_str(item))

    data = odict(zip(colnames, cols))
    return meta, data

# Reader: salt2 ============================================================= #

# SALT2 conversions upon reading:
# Names are converted to lowercase, then the following lookup table is used.
SALT2KEY_TO_KEY = {'redshift': 'z',
                   'z_heliocentric': 'z_helio',
                   'filter': 'band'}
def _salt2_rename_keys(d):
    newd = odict()
    for key, val in d.iteritems():
        key = key.lower()
        if key in SALT2KEY_TO_KEY:
            key = SALT2KEY_TO_KEY[key]
        newd[key] = val
    return newd

def _read_salt2(f, **kwargs):
    """Read a new-style SALT2 file.

    Such a file has metadata on lines starting with '@' and column names
    on lines starting with '#' and containing a ':' after the column name.
    There is optionally a line containing '#end' before the start of data.
    """

    raw = kwargs.get('raw', False)

    meta = odict()
    colnames = []
    cols = []
    readingdata = False
    for line in f:

        # strip leading & trailing whitespace & newline
        line = line.strip()
        if len(line) == 0: continue

        if not readingdata:
            # Read metadata
            if line[0] == '@':
                pos = line.find(' ')  # Find first space.
                if pos in [-1, 1]:  # Space must exist and key must exist.
                    raise ValueError('Incorrectly formatted metadata line: '+
                                     line)
                meta[line[1:pos]] = _cast_str(line[pos:])
                continue

            # Read header line
            if line[0] == '#':
                pos = line.find(':')
                if pos in [-1, 1]: continue  # comment line
                colname = line[1:pos].strip()
                if colname == 'end': continue
                colnames.append(colname)
                cols.append([])
                continue

            # If the first non-whitespace character is not '@' or '#',
            # assume the line is the first data line.
            readingdata = True

        # strip comments
        pos = line.find('#')
        if pos > -1:
            line = line[:pos]
        if len(line) == 0: continue

        # Now we're reading data
        items = line.split()
        for col, item in zip(cols, items):
            col.append(_cast_str(item))

    data = odict(zip(colnames, cols))

    # convert keys, if requested
    if not raw:
        meta = _salt2_rename_keys(meta)
        data = _salt2_rename_keys(data)

    return meta, data

# Reader: json ============================================================== #
def _read_json(f, **kwargs):
    t = json.load(f, encoding=sys.getdefaultencoding())
    d = {}
    for key, value in t['data'].items():
        d[key.encode('ascii')] = value
    return t['meta'], d

# Reader: fits ============================================================== #

def _read_fits(f, **kwargs):
    try:
        from astropy.io import fits
    except ImportError:
        try:
            import pyfits as fits
        except ImportError:
            raise ValueError("Either pyfits or astropy is required for 'fits'"
                             " format")
    hdulist = fits.open(f)

    # Parse all table objects
    tables = odict()
    for ihdu, hdu_item in enumerate(hdulist):
        if isinstance(hdu_item, (fits.TableHDU, fits.BinTableHDU,
                                 fits.GroupsHDU)):
            tables[ihdu] = hdu_item

    if len(tables) > 1:
        hdu = tables.keys()[0]
        warn("hdu= was not specified but multiple tables"
             " are present, reading in first available"
             " table (hdu={0})".format(hdu))
        hdu = hdulist.index_of(hdu)
        if hdu in tables:
            table = tables[hdu]
        else:
            raise ValueError("No table found in hdu={0}".format(hdu))

    elif len(tables) == 1:
        table = tables[tables.keys()[0]]
    else:
        raise ValueError("No table found")

    # Create a metadata dictionary.
    meta = odict()
    for key, value, comment in table.header.cards:
        if key in ['COMMENT', 'HISTORY']:
            if key in meta:
                meta[key].append(value)
            else:
                meta[key] = [value]

        elif key in meta:  # key is duplicate
            if isinstance(meta[key], list):
                meta[key].append(value)
            else:
                meta[key] = [meta[key], value]

        elif (is_column_keyword(key.upper()) or
              key.upper() in REMOVE_KEYWORDS):
            pass

        else:
            meta[key] = value

    data = table.data.view(np.ndarray)
    hdulist.close()

    return meta, data

# All readers =============================================================== #
READERS = {'csv': _read_csv,
           'json': _read_json,
           'fits': _read_fits,
           'salt2': _read_salt2}
def read_lc(fname, fmt='csv', array=True, **kwargs):
    """Read light curve data.

    Parameters
    ----------
    fname : str
        Filename
    fmt : {'csv', 'salt2', 'json'}, optional
        Format of file. Default is 'csv'. 'salt2' is the new format available
        in snfit version >= 2.3.0.
    array : bool, optional
        If True (default), returned data is a numpy array, Otherwise, data
        is a dictionary of lists (representing columns).
    delim : str, optional
        **[csv only]** Used to split entries on a line. Default is `None`.
        Extra whitespace is ignored.
    metachar : str, optional
        **[csv only]** Lines whose first non-whitespace character is `metachar`
        are treated as metadata lines, where the key and value are split on
        the first whitespace. Default is '@'
    commentchar : str, optional
        **[csv only]** One-character string indicating a comment. Default is
        '#'.
    raw : bool, optional
        **[salt2 only]** By default, the SALT2 reader converts all column names
        and metadata keys to lowercase and renames a few of them. Set to True
        to override this. Default is False.

    Returns
    -------
    meta : dict
        A (possibly empty) dictionary of metadata in the file.
    data : `~numpy.ndarray` or dict
        Data.
    """

    if fmt not in READERS:
        raise ValueError("Reader not defined for format '{}'. Options: "
                         .format(fmt) + ", ".join(READERS.keys()))

    with open(fname, 'rb') as f:
        meta, data = READERS[fmt](f, **kwargs)

    if array and not isinstance(data, np.ndarray):
        data = dict_to_array(data)

    return meta, data

# =========================================================================== #
# Writers                                                                     #
# =========================================================================== #

# Writer: csv =============================================================== #
def _write_csv(f, data, meta, **kwargs):

    delim = kwargs.get('delim', ' ')
    metachar = kwargs.get('metachar', '@')

    if meta is not None:
        for key, val in meta.iteritems():
            f.write('{}{}{}{}\n'.format(metachar, key, delim, str(val)))

    keys = data.dtype.names
    length = len(data)

    f.write(delim.join(keys))
    f.write('\n')
    for i in range(length):
        f.write(delim.join([str(data[key][i]) for key in keys]))
        f.write('\n')

# Writer: salt2 ============================================================= #
KEY_TO_SALT2KEY_META = {
    'Z': 'REDSHIFT',              # Not sure if this is used.
    'Z_HELIOCENTRIC': 'Z_HELIO',
    'MAGSYS': 'MagSys',
    'Z_SOURCE': 'z_source'}
KEY_TO_SALT2KEY_COLUMN = {
    'Mjd': 'Date',
    'Time': 'Date',
    'Flux': 'FluxPsf',
    'Fluxpsf': 'FluxPsf',
    'Fluxerr': 'FluxPsferr',
    'Fluxpsferr': 'FluxPsferr',
    'Airmass': 'AirMass',
    'Zp': 'ZP',
    'Zpsys': 'MagSys',
    'Magsys': 'MagSys',
    'Band': 'Filter'}

def _write_salt2(f, data, meta, **kwargs):

    raw = kwargs.get('raw', False)
    pedantic = kwargs.get('pedantic', True)

    if meta is not None:
        for key, val in meta.iteritems():
            if not raw:
                key = key.upper()
                key = KEY_TO_SALT2KEY_META.get(key, key)
            f.write('@{} {}\n'.format(key, str(val)))

    keys = data.dtype.names
    length = len(data)

    # Write column names
    keys_as_written = []
    for key in keys:
        if not raw:
            key = key.capitalize()
            key = KEY_TO_SALT2KEY_COLUMN.get(key, key)
        f.write('#{} :\n'.format(key))
        keys_as_written.append(key)
    f.write('#end :\n')

    # Check that necessary fields exist
    if pedantic:
        if not ('Filter' in keys_as_written and 'MagSys' in keys_as_written):
            raise ValueError('photometry data missing required some fields '
                             ': Filter, MagSys')

    # Write the data itself
    for i in range(length):
        f.write(' '.join([str(data[key][i]) for key in keys]))
        f.write('\n')

# Writer: snana ============================================================= #
KEY_TO_SNANAKEY_COLUMN = {
    'TIME': 'MJD',
    'DATE': 'MJD',
    'FILTER': 'FLT',
    'BAND': 'FLT',
    'FLUX': 'FLUXCAL',
    'FLUXERR': 'FLUXCALERR',
    'ZP': 'ZPT',
    'ZEROPOINT': 'ZPT'
    }
KEY_TO_SNANAKEY_META = {
    'DEC': 'DECL'
    }
SNANA_REQUIRED_META = ['RA', 'DECL', 'SURVEY', 'FILTERS', 'MWEBV']
SNANA_REQUIRED_COLUMN = ['MJD', 'FLT', 'FLUXCAL', 'FLUXCALERR', 'ZPT']
def _write_snana(f, data, meta, **kwargs):

    raw = kwargs.get('raw', False)
    pedantic = kwargs.get('pedantic', True)

    # Write metadata
    keys_as_written = []
    if meta is not None:
        for key, val in meta.iteritems():
            if not raw:
                key = key.upper()
                key = KEY_TO_SNANAKEY_META.get(key, key)
            f.write('{}: {}\n'.format(key, str(val)))
            keys_as_written.append(key)

    # Check that necessary metadata was written
    if pedantic:
        for key in SNANA_REQUIRED_META:
            if key not in keys_as_written:
                raise ValueError('Missing required metadata kw: ' + key)

    # Get column names and data length
    keys = data.dtype.names
    length = len(data)

    # Convert column names
    keys_to_write = []
    for key in keys:
        if not raw:
            key = key.upper()
            key = KEY_TO_SNANAKEY_COLUMN.get(key, key)
        keys_to_write.append(key)

    # Check that necessary column names are included
    if pedantic:
        for key in SNANA_REQUIRED_COLUMN:
            if key not in keys_to_write:
                raise ValueError('Missing required column name: ' + key)

    # Write the header
    f.write('\n'
            '# ==========================================\n'
            '# TERSE LIGHT CURVE OUTPUT:\n'
            '#\n'
            'NOBS: {:d}\n'
            'NVAR: {:d}\n'
            'VARLIST: {}\n'
            .format(length, len(keys), ' '.join(keys_to_write)))

    # Write data
    for i in range(length):
        f.write('OBS: ')
        f.write(' '.join([str(data[key][i]) for key in keys]))
        f.write('\n')

# Writer: json ============================================================== #
def _write_json(f, data, meta, **kwargs):

    # Build a dictionary of pure-python objects
    output = odict([('meta', meta),
                    ('data', odict())])
    for key in data.dtype.names:
        output['data'][key] = data[key].tolist()
    json.dump(output, f, encoding=sys.getdefaultencoding())
    del output

# Writer: fits ============================================================== #

# Adapted from astropy.io.fits.connect.write_table_fits

# Keywords to remove for all tables that are read in
REMOVE_KEYWORDS = ['XTENSION', 'BITPIX', 'NAXIS', 'NAXIS1', 'NAXIS2',
                   'PCOUNT', 'GCOUNT', 'TFIELDS']

# Column-specific keywords
COLUMN_KEYWORDS = ['TFORM[0-9]+',
                   'TBCOL[0-9]+',
                   'TSCAL[0-9]+',
                   'TZERO[0-9]+',
                   'TNULL[0-9]+',
                   'TTYPE[0-9]+',
                   'TUNIT[0-9]+',
                   'TDISP[0-9]+',
                   'TDIM[0-9]+',
                   'THEAP']

def is_column_keyword(keyword):
    for c in COLUMN_KEYWORDS:
        if re.match(c, keyword) is not None:
            return True
    return False

def _write_fits(f, data, meta, **kwargs):
    try:
        from astropy.io import fits
    except ImportError:
        try:
            import pyfits as fits
        except ImportError:
            raise ValueError("Either pyfits or astropy is required for 'fits'"
                             " format")

    table_hdu = fits.BinTableHDU(data)

    for key, value in meta.items():
        if is_column_keyword(key.upper()) or key.upper() in REMOVE_KEYWORDS:
            warn("Meta-data keyword {0} will be ignored since it "
                 "conflicts with a FITS reserved keyword".format(key),
                 RuntimeWarning)

        if isinstance(value, list):
            for item in value:
                try:
                    table_hdu.header.append((key, item))
                except ValueError:
                    warn("Attribute `{0}` of type {1} cannot be written "
                         "to FITS files - skipping".format(key, type(value)),
                         RuntimeWarning)
        else:
            try:
                table_hdu.header[key] = value
            except ValueError:
                warn("Attribute `{0}` of type {1} cannot be written to "
                     "FITS files - skipping".format(key, type(value)),
                     RuntimeWarning)

    table_hdu.writeto(f)

# All writers =============================================================== #
WRITERS = {'csv': _write_csv,
           'salt2': _write_salt2,
           'snana': _write_snana,
           'json': _write_json,
           'fits': _write_fits}

def write_lc(data, fname, meta=None, fmt='csv', **kwargs):
    """Write light curve data.

    Parameters
    ----------
    data : `~numpy.ndarray` or dict
        Data.
    fname : str
        Filename.
    meta : dict, optional
        A (possibly empty) dictionary of metadata. Default is None.
    fmt : {'csv', 'salt2', 'snana', 'json', 'fits'}, optional
        Format of file. Default is 'csv'. 'salt2' is the new format available
        in snfit version >= 2.3.0.
    delim : str, optional
        **[csv only]** Character used to separate entries on a line.
        Default is ' '.
    metachar : str, optional
        **[csv only]** Metadata designator. Default is '@'.
    raw : bool, optional
        **[salt2, snana]** By default, the SALT2 and SNANA writers rename
        some metadata keys and column names in order to comply with what
        snfit and SNANA expect. Set to True to override this.
        Default is False.
    pedantic : bool, optional
        **[salt2, snana]** If True, check that output column names and header
        keys comply with expected formatting, and raise a ValueError if not.
        It is probably a good idea to set to False when raw is True.
        Default is True.
    """

    if fmt not in WRITERS:
        raise ValueError("Writer not defined for format '{}'. Options: "
                         .format(fmt) + ", ".join(WRITERS.keys()))

    if not isinstance(data, np.ndarray):
        data = dict_to_array(data)

    if meta is None:
        meta = {}

    with open(fname, 'wb') as f:
        WRITERS[fmt](f, data, meta, **kwargs)
<|endoftext|>


<commit_before>
#
# This source file is part of appleseed.
# Visit http://appleseedhq.net/ for additional information and resources.
#
# This software is released under the MIT license.
#
# Copyright (c) 2013 Franz Beaune, Joel Daniels, Esteban Tovagliari.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#

import bpy
import nodeitems_utils
from bpy.types          import NodeTree
from bpy.app.handlers   import persistent
from ...util            import addon_dir, join_names_underscore
import os

#--------------------------------
# Class for appleseed node tree.
#--------------------------------
class AppleseedNodeTree( NodeTree):
    '''Appleseed Node Editor'''
    bl_idname = 'AppleseedNodeTree'
    bl_label = 'Appleseed Node Tree'
    bl_icon = 'NODETREE'

    @classmethod
    def poll( cls, context):
        renderer = context.scene.render.engine
        return renderer == 'APPLESEED_RENDER'

#--------------------------------
# Base class for appleseed nodes.
#--------------------------------
class AppleseedNode:
    @classmethod
    def poll( cls, context):
        renderer = context.scene.render.engine
        return context.bl_idname == "AppleseedNodeTree" and renderer == 'APPLESEED_RENDER'

    def get_node_name( self):
        '''
        Return the node's name, including appended pointer
        '''
        return join_names_underscore( self.name, str(self.as_pointer()))

    def traverse_tree( self, material_node):
        '''
        Iterate inputs and traverse the tree backward if any inputs are connected.
        Nodes are added to a list attribute of the material output node.
        '''
        for socket in self.inputs:
            if socket.is_linked:
                linked_node = socket.links[0].from_node
                linked_node.traverse_tree( material_node)
        material_node.tree.append( self)

#--------------------------------
# Base class for appleseed sockets.
#--------------------------------
class AppleseedSocket( object):
    # Set to default None.
    socket_value = None

    def get_socket_value( self, texture_only = True):
        '''
        Method to return socket's value, if not linked.
        If linked, return the name of the node with appended pointer.
        '''
        if self.is_linked:
            linked_node = self.links[0].from_node
            if texture_only and linked_node.node_type == 'texture':
                # The socket only accepts image textures.
                return linked_node.get_node_name() + "_inst"
            if not texture_only:
                return linked_node.get_node_name()
        # Return socket value if not linked, or if the incoming node is incompatible.
        return self.socket_value

#--------------------------------
# Node category for extending the Add menu, toolbar panels
#   and search operator
# Base class for node categories
#--------------------------------
class AppleseedNodeCategory( nodeitems_utils.NodeCategory):
    @classmethod
    def poll(cls, context):
        renderer = context.scene.render.engine
        return context.space_data.tree_type == 'AppleseedNodeTree' and renderer == 'APPLESEED_RENDER'

#--------------------------------
# appleseed node categories
# identifier, label, items list
#--------------------------------
appleseed_node_categories = [
    AppleseedNodeCategory("BSDF", "BSDF", items = [
        nodeitems_utils.NodeItem( "AppleseedAshikhminNode"),
        nodeitems_utils.NodeItem( "AppleseedDiffuseBTDFNode"),
        nodeitems_utils.NodeItem( "AppleseedDisneyNode"),
        nodeitems_utils.NodeItem( "AppleseedKelemenNode"),
        nodeitems_utils.NodeItem( "AppleseedLambertianNode"),
        nodeitems_utils.NodeItem( "AppleseedMicrofacetNode"),
        nodeitems_utils.NodeItem( "AppleseedOrenNayarNode"),
        nodeitems_utils.NodeItem( "AppleseedSpecBRDFNode"),
        nodeitems_utils.NodeItem( "AppleseedSpecBTDFNode"),
        nodeitems_utils.NodeItem( "AppleseedBlendNode")]),
    AppleseedNodeCategory( "TEXTURES", "Texture", items = [
        nodeitems_utils.NodeItem( "AppleseedTexNode"),
        nodeitems_utils.NodeItem( "AppleseedNormalNode")]),
    AppleseedNodeCategory( "OUTPUTS", "Output", items = [
        nodeitems_utils.NodeItem( "AppleseedMaterialNode")])]

@persistent
def appleseed_scene_loaded(dummy):
    # Load images as icons
    icon16 = bpy.data.images.get('APPLESEED16')
    icon32 = bpy.data.images.get('APPLESEED32')
    if icon16 is None:
        img = bpy.data.images.load(os.path.join( os.path.join( addon_dir, 'icons'), 'APPLESEED16.png'))
        img.name = 'APPLESEED16'
        img.use_alpha = True
        img.user_clear()
    # remove scene_update handler
    elif "APPLESEED16" not in icon16.keys():
        icon16["APPLESEED16"] = True
        for f in bpy.app.handlers.scene_update_pre:
            if f.__name__ == "appleseed_scene_loaded":
                bpy.app.handlers.scene_update_pre.remove(f)
    if icon32 is None:
        img = bpy.data.images.load(os.path.join( os.path.join( addon_dir, 'icons'), 'APPLESEED32.png'))
        img.name = 'APPLESEED32'
        img.use_alpha = True
        img.user_clear() # Won't get saved into .blend files
    # remove scene_update handler
    elif "APPLESEED32" not in icon32.keys():
        icon32["APPLESEED32"] = True
        for f in bpy.app.handlers.scene_update_pre:
            if f.__name__ == "appleseed_scene_loaded":
                bpy.app.handlers.scene_update_pre.remove(f)

# Load the modules after classes have been created.
from .          import ashikhmin_brdf
from .          import bsdf_blend
from .          import diffuse_btdf
from .          import disney_brdf
from .          import kelemen_brdf
from .          import lambertian_brdf
from .          import microfacet_brdf
from .          import orennayar_brdf
from .          import specular_brdf
from .          import specular_btdf
from .          import texture
from .          import normal
from .          import material

def register():
    bpy.app.handlers.load_post.append( appleseed_scene_loaded)
    bpy.app.handlers.scene_update_pre.append( appleseed_scene_loaded)
    nodeitems_utils.register_node_categories("APPLESEED", appleseed_node_categories)
    bpy.utils.register_class( AppleseedNodeTree)
    ashikhmin_brdf.register()
    bsdf_blend.register()
    diffuse_btdf.register()
    disney_brdf.register()
    kelemen_brdf.register()
    lambertian_brdf.register()
    microfacet_brdf.register()
    orennayar_brdf.register()
    specular_brdf.register()
    specular_btdf.register()
    texture.register()
    normal.register()
    material.register()

def unregister():
    nodeitems_utils.unregister_node_categories("APPLESEED")
    bpy.utils.unregister_class( AppleseedNodeTree)
    ashikhmin_brdf.unregister()
    bsdf_blend.unregister()
    disney_brdf.unregister()
    diffuse_btdf.unregister()
    kelemen_brdf.unregister()
    lambertian_brdf.unregister()
    microfacet_brdf.unregister()
    orennayar_brdf.unregister()
    specular_brdf.unregister()
    specular_btdf.unregister()
    texture.unregister()
    material.unregister()
    normal.unregister()
    bpy.app.handlers.load_post.remove( appleseed_scene_loaded)
<commit_msg>Disable specular btdf from node list<commit_after>
#
# This source file is part of appleseed.
# Visit http://appleseedhq.net/ for additional information and resources.
#
# This software is released under the MIT license.
#
# Copyright (c) 2013 Franz Beaune, Joel Daniels, Esteban Tovagliari, Luke Kliber.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#

import bpy
import nodeitems_utils
from bpy.types          import NodeTree
from bpy.app.handlers   import persistent
from ...util            import addon_dir, join_names_underscore
import os

#--------------------------------
# Class for appleseed node tree.
#--------------------------------
class AppleseedNodeTree( NodeTree):
    '''Appleseed Node Editor'''
    bl_idname = 'AppleseedNodeTree'
    bl_label = 'Appleseed Node Tree'
    bl_icon = 'NODETREE'

    @classmethod
    def poll( cls, context):
        renderer = context.scene.render.engine
        return renderer == 'APPLESEED_RENDER'

#--------------------------------
# Base class for appleseed nodes.
#--------------------------------
class AppleseedNode:
    @classmethod
    def poll( cls, context):
        renderer = context.scene.render.engine
        return context.bl_idname == "AppleseedNodeTree" and renderer == 'APPLESEED_RENDER'

    def get_node_name( self):
        '''
        Return the node's name, including appended pointer
        '''
        return join_names_underscore( self.name, str(self.as_pointer()))

    def traverse_tree( self, material_node):
        '''
        Iterate inputs and traverse the tree backward if any inputs are connected.
        Nodes are added to a list attribute of the material output node.
        '''
        for socket in self.inputs:
            if socket.is_linked:
                linked_node = socket.links[0].from_node
                linked_node.traverse_tree( material_node)
        material_node.tree.append( self)

#--------------------------------
# Base class for appleseed sockets.
#--------------------------------
class AppleseedSocket( object):
    # Set to default None.
    socket_value = None

    def get_socket_value( self, texture_only = True):
        '''
        Method to return socket's value, if not linked.
        If linked, return the name of the node with appended pointer.
        '''
        if self.is_linked:
            linked_node = self.links[0].from_node
            if texture_only and linked_node.node_type == 'texture':
                # The socket only accepts image textures.
                return linked_node.get_node_name() + "_inst"
            if not texture_only:
                return linked_node.get_node_name()
        # Return socket value if not linked, or if the incoming node is incompatible.
        return self.socket_value

#--------------------------------
# Node category for extending the Add menu, toolbar panels
#   and search operator
# Base class for node categories
#--------------------------------
class AppleseedNodeCategory( nodeitems_utils.NodeCategory):
    @classmethod
    def poll(cls, context):
        renderer = context.scene.render.engine
        return context.space_data.tree_type == 'AppleseedNodeTree' and renderer == 'APPLESEED_RENDER'

#--------------------------------
# appleseed node categories
# identifier, label, items list
#--------------------------------
appleseed_node_categories = [
    AppleseedNodeCategory("BSDF", "BSDF", items = [
        nodeitems_utils.NodeItem( "AppleseedAshikhminNode"),
        nodeitems_utils.NodeItem( "AppleseedDiffuseBTDFNode"),
        nodeitems_utils.NodeItem( "AppleseedDisneyNode"),
        nodeitems_utils.NodeItem( "AppleseedKelemenNode"),
        nodeitems_utils.NodeItem( "AppleseedLambertianNode"),
        nodeitems_utils.NodeItem( "AppleseedMicrofacetNode"),
        nodeitems_utils.NodeItem( "AppleseedOrenNayarNode"),
        nodeitems_utils.NodeItem( "AppleseedSpecBRDFNode"),
        # nodeitems_utils.NodeItem( "AppleseedSpecBTDFNode"),
        nodeitems_utils.NodeItem( "AppleseedBlendNode")]),
    AppleseedNodeCategory( "TEXTURES", "Texture", items = [
        nodeitems_utils.NodeItem( "AppleseedTexNode"),
        nodeitems_utils.NodeItem( "AppleseedNormalNode")]),
    AppleseedNodeCategory( "OUTPUTS", "Output", items = [
        nodeitems_utils.NodeItem( "AppleseedMaterialNode")])]

@persistent
def appleseed_scene_loaded(dummy):
    # Load images as icons
    icon16 = bpy.data.images.get('APPLESEED16')
    icon32 = bpy.data.images.get('APPLESEED32')
    if icon16 is None:
        img = bpy.data.images.load(os.path.join( os.path.join( addon_dir, 'icons'), 'APPLESEED16.png'))
        img.name = 'APPLESEED16'
        img.use_alpha = True
        img.user_clear()
    # remove scene_update handler
    elif "APPLESEED16" not in icon16.keys():
        icon16["APPLESEED16"] = True
        for f in bpy.app.handlers.scene_update_pre:
            if f.__name__ == "appleseed_scene_loaded":
                bpy.app.handlers.scene_update_pre.remove(f)
    if icon32 is None:
        img = bpy.data.images.load(os.path.join( os.path.join( addon_dir, 'icons'), 'APPLESEED32.png'))
        img.name = 'APPLESEED32'
        img.use_alpha = True
        img.user_clear() # Won't get saved into .blend files
    # remove scene_update handler
    elif "APPLESEED32" not in icon32.keys():
        icon32["APPLESEED32"] = True
        for f in bpy.app.handlers.scene_update_pre:
            if f.__name__ == "appleseed_scene_loaded":
                bpy.app.handlers.scene_update_pre.remove(f)

# Load the modules after classes have been created.
from .          import ashikhmin_brdf
from .          import bsdf_blend
from .          import diffuse_btdf
from .          import disney_brdf
from .          import kelemen_brdf
from .          import lambertian_brdf
from .          import microfacet_brdf
from .          import orennayar_brdf
from .          import specular_brdf
from .          import specular_btdf
from .          import texture
from .          import normal
from .          import material

def register():
    bpy.app.handlers.load_post.append( appleseed_scene_loaded)
    bpy.app.handlers.scene_update_pre.append( appleseed_scene_loaded)
    nodeitems_utils.register_node_categories("APPLESEED", appleseed_node_categories)
    bpy.utils.register_class( AppleseedNodeTree)
    ashikhmin_brdf.register()
    bsdf_blend.register()
    diffuse_btdf.register()
    disney_brdf.register()
    kelemen_brdf.register()
    lambertian_brdf.register()
    microfacet_brdf.register()
    orennayar_brdf.register()
    specular_brdf.register()
    specular_btdf.register()
    texture.register()
    normal.register()
    material.register()

def unregister():
    nodeitems_utils.unregister_node_categories("APPLESEED")
    bpy.utils.unregister_class( AppleseedNodeTree)
    ashikhmin_brdf.unregister()
    bsdf_blend.unregister()
    disney_brdf.unregister()
    diffuse_btdf.unregister()
    kelemen_brdf.unregister()
    lambertian_brdf.unregister()
    microfacet_brdf.unregister()
    orennayar_brdf.unregister()
    specular_brdf.unregister()
    specular_btdf.unregister()
    texture.unregister()
    material.unregister()
    normal.unregister()
    bpy.app.handlers.load_post.remove( appleseed_scene_loaded)
<|endoftext|>


<commit_before>
import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvRec(nn.Module):
    """
    Implementation of Efficient Char-level Document Classification
    by Combining Convolution and Recurrent Layers:
    https://arxiv.org/pdf/1602.00367.pdf

    Parameters:
    -----------
    n_classes: int, number of classes
    vocab: int, vocabulary length
    emb_dim: int, embedding dimention
    out_channels: number of channels for all kernels
       This can't vary across filters their output congruency
    kernel_sizes: tuple of int, one for each kernel, i.e. number of kernels
       will equal the length of this argument. In practice, this parameter
       only controls the width of each filter since the height is fixed to
       the dimension of the embeddings to ensure a kernel output height of 1
       (the kernel output width will vary depending on the input seq_len,
       but will as well get max pooled over)
    hid_dim: int, RNN hidden dimension
    cell: str, one of 'LSTM', 'GRU', 'RNN'
    num_layers: int, depth of the RNN
    bidi: bool, whether to use bidirectional RNNs
    dropout: float
    act: str, activation function after the convolution
    padding_idx: int or None, zero the corresponding output embedding
    """
    def __init__(self, n_classes,
                 # embedding parameters
                 vocab, emb_dim,
                 # cnn parameters
                 out_channels, kernel_sizes,
                 # rnn parameters
                 hid_dim, cell='LSTM', num_layers=1, bidi=True,
                 # rest parameters
                 dropout=0.0, act='relu', padding_idx=None):
        self.vocab = vocab
        self.emb_dim = emb_dim

        self.out_channels = out_channels
        self.kernel_sizes = kernel_sizes

        self.hid_dim = hid_dim
        self.cell = cell
        self.num_layers = num_layers
        self.bidi = bidi

        self.dropout = dropout
        self.act = act

        super(ConvRNN, self).__init__()

        # embeddings
        self.embeddings = nn.Embedding(
            self.vocab, self.emb_dim, padding_idx=padding_idx)

        # convs
        self.convs = [nn.Conv2d(1,  # text only has one channel
                                self.out_channels,
                                # number of filters WxH for each kernel (H is
                                # fixed to emb_dim to ensure H_out equals one)
                                (self.emb_dim, kernel_size))
                      for kernel_size in self.kernel_sizes]

        # RNN
        self.rnn = getattr(nn, self.cell)(
            self.out_channels, self.hid_dim, self.num_layers,
            dropout=self.dropout, bidirectional=self.bidi)

        # projection
        self.proj = nn.Linear(
            2 * self.hid_dim,
            n_classes)

    def forward(self, inp):
        # Embedding
        emb = self.embeddings(inp).t()  # (batch x seq_len x emb_dim)
        emb = emb.transpose(1, 2).unsqueeze(1)  # (batch x 1 x emb_dim x seq_len)

        # CNN
        # [(batch x out_channels x H_out (= 1) x W_out), ...]
        conv_out = [getattr(F, self.act)(conv(emb)) for conv in self.convs]
        # [(batch x out_channels x W_out), ...]
        conv_out = [out_i.squeeze(2) for out_i in conv_out]
        # [(batch x out_channels x 1), ...]
        pool_out = [F.max_pool1d(out_i, out_i.size(2)) for out_i in conv_out]
        # [(batch x out_channels), ...]
        pool_out = [out_i.squeeze(2) for out_i in pool_out]
        pool_out = torch.stack(pool_out)  # (num_kernels x batch x out_channels)
        pool_out = F.dropout(out, dropout=self.dropout, training=self.training)

        # RNN
        rnn_out, _ = self.rnn(pool_out)
        rnn_out = rnn_out[-1]   # take only last layer & last state output
        return F.log_softmax(self.proj(h_n)
<commit_msg>ConvRec WIP<commit_after>
import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvRec(nn.Module):
    """
    Implementation of Efficient Char-level Document Classification
    by Combining Convolution and Recurrent Layers:
    https://arxiv.org/pdf/1602.00367.pdf

    Parameters:
    -----------
    n_classes: int, number of classes
    vocab: int, vocabulary length
    emb_dim: int, embedding dimention
    out_channels: number of channels for all kernels
       This can't vary across filters their output congruency
    kernel_sizes: tuple of int, one for each kernel, i.e. number of kernels
       will equal the length of this argument. In practice, this parameter
       only controls the width of each filter since the height is fixed to
       the dimension of the embeddings to ensure a kernel output height of 1
       (the kernel output width will vary depending on the input seq_len,
       but will as well get max pooled over)
    hid_dim: int, RNN hidden dimension
    cell: str, one of 'LSTM', 'GRU', 'RNN'
    num_layers: int, depth of the RNN
    bidi: bool, whether to use bidirectional RNNs
    dropout: float
    act: str, activation function after the convolution
    padding_idx: int or None, zero the corresponding output embedding
    """
    def __init__(self, n_classes,
                 # embedding parameters
                 vocab, emb_dim,
                 # cnn parameters
                 out_channels, kernel_sizes,
                 # rnn parameters
                 hid_dim, cell='LSTM', num_layers=1, bidi=True,
                 # rest parameters
                 dropout=0.0, act='ReLU', padding_idx=None):
        self.vocab = vocab
        self.emb_dim = emb_dim

        self.in_channels = 1    # text has only a channel
        self.out_channels = out_channels
        self.kernel_sizes = kernel_sizes

        self.hid_dim = hid_dim
        self.cell = cell
        self.num_layers = num_layers
        self.bidi = bidi

        self.dropout = dropout
        self.act = act

        super(ConvRNN, self).__init__()

        # Embeddings
        self.embeddings = nn.Embedding(
            self.vocab, self.emb_dim, padding_idx=padding_idx)

        # CNN
        self.convs = []
        for kernel_size in self.kernel_sizes:
            self.convs.append(nn.Sequential(
                nn.Conv2d(self.in_channels, self.out_channels,
                          # number of filters HxW for each kernel (H is
                          # fixed to emb_dim to ensure H_out equals one)
                          (self.emb_dim, kernel_size)),
                getattr(nn, self.act)))

        # RNN
        self.rnn = getattr(nn, self.cell)(
            self.out_channels, self.hid_dim, self.num_layers,
            dropout=self.dropout, bidirectional=self.bidi)

        # Projection
        self.proj = nn.Linear(
            2 * self.hid_dim,
            n_classes)

    def forward(self, inp):
        # Embedding
        emb = self.embeddings(inp).t()  # (batch x seq_len x emb_dim)
        emb = emb.transpose(1, 2).unsqueeze(1)  # (batch x 1 x emb_dim x seq_len)

        # CNN
        for conv in self.convs:
            # (batch x out_channels x 1 x W_out)
            out = conv(emb)
            # (batch x out_channels x )
            out = F.max_pool1d(out, (1, out.size(3)))
            out = F.dropout(out, p=self.dropout, training=self.training)
            emb = out.view()
        conv_out = [getattr(F, self.act)(conv(emb)) for conv in self.convs]
        # [(batch x out_channels x W_out), ...]
        conv_out = [out_i.squeeze(2) for out_i in conv_out]
        # [(batch x out_channels x 1), ...]
        pool_out = [F.max_pool1d(out_i, out_i.size(2)) for out_i in conv_out]
        # [(batch x out_channels), ...]
        pool_out = [out_i.squeeze(2) for out_i in pool_out]
        pool_out = torch.stack(pool_out)  # (num_kernels x batch x out_channels)
        pool_out = F.dropout(out, dropout=self.dropout, training=self.training)

        # RNN
        rnn_out, _ = self.rnn(pool_out)
        rnn_out = rnn_out[-1]   # take only last layer & last state output
        return F.log_softmax(self.proj(h_n)
<|endoftext|>


<commit_before><commit_msg>rename simSystem in SimBuilding()<commit_after><|endoftext|>


<commit_before>"""
Socket Threaded.

Aims to provide some simplified abstraction over raw sockets such as making
concurrent socket connection easier.
"""

from __future__ import (
    absolute_import,
    division,
    print_function,
    unicode_literals
)
import socket

import sys
import threading

if sys.version_info.major == 2:
    chr = unichr
    string_type = basestring
    import Queue as queue
elif sys.version_info.major == 3:
    # chr should assume Unicode
    string_type = str
    import queue

def receive_all(sock, n):
    """
    Helper function to fully receive an arbitrary amount of data from a socket.

    :param sock: Socket connection
    :param n: Number of bytes to receive
    :return: Received data
    """
    data = b''
    while len(data) < n:
        packet = sock.recv(n - len(data))
        if not packet:
            break
            #return None
        data += packet
    return data

class SocketCommand(object):
    """
    Command object for communicating with SocketThread.
    command must be either:     payload must be:

    SocketCommand.CONNECT       (host, port) tuple
    SocketCommand.SEND          Binary data string
    SocketCommand.RECEIVE       Number of bytes to receive
    SocketCommand.CLOSE         None
    """

    # The available socket commands:
    CONNECT = "connect"
    SEND = "send"
    RECEIVE = "receive"
    CLOSE = "close"

    def __init__(self, command, payload=None):
        self.command = command
        self.payload = payload

class SocketReply(object):
    """
    Reply object for communicating with SocketThread.
    The reply must be either:       payload must be:

    SocketReply.ERROR               The Error object
    SocketReply.SUCCESS
    """
    # TODO: write description for SocketReply.SUCCESS above

    # The available socket replies
    SUCCESS = "success"
    ERROR = "error"

    def __init__(self, status, payload=None):
        self.status = status
        self.payload = payload

class SocketThread(threading.Thread):
    """
    Implements the threading.Thread interface
    """
    # TODO: should we keep track of what command corresponds to what reply?
    # TODO: What if message first sends a length prefix? implement this?

    def __init__(self):
        super(SocketThread, self).__init__()

        self.command_queue = queue.Queue()
        self.reply_queue = queue.Queue()
        self.socket = None

        self.alive = threading.Event()
        self.start()  # TODO: Should we really call this here?
        self.alive.set()

    def run(self):
        while self.alive.isSet():
            try:
                # Use a timeout value so we don't block forever
                cmd = self.command_queue.get(block=True, timeout=0.1)
                if cmd.command == SocketCommand.CONNECT:
                    address = cmd.payload
                    self._handle_CONNECT(address)
                elif cmd.command == SocketCommand.CLOSE:
                    self._handle_CLOSE()
                elif cmd.command == SocketCommand.SEND:
                    payload = cmd.payload
                    self._handle_SEND(payload)
                elif cmd.command == SocketCommand.RECEIVE:
                    number_of_bytes = cmd.payload
                    self._handle_RECEIVE(number_of_bytes)
                else:
                    # TODO: Handle invalid command
                    pass

                # TODO: should we call command_queue.task_done() here?
            except queue.Empty:
                continue

    def join(self, timeout=None):
        self.alive.clear()
        threading.Thread.join(self, timeout)

    def connect(self, address):
        """
        Connects the socket the the given address
        :param address: (host, port) Tuple
        """
        self.command_queue.put(SocketCommand(SocketCommand.CONNECT, address))

    def close(self):
        """
        Closes the socket
        """
        self.command_queue.put(SocketCommand(SocketCommand.CONNECT))

    def send(self, payload):
        """
        Sends the given payload to the socket. Requires an open and valid socket
        :param payload: Byte string of data
        """
        self.command_queue.put(SocketCommand(SocketCommand.SEND, payload))

    def receive(self, n):
        """
        Receives a specified number of bytes from the socket. Requires an open
        and valid socket.
        :param n: Number of bytes to receive from the socket.
        """
        self.command_queue.put(SocketCommand(SocketCommand.RECEIVE, n))

    def get_reply(self, block=True, timeout=None):
        """
        Fetches a reply from the reply queue.
        :param block: Boolean whether the call should block until a reply i
        present
        :param timeout: If block is true wait maximum timeout number of seconds
        :return: A SocketReply object is returned if possible. Else if returns
        None.
        """
        try:
            reply = self.reply_queue.get(block=block, timeout=timeout)
            # TODO: should we call reply_queue.task_done() here?
            return reply
        except queue.Empty:
            return None

    def _handle_CONNECT(self, address):
        """
        Handles the connection command by creating a tcp connection to the given
        address.

        :param address: (host, port) tuple
        """
        self.socket = socket.socket(socket.AF_INET,
                                    socket.SOCK_STREAM)  # tcp connection

        try:
            self.socket.connect(address)
            self.reply_queue.put(SocketReply(SocketReply.SUCCESS))
        except socket.error as e:
            self.reply_queue.put(SocketReply(SocketReply.ERROR, e))

    def _handle_CLOSE(self):
        """
        Handles the close command. This requires that a socket already has been
        opened.
        """
        self.socket.close()
        self.reply_queue.put(SocketReply(SocketReply.SUCCESS))

    def _handle_SEND(self, payload):
        """
        Handles the send command. This requires an open and valid socket
        connection.
        :param payload:
        """
        try:
            self.socket.sendall(payload)
            self.reply_queue.put(SocketReply(SocketReply.SUCCESS))
        except socket.error as e:
            self.reply_queue.put(SocketReply(SocketReply.ERROR, e))

    def _handle_RECEIVE(self, n):
        """
        Handles the receive command. This requires an open and valid socket
        connection.
        :param n: number of bytes to receive
        """
        try:
            received_data = receive_all(self.socket, n)
            self.reply_queue.put(
                SocketReply(SocketReply.SUCCESS, received_data))
        except socket.error as e:
            self.reply_queue.put(SocketReply(SocketReply.ERROR, e))
<commit_msg>Added a receive_with_prefix method<commit_after>"""
Socket Threaded.

Aims to provide some simplified abstraction over raw sockets such as making
concurrent socket connection easier.
"""

from __future__ import (
    absolute_import,
    division,
    print_function,
    unicode_literals
)
import socket
import struct

import sys
import threading

if sys.version_info.major == 2:
    chr = unichr
    string_type = basestring
    import Queue as queue
elif sys.version_info.major == 3:
    # chr should assume Unicode
    string_type = str
    import queue

def receive_all(sock, n):
    """
    Helper function to fully receive an arbitrary amount of data from a socket.

    :param sock: Socket connection
    :param n: Number of bytes to receive
    :return: Received data
    """
    data = b''
    while len(data) < n:
        packet = sock.recv(n - len(data))
        if not packet:
            break
            #return None
        data += packet
    return data

class SocketCommand(object):
    """
    Command object for communicating with SocketThread.
    command must be either:             payload must be:

    SocketCommand.CONNECT               (host, port) tuple
    SocketCommand.SEND                  Binary data string
    SocketCommand.RECEIVE               Number of bytes to receive
    SocketCommand.RECEIVE_WITH_PREFIX   byte size of the prefix
    SocketCommand.CLOSE                 None
    """

    # The available socket commands:
    CONNECT = "connect"
    SEND = "send"
    RECEIVE = "receive"
    RECEIVE_WITH_PREFIX = "receive_with_prefix"
    CLOSE = "close"

    def __init__(self, command, payload=None):
        self.command = command
        self.payload = payload

class SocketReply(object):
    """
    Reply object for communicating with SocketThread.
    The reply must be either:       payload must be:

    SocketReply.ERROR               The Error object
    SocketReply.SUCCESS
    """
    # TODO: write description for SocketReply.SUCCESS above

    # The available socket replies
    SUCCESS = "success"
    ERROR = "error"

    def __init__(self, status, payload=None):
        self.status = status
        self.payload = payload

class SocketThread(threading.Thread):
    """
    Implements the threading.Thread interface
    """
    # TODO: should we keep track of what command corresponds to what reply?
    # TODO: What if message first sends a length prefix? implement this?

    def __init__(self):
        super(SocketThread, self).__init__()

        self.command_queue = queue.Queue()
        self.reply_queue = queue.Queue()
        self.socket = None

        self.alive = threading.Event()
        self.start()  # TODO: Should we really call this here?
        self.alive.set()

    def run(self):
        while self.alive.isSet():
            try:
                # Use a timeout value so we don't block forever
                cmd = self.command_queue.get(block=True, timeout=0.1)
                if cmd.command == SocketCommand.CONNECT:
                    address = cmd.payload
                    self._handle_CONNECT(address)
                elif cmd.command == SocketCommand.CLOSE:
                    self._handle_CLOSE()
                elif cmd.command == SocketCommand.SEND:
                    payload = cmd.payload
                    self._handle_SEND(payload)
                elif cmd.command == SocketCommand.RECEIVE:
                    number_of_bytes = cmd.payload
                    self._handle_RECEIVE(number_of_bytes)
                elif cmd.command == SocketCommand.RECEIVE_WITH_PREFIX:
                    prefix_size = cmd.payload
                    self._handle_RECEIVE_WITH_PREFIX(prefix_size)
                else:
                    # TODO: Handle invalid command
                    pass

                # TODO: should we call command_queue.task_done() here?
            except queue.Empty:
                continue

    def join(self, timeout=None):
        self.alive.clear()
        threading.Thread.join(self, timeout)

    def connect(self, address):
        """
        Connects the socket the the given address
        :param address: (host, port) Tuple
        """
        self.command_queue.put(SocketCommand(SocketCommand.CONNECT, address))

    def close(self):
        """
        Closes the socket
        """
        self.command_queue.put(SocketCommand(SocketCommand.CONNECT))

    def send(self, payload):
        """
        Sends the given payload to the socket. Requires an open and valid socket
        :param payload: Byte string of data
        """
        self.command_queue.put(SocketCommand(SocketCommand.SEND, payload))

    def receive(self, n):
        """
        Receives a specified number of bytes from the socket. Requires an open
        and valid socket.
        :param n: Number of bytes to receive from the socket.
        """
        self.command_queue.put(SocketCommand(SocketCommand.RECEIVE, n))

    def receive_with_prefix(self, prefix_size):
        """
        Receive a message that has a length prefix that is prefix_size bytes
        long.
        :param prefix_size: byte size of the prefix
        """
        self.command_queue.put(SocketCommand(SocketCommand.RECEIVE_WITH_PREFIX,
                                             prefix_size))

    def get_reply(self, block=True, timeout=None):
        """
        Fetches a reply from the reply queue.
        :param block: Boolean whether the call should block until a reply i
        present
        :param timeout: If block is true wait maximum timeout number of seconds
        :return: A SocketReply object is returned if possible. Else if returns
        None.
        """
        try:
            reply = self.reply_queue.get(block=block, timeout=timeout)
            # TODO: should we call reply_queue.task_done() here?
            return reply
        except queue.Empty:
            return None

    def _handle_CONNECT(self, address):
        """
        Handles the connection command by creating a tcp connection to the given
        address.

        :param address: (host, port) tuple
        """
        self.socket = socket.socket(socket.AF_INET,
                                    socket.SOCK_STREAM)  # tcp connection

        try:
            self.socket.connect(address)
            self.reply_queue.put(SocketReply(SocketReply.SUCCESS))
        except socket.error as e:
            self.reply_queue.put(SocketReply(SocketReply.ERROR, e))

    def _handle_CLOSE(self):
        """
        Handles the close command. This requires that a socket already has been
        opened.
        """
        self.socket.close()
        self.reply_queue.put(SocketReply(SocketReply.SUCCESS))

    def _handle_SEND(self, payload):
        """
        Handles the send command. This requires an open and valid socket
        connection.
        :param payload:
        """
        try:
            self.socket.sendall(payload)
            self.reply_queue.put(SocketReply(SocketReply.SUCCESS))
        except socket.error as e:
            self.reply_queue.put(SocketReply(SocketReply.ERROR, e))

    def _handle_RECEIVE(self, n):
        """
        Handles the receive command. This requires an open and valid socket
        connection.
        :param n: number of bytes to receive
        """
        try:
            received_data = receive_all(self.socket, n)
            self.reply_queue.put(
                SocketReply(SocketReply.SUCCESS, received_data))
        except socket.error as e:
            self.reply_queue.put(SocketReply(SocketReply.ERROR, e))

    def _handle_RECEIVE_WITH_PREFIX(self, prefix_size):
        """
        Handles the receive with prefix command. This require an open and valid
        socket connection.
        :param prefix_size: byte size of the prefix
        :return:
        """
        try:
            length_prefix = receive_all(self.socket, prefix_size)
            if len(length_prefix) == prefix_size:
                if prefix_size == 1:
                    message_length = struct.unpack(b"!B", length_prefix)[0]
                elif prefix_size == 2:
                    message_length = struct.unpack(b"!H", length_prefix)[0]
                elif prefix_size == 4:
                    message_length = struct.unpack(b"!L", length_prefix)[0]
                elif prefix_size == 8:
                    message_length = struct.unpack(b"!Q", length_prefix)[0]
                else:
                    message_length = struct.unpack(b"!P", length_prefix)[0]

                received_data = receive_all(self.socket, message_length)
                if len(received_data) == message_length:
                    message = (message_length, received_data)
                    self.reply_queue.put(SocketReply(SocketReply.SUCCESS,
                                                     message))
                    return
            self.reply_queue.put(SocketReply(
                SocketReply.ERROR, socket.error("Socket closed prematurely")))
        except socket.error as e:
            self.reply_queue.put(SocketReply(SocketReply.ERROR, e))

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import math
from collections import *
from itertools import *
import sys

Point = namedtuple("Point", field_names=['a', 'b'])

def distance(x, y):
    return int(math.ceil(math.sqrt((x.a - y.a)**2 + (x.b - y.b)**2)))

def nline(f):
    val = f.readline().strip()
    return val

def choose_drone(order, order_index, drone_positions, drone_turns, sim_deadline, max_turns):
    index = max(drone_positions.keys(), key=lambda d: distance(order['pos'], drone_positions[d]))
    # index = order_index % len(drone_positions)
    if drone_turns[index] + 2 * max_turns < sim_deadline:
        return index
    else:
        for index in range(len(drone_positions)):
            if drone_turns[index] + 2 * max_turns < sim_deadline:
                return index
    return None

def main(f):
    rows, columns, drone_quantity, sim_deadline, drone_max_load = map(int, nline(f).split())
    max_turns = distance(Point(0, 0), Point(rows, columns)) + 1
    nline(f)
    products = dict(enumerate(map(int, nline(f).split())))
    num_warehouses = int(nline(f))
    warehouses = defaultdict(dict)
    for index in range(num_warehouses):
        warehouses[index]['pos'] = Point(*map(int, nline(f).split()))
        warehouses[index]['products'] = list(map(int, nline(f).split()))
    num_orders = int(nline(f))
    orders = defaultdict(dict)
    for index in range(num_orders):
        orders[index]['pos'] = Point(*map(int, nline(f).split()))
        nline(f)
        orders[index]['products'] = list(sorted(map(int, nline(f).split())))
        orders[index]['product_qty'] = defaultdict(int)
        for p in orders[index]['products']:
            orders[index]['product_qty'][p] += 1

    retired_drones = set()
    drone_positions = {x: warehouses[0]['pos'] for x in range(drone_quantity)}
    drone_turns = {x: 0 for x in range(drone_quantity)}
    moves = []
    for order_index, order in sorted(orders.items(), key=lambda o: len(o[1]['products'])):
        for order_product, qty in order['product_qty'].items():
            drone = choose_drone(order, order_index, drone_positions, drone_turns, sim_deadline, max_turns)
            potential_warehouses = [i for i, w in warehouses.items() if w['products'][order_product]]
            w = min(potential_warehouses, key=lambda i: distance(order['pos'], warehouses[i]['pos']))
            if drone is None:
                break
            max_qty = drone_max_load // products[order_product]
            qty_left = min(max_qty, qty)

            while qty_left > 0:
                potential_warehouses = [i for i, w in warehouses.items() if w['products'][order_product]]
                w = min(potential_warehouses,
                        key=lambda i: distance(drone_positions[drone], warehouses[i]['pos']) + \
                        distance(warehouses[i]['pos'], order['pos']))
                quantity = min(qty_left, warehouses[w]['products'][order_product])
                qty_left -= quantity

                moves.append("{} L {} {} {}".format(drone, w, order_product, quantity))
                drone_turns[drone] += 1 + distance(drone_positions[drone], warehouses[w]['pos'])
                drone_positions[drone] = warehouses[w]['pos']
                warehouses[w]['products'][order_product] -= quantity
                moves.append("{} D {} {} {}".format(drone, order_index, order_product, quantity))
                drone_turns[drone] += 1 + distance(drone_positions[drone], order['pos'])
                drone_positions[drone] = order['pos']
    print(len(moves))
    for move in moves:
        print(move)

if __name__ == '__main__':
    main(open(sys.argv[1]))
<commit_msg>WIP: sort by shortest distance from warehouse to delivery, then assign drone<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import math
from collections import *
from itertools import *
import sys

Point = namedtuple("Point", field_names=['a', 'b'])

class Move:
    def __init__(self, load, deliver):
        self.load = load
        self.deliver = deliver

class Load:
    def __init__(self, drone, warehouse, warehouse_index, product_type, quantity):
        self.drone = drone
        self.warehouse = warehouse
        self.warehouse_index = warehouse_index
        self.product_type = product_type
        self.quantity = quantity

    def __str__(self):
        return "{} L {} {} {}".format(self.drone, self.warehouse, self.product_type, self.quantity)

class Deliver:
    def __init__(self, drone, order, order_index, product_type, quantity):
        self.drone = drone
        self.order = order
        self.order_index = order_index
        self.product_type = product_type
        self.quantity = quantity

    def __str__(self):
        return "{} D {} {} {}".format(self.drone, self.order, self.product_type, self.quantity)

def distance(x, y):
    return int(math.ceil(math.sqrt((x.a - y.a)**2 + (x.b - y.b)**2)))

def nline(f):
    val = f.readline().strip()
    return val

def choose_drone(order, order_index, drone_positions, drone_turns, sim_deadline, max_turns):
    index = max(drone_positions.keys(), key=lambda d: distance(order['pos'], drone_positions[d]))
    # index = order_index % len(drone_positions)
    if drone_turns[index] + 2 * max_turns < sim_deadline:
        return index
    else:
        for index in range(len(drone_positions)):
            if drone_turns[index] + 2 * max_turns < sim_deadline:
                return index
    return None

def main(f):
    rows, columns, drone_quantity, sim_deadline, drone_max_load = map(int, nline(f).split())
    max_turns = distance(Point(0, 0), Point(rows, columns)) + 1
    nline(f)
    products = dict(enumerate(map(int, nline(f).split())))
    num_warehouses = int(nline(f))
    warehouses = defaultdict(dict)
    for index in range(num_warehouses):
        warehouses[index]['pos'] = Point(*map(int, nline(f).split()))
        warehouses[index]['products'] = list(map(int, nline(f).split()))
    num_orders = int(nline(f))
    orders = defaultdict(dict)
    for index in range(num_orders):
        orders[index]['pos'] = Point(*map(int, nline(f).split()))
        nline(f)
        orders[index]['products'] = list(sorted(map(int, nline(f).split())))
        orders[index]['product_qty'] = defaultdict(int)
        for p in orders[index]['products']:
            orders[index]['product_qty'][p] += 1

    retired_drones = set()
    drone_positions = {x: warehouses[0]['pos'] for x in range(drone_quantity)}
    drone_turns = {x: 0 for x in range(drone_quantity)}
    moves = []
    for order_index, order in sorted(orders.items(), key=lambda o: len(o[1]['products'])):
        for order_product, qty in order['product_qty'].items():
            drone = choose_drone(order, order_index, drone_positions, drone_turns, sim_deadline, max_turns)
            potential_warehouses = [i for i, w in warehouses.items() if w['products'][order_product]]
            w = min(potential_warehouses, key=lambda i: distance(order['pos'], warehouses[i]['pos']))
            if drone is None:
                break
            max_qty = drone_max_load // products[order_product]
            qty_left = min(max_qty, qty)

            while qty_left > 0:
                potential_warehouses = [i for i, w in warehouses.items() if w['products'][order_product]]
                w = min(potential_warehouses,
                        key=lambda i: distance(drone_positions[drone], warehouses[i]['pos']) + \
                        distance(warehouses[i]['pos'], order['pos']))
                quantity = min(qty_left, warehouses[w]['products'][order_product])
                qty_left -= quantity

                moves.append(Move(
                    Load(drone, warehouses[w], w, order_product, quantity),
                    Deliver(drone, order, order_index, order_product, quantity)
                ))

                drone_turns[drone] += 1 + distance(drone_positions[drone], warehouses[w]['pos'])
                drone_positions[drone] = warehouses[w]['pos']
                warehouses[w]['products'][order_product] -= quantity
                drone_turns[drone] += 1 + distance(drone_positions[drone], order['pos'])
                drone_positions[drone] = order['pos']

    moves_per_order = defaultdict(list)
    for m in moves:
        moves_per_order[m.deliver.order_index].append(m)
    for i, move_list in moves_per_order.items():
        move_list.sort(key=lambda m: distance(m.deliver.order['pos'], m.load.warehouse['pos']))

    order_moves = []
    for m in moves:
        # Assign drones
        pass

    print(len(moves))
    for move in moves:
        print(move.load)
        print(move.deliver)

if __name__ == '__main__':
    main(open(sys.argv[1]))
<|endoftext|>


<commit_before>from django.db import models

# Create your models here.
<commit_msg>Create models for Genre and Song<commit_after>from django.db import models

class Genre(models.Model):
	name = models.CharField(max_length=50)

	def __unicode__(self):
		return self.name

class Song(models.Model):
	title = models.CharField(max_length=100)
	artist = models.CharField(max_length=100)
	year = models.DateField()
	genre = models.ForeignKey(Genre)
	up = models.PositiveIntegerField(default=0)
	down = models.PositiveIntegerField(default=0)
	score = models.IntegerField(default=0)

	def __unicode__(self):
		return self.title
<|endoftext|>


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

"""

Takes MESA grids to create isochrones and organizes all the output files
into a nice directory structure.

The directory structure is as follows:
    top level directory --> MIST_vXX/FIDUCIAL/feh_pX.X_afe_pX.X/
    four subdirectories --> tracks/    eeps/    inlists/    isochrones/

Args:
    runname: the name of the grid

"""

import glob
import os
import sys
import csv
import reformat_massname
import subprocess
import datetime

work_dir = os.environ['MESAWORK_DIR']

if __name__ == "__main__":

    runname = sys.argv[1]
    do_organize(runname)

def gen_summary(rawdirname):

    """

    Retrieves various information about the MESA run and writes to a summary file

    Args:
        rawdirname: the name of the grid with the suffix '_raw'

    Outputs:
        None

    """

    #Outputs from the cluster
    listerrfiles = glob.glob(os.path.join(work_dir, rawdirname) + '/*/*.e')
    listoutfiles = glob.glob(os.path.join(work_dir, rawdirname) + '/*/*.o')

    #Dictionary to store the information about the MESA run
    stat_summary = {}

    #Loop over each model
    for index, file in enumerate(listerrfiles):

        #Declare status and also initialize each iteration
        status = ''

        #Extract the mass of the model
        if 'M_dir' in file:
            mass = file.split("/")[-2].rstrip('M_dir/')
        else:
            mass = file.split("/")[-2].split('M_')[0] + '_' + file.split("/")[-2].split('M_')[1].rstrip('_dir')

        with open(file, 'r') as errfile:
            errcontent = errfile.readlines()
        with open(listoutfiles[index], 'r') as outfile:
            outcontent = outfile.readlines()

        status = ''
        termination_reason = ''

        #Check for error messages
        if (len(errcontent) >= 1):
            if 'DUE TO TIME LIMIT ***' in errcontent[1]:
                status = 'FAILED'
                reason = 'need_more_time'
            else:
                status = 'FAILED'
                reason = 'unknown_error'

        #Retrieve the stopping reasons
        else:
            for line in outcontent[-30:]:
                if 'termination code' in line:
                    termination_reason = line.split('termination code: ')[1].split('\n')[0]
                    break
                if 'failed in do_relax_num_steps' in line:
                    termination_reason = 'failed_during_preMS'
                    break
            for line in outcontent[-50:]:
                if (' stopping because of convergence problems' in line) or \
                           ('terminated evolution: convergence problems' in line):
                    status = 'FAILED'
                    reason = termination_reason.replace(' ', '_')
                if (line == outcontent[-1]) & (status == ''):
                    status = 'OK'
                    reason = termination_reason.replace(' ', '_')

        #Retrieve the run time information
        dates = subprocess.Popen('grep [0-9][0-9]:[0-9][0-9]:[0-9][0-9] ' + listoutfiles[index], shell=True, stdout=subprocess.PIPE)
        try:
            startdate, enddate = dates.stdout
            startdatelist = startdate.rstrip('\n').split(' ')
            enddatelist = enddate.rstrip('\n').split(' ')

            #For single-digit dates
            if '' in startdatelist:
                startdatelist.remove('')
            if '' in enddatelist:
                enddatelist.remove('')

            start = datetime.timedelta(int(startdatelist[2]), int(startdatelist[3].split(':')[-1]), 0,0,int(startdatelist[3].split(':')[-2]), int(startdatelist[3].split(':')[-3]))
            end = datetime.timedelta(int(enddatelist[2]), int(enddatelist[3].split(':')[-1]), 0,0,int(enddatelist[3].split(':')[-2]), int(enddatelist[3].split(':')[-3]))
            runtime = str(datetime.timedelta(seconds=datetime.timedelta.total_seconds(end-start)))

        #If there is no end date
        except ValueError:
            runtime = 'exceeded_req_time'

        #Populate the stat_summary dictionary
        stat_summary[mass] = "{:10}".format(status) + "{:50}".format(reason) + "{:25}".format(runtime)

    keys = stat_summary.keys()
    #Sort by mass in ascending order
    keys.sort()

    #Write to a file
    summary_filename = "tracks_summary.txt"
    f = csv.writer(open(summary_filename, 'w'), delimiter='\t')
    f.writerow(["{:15}".format('#Mass'), "{:10}".format('Status') + "{:50}".format('Reason') + "{:25}".format('Runtime')])
    f.writerow(['','','',''])

    for key in keys:
        f.writerow(["{:15}".format(key), stat_summary[key]])

def sort_histfiles(rawdirname):

    """

    Organizes the history files.

    Args:
        rawdirname: the name of the grid with the suffix '_raw'

    Outputs:
        None

    """

    #Get the list of history files (tracks)
    listofhist = glob.glob(os.path.join(work_dir, os.path.join(rawdirname+'/*/LOGS/*.data')))

    #Make the track directory in the new reduced MESA run directory
    new_parentdirname = rawdirname.split("_raw")[0]
    os.system("cd " + new_parentdirname)
    histfiles_dirname = os.path.join(new_parentdirname, "tracks")
    os.mkdir(histfiles_dirname)

    #Rename & copy the history files over
    for histfile in listofhist:
        if 'M_history.data' in histfile:
            unformat_mass_string = histfile.split('LOGS/')[1].split('M_history.data')[0]
            newhistfilename = histfile.split('LOGS')[0]+'LOGS/'+reformat_massname.reformat_massname(unformat_mass_string)+'M.track'
        else:
            unformat_mass_string = histfile.split('LOGS/')[1].split('_history.data')[0].split('M_')[0]
            bc_name = histfile.split('LOGS/')[1].split('M_')[1].split('_history.data')[0]
            newhistfilename = histfile.split('LOGS')[0]+'LOGS/'+reformat_massname.reformat_massname(unformat_mass_string)+'M_' + bc_name + '.track'
        os.system("cp " + histfile + " " + newhistfilename)
        os.system("mv " + newhistfilename + " " + histfiles_dirname)

def save_inlists(rawdirname):

    """

    Organizes the inlist files.

    Args:
        rawdirname: the name of the grid with the suffix '_raw'

    Outputs:
        None

    """

    #Nake the inlist directory in the new reduced MESA run directory
    new_parentdirname = rawdirname.split("_raw")[0]
    os.system("cd " + new_parentdirname)
    inlistfiles_dirname = os.path.join(new_parentdirname, "inlists")
    os.mkdir(inlistfiles_dirname)

    #Copy the inlist files from the general inlist directory in MESAWORK_DIR to the newly created inlist directory
    os.system("cp " + os.path.join(work_dir, "inlists/inlists_"+new_parentdirname+"/*") + " " + inlistfiles_dirname)

def do_organize(runname):

    """

    Wrapper for the various routines to reduce a directory of MESA runs to
    a MIST directory.

    Args:
        runname: the name of the grid

    Outputs:
        None

    """

    #Rename the run directory XXX as XXX_raw
    rawdirname = runname+"_raw"
    os.system("mv " + os.path.join(work_dir,runname) + " " + os.path.join(work_dir,rawdirname))

    #The XXX directory will contain the organized, reduced information
    os.mkdir(runname)

    #Make the eeps directory that will be filled in later
    os.mkdir(os.path.join(runname, "eeps"))

    #Make the isochrones directory that will be filled in later
    os.mkdir(os.path.join(runname, "isochrones"))

    print "************************************************************"
    print "****************SORTING THE HISTORY FILES*******************"
    print "************************************************************"
    sort_histfiles(rawdirname)

    print "************************************************************"
    print "****************GENERATING A SUMMARY FILE*******************"
    print "************************************************************"
    gen_summary(rawdirname)

    #Move the summary file to the tracks directory
    os.system("mv tracks_summary.txt " + os.path.join(runname, "tracks"))

    print "************************************************************"
    print "****************SORTING THE INLIST FILES********************"
    print "************************************************************"
    save_inlists(rawdirname)

    print "************************************************************"
    print "**********************MAKE ISOCHRONES***********************"
    print "************************************************************"
    mesa2fsp(runname)

    print "************************************************************"
    print "****************COMPRESSING THE DIRECTORY*******************"
    print "************************************************************"
    os.system("tar -zcvf " + runname + ".tar.gz " + runname)
    os.system("mv " + runname + "* " + work_dir)

<commit_msg>can now recognize errors due to memory being exceeded<commit_after>#!/usr/bin/env python

"""

Takes MESA grids to create isochrones and organizes all the output files
into a nice directory structure.

The directory structure is as follows:
    top level directory --> MIST_vXX/FIDUCIAL/feh_pX.X_afe_pX.X/
    four subdirectories --> tracks/    eeps/    inlists/    isochrones/

Args:
    runname: the name of the grid

"""

import glob
import os
import sys
import csv
import reformat_massname
import subprocess
import datetime

work_dir = os.environ['MESAWORK_DIR']

if __name__ == "__main__":

    runname = sys.argv[1]
    do_organize(runname)

def gen_summary(rawdirname):

    """

    Retrieves various information about the MESA run and writes to a summary file

    Args:
        rawdirname: the name of the grid with the suffix '_raw'

    Outputs:
        None

    """

    #Outputs from the cluster
    listerrfiles = glob.glob(os.path.join(work_dir, rawdirname) + '/*/*.e')
    listoutfiles = glob.glob(os.path.join(work_dir, rawdirname) + '/*/*.o')

    #Dictionary to store the information about the MESA run
    stat_summary = {}

    #Loop over each model
    for index, file in enumerate(listerrfiles):

        #Declare status and also initialize each iteration
        status = ''

        #Extract the mass of the model
        if 'M_dir' in file:
            mass = file.split("/")[-2].rstrip('M_dir/')
        else:
            mass = file.split("/")[-2].split('M_')[0] + '_' + file.split("/")[-2].split('M_')[1].rstrip('_dir')

        with open(file, 'r') as errfile:
            errcontent = errfile.readlines()
        with open(listoutfiles[index], 'r') as outfile:
            outcontent = outfile.readlines()

        status = ''
        termination_reason = ''

        #Check for error messages
        if (len(errcontent) > 0):
            for line in errcontent:
                if 'DUE TO TIME LIMIT ***' in line:
                    status = 'FAILED'
                    reason = 'need_more_time'
                    break
                elif 'exceeded memory limit' in line:
                    status = 'FAILED'
                    reason = 'memory_exceed'
                    break
                else:
                    status = 'FAILED'
                    reason = 'unknown_error'
                    break

        #Retrieve the stopping reasons
        else:
            for line in outcontent[-30:]:
                if 'termination code' in line:
                    termination_reason = line.split('termination code: ')[1].split('\n')[0]
                    break
                if 'failed in do_relax_num_steps' in line:
                    termination_reason = 'failed_during_preMS'
                    break
            for line in outcontent[-50:]:
                if (' stopping because of convergence problems' in line) or \
                           ('terminated evolution: convergence problems' in line):
                    status = 'FAILED'
                    reason = termination_reason.replace(' ', '_')
                if (line == outcontent[-1]) & (status == ''):
                    status = 'OK'
                    reason = termination_reason.replace(' ', '_')

        #Retrieve the run time information
        dates = subprocess.Popen('grep [0-9][0-9]:[0-9][0-9]:[0-9][0-9] ' + listoutfiles[index], shell=True, stdout=subprocess.PIPE)
        try:
            startdate, enddate = dates.stdout
            startdatelist = startdate.rstrip('\n').split(' ')
            enddatelist = enddate.rstrip('\n').split(' ')

            #For single-digit dates
            if '' in startdatelist:
                startdatelist.remove('')
            if '' in enddatelist:
                enddatelist.remove('')

            start = datetime.timedelta(int(startdatelist[2]), int(startdatelist[3].split(':')[-1]), 0,0,int(startdatelist[3].split(':')[-2]), int(startdatelist[3].split(':')[-3]))
            end = datetime.timedelta(int(enddatelist[2]), int(enddatelist[3].split(':')[-1]), 0,0,int(enddatelist[3].split(':')[-2]), int(enddatelist[3].split(':')[-3]))
            runtime = str(datetime.timedelta(seconds=datetime.timedelta.total_seconds(end-start)))

        #If there is no end date
        except ValueError:
            runtime = 'exceeded_req_time'

        #Populate the stat_summary dictionary
        stat_summary[mass] = "{:10}".format(status) + "{:50}".format(reason) + "{:25}".format(runtime)

    keys = stat_summary.keys()
    #Sort by mass in ascending order
    keys.sort()

    #Write to a file
    summary_filename = "tracks_summary.txt"
    f = csv.writer(open(summary_filename, 'w'), delimiter='\t')
    f.writerow(["{:15}".format('#Mass'), "{:10}".format('Status') + "{:50}".format('Reason') + "{:25}".format('Runtime')])
    f.writerow(['','','',''])

    for key in keys:
        f.writerow(["{:15}".format(key), stat_summary[key]])

def sort_histfiles(rawdirname):

    """

    Organizes the history files.

    Args:
        rawdirname: the name of the grid with the suffix '_raw'

    Outputs:
        None

    """

    #Get the list of history files (tracks)
    listofhist = glob.glob(os.path.join(work_dir, os.path.join(rawdirname+'/*/LOGS/*.data')))

    #Make the track directory in the new reduced MESA run directory
    new_parentdirname = rawdirname.split("_raw")[0]
    os.system("cd " + new_parentdirname)
    histfiles_dirname = os.path.join(new_parentdirname, "tracks")
    os.mkdir(histfiles_dirname)

    #Rename & copy the history files over
    for histfile in listofhist:
        if 'M_history.data' in histfile:
            unformat_mass_string = histfile.split('LOGS/')[1].split('M_history.data')[0]
            newhistfilename = histfile.split('LOGS')[0]+'LOGS/'+reformat_massname.reformat_massname(unformat_mass_string)+'M.track'
        else:
            unformat_mass_string = histfile.split('LOGS/')[1].split('_history.data')[0].split('M_')[0]
            bc_name = histfile.split('LOGS/')[1].split('M_')[1].split('_history.data')[0]
            newhistfilename = histfile.split('LOGS')[0]+'LOGS/'+reformat_massname.reformat_massname(unformat_mass_string)+'M_' + bc_name + '.track'
        os.system("cp " + histfile + " " + newhistfilename)
        os.system("mv " + newhistfilename + " " + histfiles_dirname)

def save_inlists(rawdirname):

    """

    Organizes the inlist files.

    Args:
        rawdirname: the name of the grid with the suffix '_raw'

    Outputs:
        None

    """

    #Nake the inlist directory in the new reduced MESA run directory
    new_parentdirname = rawdirname.split("_raw")[0]
    os.system("cd " + new_parentdirname)
    inlistfiles_dirname = os.path.join(new_parentdirname, "inlists")
    os.mkdir(inlistfiles_dirname)

    #Copy the inlist files from the general inlist directory in MESAWORK_DIR to the newly created inlist directory
    os.system("cp " + os.path.join(work_dir, "inlists/inlists_"+new_parentdirname+"/*") + " " + inlistfiles_dirname)

def do_organize(runname):

    """

    Wrapper for the various routines to reduce a directory of MESA runs to
    a MIST directory.

    Args:
        runname: the name of the grid

    Outputs:
        None

    """

    #Rename the run directory XXX as XXX_raw
    rawdirname = runname+"_raw"
    os.system("mv " + os.path.join(work_dir,runname) + " " + os.path.join(work_dir,rawdirname))

    #The XXX directory will contain the organized, reduced information
    os.mkdir(runname)

    #Make the eeps directory that will be filled in later
    os.mkdir(os.path.join(runname, "eeps"))

    #Make the isochrones directory that will be filled in later
    os.mkdir(os.path.join(runname, "isochrones"))

    print "************************************************************"
    print "****************SORTING THE HISTORY FILES*******************"
    print "************************************************************"
    sort_histfiles(rawdirname)

    print "************************************************************"
    print "****************GENERATING A SUMMARY FILE*******************"
    print "************************************************************"
    gen_summary(rawdirname)

    #Move the summary file to the tracks directory
    os.system("mv tracks_summary.txt " + os.path.join(runname, "tracks"))

    print "************************************************************"
    print "****************SORTING THE INLIST FILES********************"
    print "************************************************************"
    save_inlists(rawdirname)

    print "************************************************************"
    print "**********************MAKE ISOCHRONES***********************"
    print "************************************************************"
    mesa2fsp(runname)

    print "************************************************************"
    print "****************COMPRESSING THE DIRECTORY*******************"
    print "************************************************************"
    os.system("tar -zcvf " + runname + ".tar.gz " + runname)
    os.system("mv " + runname + "* " + work_dir)

<|endoftext|>