<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Read and save event lists from FITS files."""
from __future__ import (absolute_import, unicode_literals, division,
                        print_function)

from .base import mp_root, read_header_key, ref_mjd
from .io import save_events
from .io import MP_FILE_EXTENSION
import numpy as np
import logging
import os

def load_gtis(fits_file, gtistring=None):
    """Load GTI from HDU EVENTS of file fits_file."""
    from astropy.io import fits as pf
    import numpy as np

    if gtistring is None:
        gtistring = 'GTI'
    logging.info("Loading GTIS from file %s" % fits_file)
    lchdulist = pf.open(fits_file, checksum=True)
    lchdulist.verify('warn')

    gtitable = lchdulist[gtistring].data
    gti_list = np.array([[a, b]
                         for a, b in zip(gtitable.field('START'),
                                         gtitable.field('STOP'))],
                        dtype=np.longdouble)
    lchdulist.close()
    return gti_list

def load_events_and_gtis(fits_file, return_limits=False,
                         additional_columns=None, gtistring=None,
                         gti_file=None, hduname='EVENTS', column='TIME'):
    """Load event lists and GTIs from one or more files.

    Loads event list from HDU EVENTS of file fits_file, with Good Time
    intervals. Optionally, returns additional columns of data from the same
    HDU of the events.

    Parameters
    ----------
    fits_file : str
    return_limits: bool, optional
        Return the TSTART and TSTOP keyword values
    additional_columns: list of str, optional
        A list of keys corresponding to the additional columns to extract from
        the event HDU (ex.: ['PI', 'X'])

    Returns
    -------
    ev_list : array-like
    gtis: [[gti0_0, gti0_1], [gti1_0, gti1_1], ...]
    additional_data: dict
        A dictionary, where each key is the one specified in additional_colums.
        The data are an array with the values of the specified column in the
        fits file.
    t_start : float
    t_stop : float
    """
    from astropy.io import fits as pf

    lchdulist = pf.open(fits_file)

    hdunames = [h.name for h in lchdulist]

    # Load data table
    try:
        lctable = lchdulist[hduname].data
    except:  # pragma: no cover
        logging.warning('HDU %s not found. Trying first extension' % hduname)
        lctable = lchdulist[1].data

    # Read event list
    ev_list = np.array(lctable.field(column), dtype=np.longdouble)

    # Read TIMEZERO keyword and apply it to events
    try:
        timezero = np.longdouble(lchdulist[1].header['TIMEZERO'])
    except:  # pragma: no cover
        logging.warning("TIMEZERO is 0")
        timezero = np.longdouble(0.)

    if timezero != 0.:
        logging.warning("TIMEZERO != 0, correcting")
        ev_list += timezero

    # Read TSTART, TSTOP from header
    try:
        t_start = np.longdouble(lchdulist[1].header['TSTART'])
        t_stop = np.longdouble(lchdulist[1].header['TSTOP'])
    except:  # pragma: no cover
        logging.warning("Tstart and Tstop error. using defaults")
        t_start = ev_list[0]
        t_stop = ev_list[-1]

    # Read and handle GTI extension
    if gtistring is None:
        accepted_gtistrings = ['GTI', 'STDGTI']
    else:
        accepted_gtistrings = [gtistring]

    if gti_file is None:
        # Select first GTI with accepted name
        try:
            gtiextn = [ix for ix, x in enumerate(hdunames)
                       if x in accepted_gtistrings][0]
            gtiext = lchdulist[gtiextn]
            gtitable = gtiext.data

            colnames = [col.name for col in gtitable.columns.columns]
            # Default: NuSTAR: START, STOP. Otherwise, try RXTE: Start, Stop
            if 'START' in colnames:
                startstr, stopstr = 'START', 'STOP'
            else:
                startstr, stopstr = 'Start', 'Stop'

            gtistart = np.array(gtitable.field(startstr), dtype=np.longdouble)
            gtistop = np.array(gtitable.field(stopstr), dtype=np.longdouble)
            gti_list = np.array([[a, b]
                                 for a, b in zip(gtistart,
                                                 gtistop)],
                                dtype=np.longdouble)

        except:  # pragma: no cover
            logging.warning("%s Extension not found or invalid " % gtistring +
                            "in %s!! Please check!!" % fits_file)
            gti_list = np.array([[t_start, t_stop]],
                                dtype=np.longdouble)
    else:
        gti_list = load_gtis(gti_file, gtistring)

    if additional_columns is not None:
        additional_data = {}
        for a in additional_columns:
            try:
                additional_data[a] = np.array(lctable.field(a))
            except:  # pragma: no cover
                if a == 'PI':
                    logging.warning('Column PI not found. Trying with PHA')
                    additional_data[a] = np.array(lctable.field('PHA'))
                else:
                    raise Exception('Column' + a + 'not found')

    lchdulist.close()

    if return_limits:
        if additional_columns is not None:
            return ev_list, gti_list, additional_data, t_start, t_stop
        else:
            return ev_list, gti_list, t_start, t_stop
    else:
        if additional_columns is not None:
            return ev_list, gti_list, additional_data
        else:
            return ev_list, gti_list

def treat_event_file(filename, noclobber=False, gti_split=False,
                     min_length=4):
    """Read data from an event file, with no external GTI information.

    Parameters
    ----------
    filename : str

    Other Parameters
    ----------------
    noclobber: bool
        if a file is present, do not overwrite it
    gti_split: bool
        split the file in multiple chunks, containing one GTI each
    min_length: float
        minimum length of GTIs accepted (only if gti_split is True)
    """
    logging.info('Opening %s' % filename)
    outfile = mp_root(filename) + '_ev' + MP_FILE_EXTENSION
    if noclobber and os.path.exists(outfile) and (not gti_split):
        print(
            '{0} exists, and noclobber option used. Skipping'.format(outfile))
        return

    instr = read_header_key(filename, 'INSTRUME')
    additional_columns = ['PI']
    if instr == 'PCA':
        additional_columns.append('PCUID')

    mjdref = ref_mjd(filename)
    events, gtis, additional, tstart, tstop = \
        load_events_and_gtis(filename,
                             additional_columns=additional_columns,
                             return_limits=True)

    pis = additional['PI']
    out = {'time': events,
           'GTI': gtis,
           'PI': pis,
           'MJDref': mjdref,
           'Tstart': tstart,
           'Tstop': tstop,
           'Instr': instr
           }

    if instr == "PCA":
        out['PCU'] = np.array(additional['PCUID'], dtype=np.byte)

    if gti_split:
        for ig, g in enumerate(gtis):
            length = g[1] - g[0]
            if length < min_length:
                print("GTI shorter than {0} s; skipping".format(min_length))
                continue

            outfile_local = \
                '{0}_{1}'.format(outfile.replace(MP_FILE_EXTENSION,
                                                 ''), ig) + \
                MP_FILE_EXTENSION
            if noclobber and os.path.exists(outfile_local):
                print('{0} exists, '.format(outfile_local) +
                      'and noclobber option used. Skipping')
                return
            out_local = out.copy()
            good = np.logical_and(events >= g[0], events < g[1])
            if not np.any(good):
                print("This GTI has no valid events; skipping")
                continue
            out_local['time'] = events[good]
            out_local['Tstart'] = g[0]
            out_local['Tstop'] = g[1]
            out_local['PI'] = pis[good]
            out_local['GTI'] = np.array([g], dtype=np.longdouble)
            save_events(out_local, outfile_local)
        pass
    else:
        save_events(out, outfile)

def _wrap_fun(arglist):
    f, kwargs = arglist
    return treat_event_file(f, **kwargs)

def main(args=None):
    import argparse
    from multiprocessing import Pool

    description = ('Read a cleaned event files and saves the relevant '
                   'information in a standard format')
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("files", help="List of files", nargs='+')
    parser.add_argument("--loglevel",
                        help=("use given logging level (one between INFO, "
                              "WARNING, ERROR, CRITICAL, DEBUG; "
                              "default:WARNING)"),
                        default='WARNING',
                        type=str)
    parser.add_argument("--nproc",
                        help=("Number of processors to use"),
                        default=1,
                        type=int)
    parser.add_argument("--noclobber",
                        help=("Do not overwrite existing event files"),
                        default=False, action='store_true')
    parser.add_argument("-g", "--gti-split",
                        help="Split event list by GTI",
                        default=False,
                        action="store_true")
    parser.add_argument("--min-length", type=int,
                        help="Minimum length of GTIs to consider",
                        default=0)
    parser.add_argument("--debug", help="use DEBUG logging level",
                        default=False, action='store_true')

    args = parser.parse_args(args)
    files = args.files

    if args.debug:
        args.loglevel = 'DEBUG'

    numeric_level = getattr(logging, args.loglevel.upper(), None)
    logging.basicConfig(filename='MPreadevents.log', level=numeric_level,
                        filemode='w')

    argdict = {"noclobber": args.noclobber, "gti_split": args.gti_split,
               "min_length": args.min_length}
    arglist = [[f, argdict] for f in files]

    if os.name == 'nt' or args.nproc == 1:
        [_wrap_fun(a) for a in arglist]
    else:
        pool = Pool(processes=args.nproc)
        for i in pool.imap_unordered(_wrap_fun, arglist):
            pass
        pool.close()
<commit_msg>Sort event list<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Read and save event lists from FITS files."""
from __future__ import (absolute_import, unicode_literals, division,
                        print_function)

from .base import mp_root, read_header_key, ref_mjd
from .io import save_events
from .io import MP_FILE_EXTENSION
import numpy as np
import logging
import os

def load_gtis(fits_file, gtistring=None):
    """Load GTI from HDU EVENTS of file fits_file."""
    from astropy.io import fits as pf
    import numpy as np

    if gtistring is None:
        gtistring = 'GTI'
    logging.info("Loading GTIS from file %s" % fits_file)
    lchdulist = pf.open(fits_file, checksum=True)
    lchdulist.verify('warn')

    gtitable = lchdulist[gtistring].data
    gti_list = np.array([[a, b]
                         for a, b in zip(gtitable.field('START'),
                                         gtitable.field('STOP'))],
                        dtype=np.longdouble)
    lchdulist.close()
    return gti_list

def load_events_and_gtis(fits_file, return_limits=False,
                         additional_columns=None, gtistring=None,
                         gti_file=None, hduname='EVENTS', column='TIME'):
    """Load event lists and GTIs from one or more files.

    Loads event list from HDU EVENTS of file fits_file, with Good Time
    intervals. Optionally, returns additional columns of data from the same
    HDU of the events.

    Parameters
    ----------
    fits_file : str
    return_limits: bool, optional
        Return the TSTART and TSTOP keyword values
    additional_columns: list of str, optional
        A list of keys corresponding to the additional columns to extract from
        the event HDU (ex.: ['PI', 'X'])

    Returns
    -------
    ev_list : array-like
    gtis: [[gti0_0, gti0_1], [gti1_0, gti1_1], ...]
    additional_data: dict
        A dictionary, where each key is the one specified in additional_colums.
        The data are an array with the values of the specified column in the
        fits file.
    t_start : float
    t_stop : float
    """
    from astropy.io import fits as pf

    lchdulist = pf.open(fits_file)

    hdunames = [h.name for h in lchdulist]

    # Load data table
    try:
        lctable = lchdulist[hduname].data
    except:  # pragma: no cover
        logging.warning('HDU %s not found. Trying first extension' % hduname)
        lctable = lchdulist[1].data

    # Read event list
    ev_list = np.array(lctable.field(column), dtype=np.longdouble)

    # Read TIMEZERO keyword and apply it to events
    try:
        timezero = np.longdouble(lchdulist[1].header['TIMEZERO'])
    except:  # pragma: no cover
        logging.warning("TIMEZERO is 0")
        timezero = np.longdouble(0.)

    if timezero != 0.:
        logging.warning("TIMEZERO != 0, correcting")
        ev_list += timezero

    # Read TSTART, TSTOP from header
    try:
        t_start = np.longdouble(lchdulist[1].header['TSTART'])
        t_stop = np.longdouble(lchdulist[1].header['TSTOP'])
    except:  # pragma: no cover
        logging.warning("Tstart and Tstop error. using defaults")
        t_start = ev_list[0]
        t_stop = ev_list[-1]

    # Read and handle GTI extension
    if gtistring is None:
        accepted_gtistrings = ['GTI', 'STDGTI']
    else:
        accepted_gtistrings = [gtistring]

    if gti_file is None:
        # Select first GTI with accepted name
        try:
            gtiextn = [ix for ix, x in enumerate(hdunames)
                       if x in accepted_gtistrings][0]
            gtiext = lchdulist[gtiextn]
            gtitable = gtiext.data

            colnames = [col.name for col in gtitable.columns.columns]
            # Default: NuSTAR: START, STOP. Otherwise, try RXTE: Start, Stop
            if 'START' in colnames:
                startstr, stopstr = 'START', 'STOP'
            else:
                startstr, stopstr = 'Start', 'Stop'

            gtistart = np.array(gtitable.field(startstr), dtype=np.longdouble)
            gtistop = np.array(gtitable.field(stopstr), dtype=np.longdouble)
            gti_list = np.array([[a, b]
                                 for a, b in zip(gtistart,
                                                 gtistop)],
                                dtype=np.longdouble)

        except:  # pragma: no cover
            logging.warning("%s Extension not found or invalid " % gtistring +
                            "in %s!! Please check!!" % fits_file)
            gti_list = np.array([[t_start, t_stop]],
                                dtype=np.longdouble)
    else:
        gti_list = load_gtis(gti_file, gtistring)

    if additional_columns is not None:
        additional_data = {}
        for a in additional_columns:
            try:
                additional_data[a] = np.array(lctable.field(a))
            except:  # pragma: no cover
                if a == 'PI':
                    logging.warning('Column PI not found. Trying with PHA')
                    additional_data[a] = np.array(lctable.field('PHA'))
                else:
                    raise Exception('Column' + a + 'not found')

    lchdulist.close()

    # Sort event list
    order = np.argsort(ev_list)
    ev_list = ev_list[order]
    if additional_columns is not None:
        additional_data = dict([(k, additional_data[k][order])
                                for k in additional_columns])

    if return_limits:
        if additional_columns is not None:
            return ev_list, gti_list, additional_data, t_start, t_stop
        else:
            return ev_list, gti_list, t_start, t_stop
    else:
        if additional_columns is not None:
            return ev_list, gti_list, additional_data
        else:
            return ev_list, gti_list

def treat_event_file(filename, noclobber=False, gti_split=False,
                     min_length=4):
    """Read data from an event file, with no external GTI information.

    Parameters
    ----------
    filename : str

    Other Parameters
    ----------------
    noclobber: bool
        if a file is present, do not overwrite it
    gti_split: bool
        split the file in multiple chunks, containing one GTI each
    min_length: float
        minimum length of GTIs accepted (only if gti_split is True)
    """
    logging.info('Opening %s' % filename)
    outfile = mp_root(filename) + '_ev' + MP_FILE_EXTENSION
    if noclobber and os.path.exists(outfile) and (not gti_split):
        print(
            '{0} exists, and noclobber option used. Skipping'.format(outfile))
        return

    instr = read_header_key(filename, 'INSTRUME')
    additional_columns = ['PI']
    if instr == 'PCA':
        additional_columns.append('PCUID')

    mjdref = ref_mjd(filename)
    events, gtis, additional, tstart, tstop = \
        load_events_and_gtis(filename,
                             additional_columns=additional_columns,
                             return_limits=True)

    pis = additional['PI']
    out = {'time': events,
           'GTI': gtis,
           'PI': pis,
           'MJDref': mjdref,
           'Tstart': tstart,
           'Tstop': tstop,
           'Instr': instr
           }

    if instr == "PCA":
        out['PCU'] = np.array(additional['PCUID'], dtype=np.byte)

    if gti_split:
        for ig, g in enumerate(gtis):
            length = g[1] - g[0]
            if length < min_length:
                print("GTI shorter than {0} s; skipping".format(min_length))
                continue

            outfile_local = \
                '{0}_{1}'.format(outfile.replace(MP_FILE_EXTENSION,
                                                 ''), ig) + \
                MP_FILE_EXTENSION
            if noclobber and os.path.exists(outfile_local):
                print('{0} exists, '.format(outfile_local) +
                      'and noclobber option used. Skipping')
                return
            out_local = out.copy()
            good = np.logical_and(events >= g[0], events < g[1])
            if not np.any(good):
                print("This GTI has no valid events; skipping")
                continue
            out_local['time'] = events[good]
            out_local['Tstart'] = g[0]
            out_local['Tstop'] = g[1]
            out_local['PI'] = pis[good]
            out_local['GTI'] = np.array([g], dtype=np.longdouble)
            save_events(out_local, outfile_local)
        pass
    else:
        save_events(out, outfile)

def _wrap_fun(arglist):
    f, kwargs = arglist
    return treat_event_file(f, **kwargs)

def main(args=None):
    import argparse
    from multiprocessing import Pool

    description = ('Read a cleaned event files and saves the relevant '
                   'information in a standard format')
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("files", help="List of files", nargs='+')
    parser.add_argument("--loglevel",
                        help=("use given logging level (one between INFO, "
                              "WARNING, ERROR, CRITICAL, DEBUG; "
                              "default:WARNING)"),
                        default='WARNING',
                        type=str)
    parser.add_argument("--nproc",
                        help=("Number of processors to use"),
                        default=1,
                        type=int)
    parser.add_argument("--noclobber",
                        help=("Do not overwrite existing event files"),
                        default=False, action='store_true')
    parser.add_argument("-g", "--gti-split",
                        help="Split event list by GTI",
                        default=False,
                        action="store_true")
    parser.add_argument("--min-length", type=int,
                        help="Minimum length of GTIs to consider",
                        default=0)
    parser.add_argument("--debug", help="use DEBUG logging level",
                        default=False, action='store_true')

    args = parser.parse_args(args)
    files = args.files

    if args.debug:
        args.loglevel = 'DEBUG'

    numeric_level = getattr(logging, args.loglevel.upper(), None)
    logging.basicConfig(filename='MPreadevents.log', level=numeric_level,
                        filemode='w')

    argdict = {"noclobber": args.noclobber, "gti_split": args.gti_split,
               "min_length": args.min_length}
    arglist = [[f, argdict] for f in files]

    if os.name == 'nt' or args.nproc == 1:
        [_wrap_fun(a) for a in arglist]
    else:
        pool = Pool(processes=args.nproc)
        for i in pool.imap_unordered(_wrap_fun, arglist):
            pass
        pool.close()
<|endoftext|>


<commit_before>import os, os.path, logging, subprocess, threading, time, errno
import tempfile
from subprocess import PIPE
from contextlib import contextmanager

def ensure_dir(dir, logger=None) :
    if not os.path.exists(dir) :
        try:
            if logger: logger.debug("Making dir: %s", dir)
            os.makedirs(dir)
            if logger: logger.info("Created dir: %s", dir)
        except OSError, e :
            #not entirely sure why we are getting errors,
            #http://docs.python.org/library/os.path.html#os.path.exists
            #says this could be related to not being able to call
            #os.stat, but I can do so from the command line python
            #just fine.
            if os.path.exists(dir) :
                if logger: logger.debug("Erred, already exists: e.errno: %s",e.errno)
                return
            else:
                raise

def withDir(dir, fn, *args,**kwargs) :
    olddir = os.getcwd()
    try :
        os.chdir(dir)
        return fn(*args,**kwargs)
    finally :
        os.chdir(olddir)

def pprint_bytes(bytes) :
    suffix = 'B'
    bytes= float(bytes)
    if bytes >= 1024 :
        bytes = bytes / 1024
        suffix = 'KB'
    if bytes >= 1024 :
        bytes = bytes / 1024
        suffix = 'MB'
    if bytes >= 1024 :
        bytes = bytes / 1024
        suffix = 'GB'
    if bytes >= 1024 :
        bytes = bytes / 1024
        suffix = 'TB'
    return '%.2f%s' % (bytes,suffix)

def exec_external( proc ):
    return subprocess.call(['python', proc])

def capturedCall(cmd, **kwargs) :
    """Do the equivelent of the subprocess.call except
    log the stderr and stdout where appropriate."""
    p= capturedPopen(cmd,**kwargs)
    rc = p.wait()
    #this is a cheap attempt to make sure the monitors
    #are scheduled and hopefully finished.
    time.sleep(0.01)
    time.sleep(0.01)
    return rc

def capturedPopen(cmd, stdin=None, stdout=None, stderr=None,
                  logger=logging,
                  stdout_level=logging.INFO,
                  stderr_level=logging.WARNING, **kwargs) :
    """Equivalent to subprocess.Popen except log stdout and stderr
    where appropriate. Also log the command being called."""
    #we use None as sigil values for stdin,stdout,stderr above so we
    # can distinguish from the caller passing in Pipe.

    if not kwargs.has_key("close_fds") and os.name == 'posix' :
        #http://old.nabble.com/subprocess.Popen-pipeline-bug--td16026600.html
        kwargs['close_fds'] = True

    if not isinstance(cmd,str) :
        cmd = [str(e) for e in cmd]

    if(logger):
        #if we are logging, record the command we're running,
        #trying to strip out passwords.
        logger.debug("Running cmd: %s",
                     isinstance(cmd,str) and cmd or subprocess.list2cmdline(cmd))

    p = subprocess.Popen(cmd, stdin=stdin,
                         stdout=(stdout or (logger and PIPE)),
                         stderr=(stderr or (logger and PIPE)),
                         **kwargs)
    if logger :
        def monitor(level, src, name) :
            #if the cmd[0] (the binary) contains a full path, just get the name
            lname = "%s.%s" % (os.path.basename(cmd[0]), name)
            if(hasattr(logger, 'name')) :
                lname = "%s.%s" % (logger.name, lname)
            sublog = logging.getLogger(lname)

            def tfn() :
                l = src.readline()
                while l != "":
                    sublog.log(level,l.strip())
                    l = src.readline()

            th = threading.Thread(target=tfn,name=lname)
            p.__setattr__("std%s_thread" % name, th)
            th.start()

        if stdout == None : monitor(stdout_level, p.stdout,"out")
        if stderr == None : monitor(stderr_level, p.stderr,"err")
    return p

@contextmanager
def guardPopen(cmd, timeout=0.1, timeout_count=2, **kwargs) :
    popen = None
    logger = kwargs.get('logger')
    try :
        popen = capturedPopen(cmd,**kwargs)
        yield popen
    finally :
        if popen:
            while popen.poll() == None and timeout and timeout_count > 0 :
                if logger:
                    logger.debug("Things don't look right yet waiting for %ss (%s tries left)",
                                 timeout, timeout_count)
                time.sleep(timeout)
                timeout_count -=1
            if popen.poll() == None:
                if logger:
                    logger.exception("Terminating %s", cmd[0])
                popen.terminate()

def selfCaptured(klass) :
    def mungekw(self,kwargs) :
        if(not kwargs.has_key("logger")): kwargs["logger"] = self.logger
        return kwargs
    def add(func) :
        def newfunc(self,cmd,**kwargs) :
            return func(cmd, **mungekw(self,kwargs))
        setattr(klass,func.__name__,newfunc)

    add(capturedCall)
    add(capturedPopen)
    add(guardPopen)
    return klass

def which(program):
    def is_exe(fpath):
        return os.path.exists(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

def stream_to_file(stream,path,bufsize=8192) :
    with open(path, "wb") as f:
        bytes=0
        while True :
            buf = stream.read(bufsize)
            if buf == "" :  break #EOF
            bytes += len(buf)
            f.write(buf)
        return bytes

@contextmanager
def mkstemp_overwrite(real_filename, conflict_overwrite=True, logger=None, **kwargs) :
    """For writing to a temporary file and then move it ontop of a
    (possibly) existing file only when finished.  This enables us to
    perform long running operations on a file that other people might
    be using and let everyone else see a consistent version"""

    mtime1 = mtime2 = None
    if os.path.exists(real_filename) :
        mtime1 = os.path.getmtime(real_filename)

    (fd,path) = tempfile.mkstemp(**kwargs)

    filelike = os.fdopen(fd,'wb')
    yield filelike
    filelike.close()

    if os.path.exists(real_filename) :
        mtime2 = os.path.getmtime(real_filename)

    if mtime1 != mtime2 and logger:
        logger.warning("Potential conflict on %r, overwrite: %s",
                           real_filename, conflict_overwrite)

    if mtime1 == mtime2 or conflict_overwrite :
        os.rename(path,real_filename)

# Copright (c) 2011  Accelerated Data Works
# All rights reserved.

# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

#     Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.

#     Redistributions in binary form must reproduce the above
#     copyright notice, this list of conditions and the following
#     disclaimer in the documentation and/or other materials provided
#     with the distribution.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
<commit_msg>Another pass over mkstemp_overwrite adding doc, logging, and cleanup. re Programming:#555 (0.25)<commit_after>import os, os.path, logging, subprocess, threading, time, errno
import tempfile
from subprocess import PIPE
from contextlib import contextmanager

def ensure_dir(dir, logger=None) :
    if not os.path.exists(dir) :
        try:
            if logger: logger.debug("Making dir: %s", dir)
            os.makedirs(dir)
            if logger: logger.info("Created dir: %s", dir)
        except OSError, e :
            #not entirely sure why we are getting errors,
            #http://docs.python.org/library/os.path.html#os.path.exists
            #says this could be related to not being able to call
            #os.stat, but I can do so from the command line python
            #just fine.
            if os.path.exists(dir) :
                if logger: logger.debug("Erred, already exists: e.errno: %s",e.errno)
                return
            else:
                raise

def withDir(dir, fn, *args,**kwargs) :
    olddir = os.getcwd()
    try :
        os.chdir(dir)
        return fn(*args,**kwargs)
    finally :
        os.chdir(olddir)

def pprint_bytes(bytes) :
    suffix = 'B'
    bytes= float(bytes)
    if bytes >= 1024 :
        bytes = bytes / 1024
        suffix = 'KB'
    if bytes >= 1024 :
        bytes = bytes / 1024
        suffix = 'MB'
    if bytes >= 1024 :
        bytes = bytes / 1024
        suffix = 'GB'
    if bytes >= 1024 :
        bytes = bytes / 1024
        suffix = 'TB'
    return '%.2f%s' % (bytes,suffix)

def exec_external( proc ):
    return subprocess.call(['python', proc])

def capturedCall(cmd, **kwargs) :
    """Do the equivelent of the subprocess.call except
    log the stderr and stdout where appropriate."""
    p= capturedPopen(cmd,**kwargs)
    rc = p.wait()
    #this is a cheap attempt to make sure the monitors
    #are scheduled and hopefully finished.
    time.sleep(0.01)
    time.sleep(0.01)
    return rc

def capturedPopen(cmd, stdin=None, stdout=None, stderr=None,
                  logger=logging,
                  stdout_level=logging.INFO,
                  stderr_level=logging.WARNING, **kwargs) :
    """Equivalent to subprocess.Popen except log stdout and stderr
    where appropriate. Also log the command being called."""
    #we use None as sigil values for stdin,stdout,stderr above so we
    # can distinguish from the caller passing in Pipe.

    if not kwargs.has_key("close_fds") and os.name == 'posix' :
        #http://old.nabble.com/subprocess.Popen-pipeline-bug--td16026600.html
        kwargs['close_fds'] = True

    if not isinstance(cmd,str) :
        cmd = [str(e) for e in cmd]

    if(logger):
        #if we are logging, record the command we're running,
        #trying to strip out passwords.
        logger.debug("Running cmd: %s",
                     isinstance(cmd,str) and cmd or subprocess.list2cmdline(cmd))

    p = subprocess.Popen(cmd, stdin=stdin,
                         stdout=(stdout or (logger and PIPE)),
                         stderr=(stderr or (logger and PIPE)),
                         **kwargs)
    if logger :
        def monitor(level, src, name) :
            #if the cmd[0] (the binary) contains a full path, just get the name
            lname = "%s.%s" % (os.path.basename(cmd[0]), name)
            if(hasattr(logger, 'name')) :
                lname = "%s.%s" % (logger.name, lname)
            sublog = logging.getLogger(lname)

            def tfn() :
                l = src.readline()
                while l != "":
                    sublog.log(level,l.strip())
                    l = src.readline()

            th = threading.Thread(target=tfn,name=lname)
            p.__setattr__("std%s_thread" % name, th)
            th.start()

        if stdout == None : monitor(stdout_level, p.stdout,"out")
        if stderr == None : monitor(stderr_level, p.stderr,"err")
    return p

@contextmanager
def guardPopen(cmd, timeout=0.1, timeout_count=2, **kwargs) :
    popen = None
    logger = kwargs.get('logger')
    try :
        popen = capturedPopen(cmd,**kwargs)
        yield popen
    finally :
        if popen:
            while popen.poll() == None and timeout and timeout_count > 0 :
                if logger:
                    logger.debug("Things don't look right yet waiting for %ss (%s tries left)",
                                 timeout, timeout_count)
                time.sleep(timeout)
                timeout_count -=1
            if popen.poll() == None:
                if logger:
                    logger.exception("Terminating %s", cmd[0])
                popen.terminate()

def selfCaptured(klass) :
    def mungekw(self,kwargs) :
        if(not kwargs.has_key("logger")): kwargs["logger"] = self.logger
        return kwargs
    def add(func) :
        def newfunc(self,cmd,**kwargs) :
            return func(cmd, **mungekw(self,kwargs))
        setattr(klass,func.__name__,newfunc)

    add(capturedCall)
    add(capturedPopen)
    add(guardPopen)
    return klass

def which(program):
    def is_exe(fpath):
        return os.path.exists(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

def stream_to_file(stream,path,bufsize=8192) :
    with open(path, "wb") as f:
        bytes=0
        while True :
            buf = stream.read(bufsize)
            if buf == "" :  break #EOF
            bytes += len(buf)
            f.write(buf)
        return bytes

@contextmanager
def mkstemp_overwrite(destination, conflict_overwrite=True, logger=None,
                      cleanup=True, **kwargs) :
    """For writing to a temporary file and then move it ontop of a
    (possibly) existing file only when finished.  This enables us to
    perform long running operations on a file that other people might
    be using and let everyone else see a consistent version

 * conflict_overwrite=True: whether or not to overwrite the file if
   it's been modified in the intermediate time.
 * logger=None: if provided log information to it.
 * cleanup=True: ensure the tempfile is delete when we exit the
   function (e.g. an error or conflict)
 * other args are passed to tempfile.mkstemp

Example:
with mkstemp_overwrite('foobar.txt') as f:
   #long time intensive processing
   f.write('stuff\n')

    """

    mtime1 = mtime2 = None
    if os.path.exists(destination) :
        mtime1 = os.path.getmtime(destination)
    if logger : logger.debug("Found timestamp(1) %s on %r", mtime1, destination)

    (fd,path) = tempfile.mkstemp(**kwargs)
    try :
        filelike = os.fdopen(fd,'wb')
        yield filelike
        filelike.close()

        if os.path.exists(destination) :
            mtime2 = os.path.getmtime(destination)
        if logger : logger.debug("Found timestamp(2) %s on %r", mtime1, destination)

        if mtime1 != mtime2 and logger:
            logger.warning("Potential conflict on %r, overwrite: %s",
                               destination, conflict_overwrite)

        if conflict_overwrite or mtime1 == mtime2 :
            #TODO: permissions?
            os.rename(path,destination)
    finally :
        if cleanup and os.path.exists(path) :
            if logger :
                logger.info("Cleaning up leftover tempfile %r", path)
            os.remove(path)

# Copright (c) 2011  Accelerated Data Works
# All rights reserved.

# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

#     Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.

#     Redistributions in binary form must reproduce the above
#     copyright notice, this list of conditions and the following
#     disclaimer in the documentation and/or other materials provided
#     with the distribution.

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
<|endoftext|>


<commit_before>"""Convenient, high-level file API"""

# Copyright (c) 2017 Aubrey Barnard.  This is free software released
# under the MIT license.  See LICENSE for details.

import io
import os.path

# TODO how handle streams? can they be effectively wrapped?
# TODO how handle resolving names of executables?
# TODO path templates? shell names/tokens/words?

def new(file):
    if isinstance(file, File):
        return file
    elif isinstance(file, str):
        return File(file)
    elif isinstance(file, io.IOBase):
        return Stream(file)
    else:
        raise TypeError('Not a filename or stream: {}'.format())

def open(path, mode='rt'):
    return new(path).open(mode=mode)

class File:
    """File paths as objects, like pathlib."""

    # TODO consider subclassing pathlib with Python 3.4

    def __init__(self, *path_components):
        num_components = len(path_components)
        if num_components == 0 or path_components[0] == '':
            self._path = '.'
        elif num_components == 1:
            self._path = path_components[0]
        else:
            self._path = os.path.join(*path_components)
        # Strip trailing separator
        if self._path.endswith(os.path.sep):
            self._path = self._path[:-1]
        # Properties for lazy construction
        self._parent = None
        self._pieces = None

    @property
    def path(self):
        return self._path

    @property
    def name(self):
        return os.path.basename(self._path)

    @property
    def parent(self):
        if self._path == '.':
            return self # Avoids storing self reference
        elif self._parent is None:
            parent, name = os.path.split(self._path)
            if parent:
                self._parent = File(parent)
            else:
                self._parent = File('.')
        return self._parent

    @property
    def pieces(self):
        if self._pieces is None:
            # Split on dots
            raw_pieces = self.name.split('.')
            # Add dots back to all pieces except the first piece
            pieces = ['.' + p for p in raw_pieces[1:]]
            # Add the first piece back if not empty
            if raw_pieces[0] != '':
                pieces.insert(0, raw_pieces[0])
            self._pieces = pieces
        return self._pieces

    @property
    def stem(self):
        return self.pieces[0]

    @property
    def suffix(self):
        pieces = self.pieces
        if len(pieces) > 1:
            return pieces[-1]
        else:
            return ''

    @property
    def suffixes(self):
        return self.pieces[1:]

    def abspath(self):
        # Note: must be method as relies on IO and may change over time
        return os.path.abspath(self._path)

    def exists(self):
        return os.path.exists(self._path)

    def is_readable(self):
        return os.access(self._path, os.R_OK, effective_ids=True)

    def is_executable(self):
        return os.access(self._path, os.X_OK, effective_ids=True)

    def is_file(self):
        return os.path.isfile(self._path)

    def is_directory(self):
        return os.path.isdir(self._path)

    def is_readable_file(self):
        return self.exists() and self.is_file() and self.is_readable()

    def is_readable_directory(self):
        return (self.exists()
                and self.is_directory()
                and self.is_readable()
                and self.is_executable())

    def assert_readable(self): # TODO better name? check/assert/?
        if not self.is_readable():
            raise ValueError('Not a readable file: {}'.format(self))

    def is_writable(self):
        if os.path.exists(self._path):
            return (os.path.isfile(self._path) and
                    os.access(self._path, os.W_OK, effective_ids=True))
        else:
            parent = self.parent.path
            return (os.path.exists(parent) and
                    os.path.isdir(parent) and
                    os.access(parent, os.W_OK, effective_ids=True))

    def assert_writable(self): # TODO better name? check/assert/?
        if not self.is_writable():
            raise ValueError('Not a writable file: {}'.format(self))

    def open(self, mode='rt'):
        # Handle compressed files
        suffix = self.suffix.lower()
        if suffix in ('bz2', 'bzip2'):
            import bz2
            return bz2.open(self._path, mode=mode)
        elif suffix in ('xz', 'lzma'):
            import lzma
            return lzma.open(self._path, mode=mode)
        # Use gzip for all common Lempel-Ziv compression suffixes
        elif suffix in ('gz', 'z', 'Z'):
            import gzip
            return gzip.open(self._path, mode=mode)
        else:
            return io.open(self._path, mode=mode)

    def generate_lines(self):
        with self.open(mode='rt') as file:
            for line in file:
                yield line

    def __repr__(self):
        return 'File({})'.format(repr(self._path))

    def __str__(self):
        return self._path

    def join(self, *path_components):
        return File(self.path, *path_components)

class Stream:

    def __init__(self, stream, name=None):
        # Check for subclass of IOBase
        if not isinstance(stream, io.IOBase):
            raise TypeError(
                'Not an instance of IOBase: {!r}'.format(stream))
        self._stream = stream
        self._name = name
        if name is None:
            if hasattr(stream, 'name'):
                self._name = stream.name
            else:
                self._name = repr(stream)

    @property
    def name(self):
        return self._name

    def is_readable(self):
        return self._stream.readable()

    def assert_readable(self):
        if not self.is_readable():
            raise ValueError('Not a readable stream: {}'.format(self))

    def is_writable(self):
        return self._stream.writable()

    def assert_writable(self):
        if not self.is_writable():
            raise ValueError('Not a writable stream: {}'.format(self))

    def open(self, mode='rt'):
        # Check mode is compatible with the existing stream
        mode = mode.lower()
        if 'r' in mode:
            self.assert_readable()
        if 'w' in mode or 'a' in mode:
            self.assert_writable()
        if '+' in mode:
            self.assert_readable()
            self.assert_writable()
        text_stream = isinstance(self._stream, io.TextIOBase)
        if text_stream and 'b' in mode:
            raise TypeError(
                'Text stream not compatible with binary mode: {}'
                .format(mode))
        if not text_stream and 't' in mode:
            raise TypeError(
                'Binary stream not compatible with text mode: {}'
                .format(mode))
        return self._stream

    def generate_lines(self):
        with self.open(mode='rt') as file:
            for line in file:
                yield line

    def __repr__(self):
        return 'Stream({})'.format(self._name)
<commit_msg>Fixed `files.File.open` to handle leading dots on the filename suffix.<commit_after>"""Convenient, high-level file API"""

# Copyright (c) 2017 Aubrey Barnard.  This is free software released
# under the MIT license.  See LICENSE for details.

import io
import os.path

# TODO how handle streams? can they be effectively wrapped?
# TODO how handle resolving names of executables?
# TODO path templates? shell names/tokens/words?

def new(file):
    if isinstance(file, File):
        return file
    elif isinstance(file, str):
        return File(file)
    elif isinstance(file, io.IOBase):
        return Stream(file)
    else:
        raise TypeError('Not a filename or stream: {}'.format())

def open(path, mode='rt'):
    return new(path).open(mode=mode)

class File:
    """File paths as objects, like pathlib."""

    # TODO consider subclassing pathlib with Python 3.4

    def __init__(self, *path_components):
        num_components = len(path_components)
        if num_components == 0 or path_components[0] == '':
            self._path = '.'
        elif num_components == 1:
            self._path = path_components[0]
        else:
            self._path = os.path.join(*path_components)
        # Strip trailing separator
        if self._path.endswith(os.path.sep):
            self._path = self._path[:-1]
        # Properties for lazy construction
        self._parent = None
        self._pieces = None

    @property
    def path(self):
        return self._path

    @property
    def name(self):
        return os.path.basename(self._path)

    @property
    def parent(self):
        if self._path == '.':
            return self # Avoids storing self reference
        elif self._parent is None:
            parent, name = os.path.split(self._path)
            if parent:
                self._parent = File(parent)
            else:
                self._parent = File('.')
        return self._parent

    @property
    def pieces(self):
        if self._pieces is None:
            # Split on dots
            raw_pieces = self.name.split('.')
            # Add dots back to all pieces except the first piece
            pieces = ['.' + p for p in raw_pieces[1:]]
            # Add the first piece back if not empty
            if raw_pieces[0] != '':
                pieces.insert(0, raw_pieces[0])
            self._pieces = pieces
        return self._pieces

    @property
    def stem(self):
        return self.pieces[0]

    @property
    def suffix(self):
        pieces = self.pieces
        if len(pieces) > 1:
            return pieces[-1]
        else:
            return ''

    @property
    def suffixes(self):
        return self.pieces[1:]

    def abspath(self):
        # Note: must be method as relies on IO and may change over time
        return os.path.abspath(self._path)

    def exists(self):
        return os.path.exists(self._path)

    def is_readable(self):
        return os.access(self._path, os.R_OK, effective_ids=True)

    def is_executable(self):
        return os.access(self._path, os.X_OK, effective_ids=True)

    def is_file(self):
        return os.path.isfile(self._path)

    def is_directory(self):
        return os.path.isdir(self._path)

    def is_readable_file(self):
        return self.exists() and self.is_file() and self.is_readable()

    def is_readable_directory(self):
        return (self.exists()
                and self.is_directory()
                and self.is_readable()
                and self.is_executable())

    def assert_readable(self): # TODO better name? check/assert/?
        if not self.is_readable():
            raise ValueError('Not a readable file: {}'.format(self))

    def is_writable(self):
        if os.path.exists(self._path):
            return (os.path.isfile(self._path) and
                    os.access(self._path, os.W_OK, effective_ids=True))
        else:
            parent = self.parent.path
            return (os.path.exists(parent) and
                    os.path.isdir(parent) and
                    os.access(parent, os.W_OK, effective_ids=True))

    def assert_writable(self): # TODO better name? check/assert/?
        if not self.is_writable():
            raise ValueError('Not a writable file: {}'.format(self))

    def open(self, mode='rt'):
        # Handle compressed files
        suffix = self.suffix.lower().lstrip('.')
        if suffix in ('bz2', 'bzip2'):
            import bz2
            return bz2.open(self._path, mode=mode)
        elif suffix in ('xz', 'lzma'):
            import lzma
            return lzma.open(self._path, mode=mode)
        # Use gzip for all common Lempel-Ziv compression suffixes
        elif suffix in ('gz', 'z', 'Z'):
            import gzip
            return gzip.open(self._path, mode=mode)
        else:
            return io.open(self._path, mode=mode)

    def generate_lines(self):
        with self.open(mode='rt') as file:
            for line in file:
                yield line

    def __repr__(self):
        return 'File({})'.format(repr(self._path))

    def __str__(self):
        return self._path

    def join(self, *path_components):
        return File(self.path, *path_components)

class Stream:

    def __init__(self, stream, name=None):
        # Check for subclass of IOBase
        if not isinstance(stream, io.IOBase):
            raise TypeError(
                'Not an instance of IOBase: {!r}'.format(stream))
        self._stream = stream
        self._name = name
        if name is None:
            if hasattr(stream, 'name'):
                self._name = stream.name
            else:
                self._name = repr(stream)

    @property
    def name(self):
        return self._name

    def is_readable(self):
        return self._stream.readable()

    def assert_readable(self):
        if not self.is_readable():
            raise ValueError('Not a readable stream: {}'.format(self))

    def is_writable(self):
        return self._stream.writable()

    def assert_writable(self):
        if not self.is_writable():
            raise ValueError('Not a writable stream: {}'.format(self))

    def open(self, mode='rt'):
        # Check mode is compatible with the existing stream
        mode = mode.lower()
        if 'r' in mode:
            self.assert_readable()
        if 'w' in mode or 'a' in mode:
            self.assert_writable()
        if '+' in mode:
            self.assert_readable()
            self.assert_writable()
        text_stream = isinstance(self._stream, io.TextIOBase)
        if text_stream and 'b' in mode:
            raise TypeError(
                'Text stream not compatible with binary mode: {}'
                .format(mode))
        if not text_stream and 't' in mode:
            raise TypeError(
                'Binary stream not compatible with text mode: {}'
                .format(mode))
        return self._stream

    def generate_lines(self):
        with self.open(mode='rt') as file:
            for line in file:
                yield line

    def __repr__(self):
        return 'Stream({})'.format(self._name)
<|endoftext|>


<commit_before>#-------------------------------------------------------------------------------------------------
#MIT License
#If you use this in something awesome, credit would be nice but isn't required
#Written by Keilan
#Website: www.scholtek.com
#Github: https://github.com/Keilan
#-------------------------------------------------------------------------------------------------

#-------------------------------------------------------------------------------------------------
#Imports
#-------------------------------------------------------------------------------------------------
import os
import sys
import random
from copy import deepcopy
from multiprocessing import Pool
import jsonpickle

import numpy
import Image, ImageDraw

#-------------------------------------------------------------------------------------------------
#Knobs and Dials
#-------------------------------------------------------------------------------------------------
POP_PER_GENERATION = 50
MUTATION_CHANCE = 0.01
ADD_GENE_CHANCE = 0.3
REM_GENE_CHANCE = 0.2
INITIAL_GENES = 50

#How often to output images and save files
GENERATIONS_PER_IMAGE = 100
GENERATIONS_PER_SAVE = 100

try:
    globalTarget = Image.open("reference.png")
except IOError:
    print "File reference.png must be located in the same directory as poly_artist.py."
    exit()

#-------------------------------------------------------------------------------------------------
#Helper Classes
#-------------------------------------------------------------------------------------------------
class Point:
    """
    A 2D point. You can add them together if you want.
    """
    def __init__(self,x,y):
        self.x = x
        self.y = y

    def __add__(self,o):
        return Point(self.x+o.x,self.y+o.y)

class Color:
    """
    A color. You can shift it by a given value.
    """
    def __init__(self,r,g,b):
        self.r = r
        self.g = g
        self.b = b

    def shift(self,r,g,b):
        self.r = max(0,min(255,self.r+r))
        self.g = max(0,min(255,self.g+g))
        self.b = max(0,min(255,self.b+b))

    def __str__(self):
        return "({},{},{})".format(self.r,self.g,self.b)

#-------------------------------------------------------------------------------------------------
#Genetic Classes
#-------------------------------------------------------------------------------------------------
class Gene:
    """
    A gene is the object that can be mutated. Genetic algorithms work by randomly mutating genes
    and then using some function to determine how "ideal" the resulting organism is.

    This one is basically a circle, with a size, color, and position on the canvas.
    """
    def __init__(self,size):
        self.size = size #The canvas size so we know the maximum position value

        self.diameter = random.randint(5,15)
        self.pos = Point(random.randint(0,size[0]),random.randint(0,size[1]))
        self.color = Color(random.randint(0,255),random.randint(0,255),random.randint(0,255))
        self.params = ["diameter","pos","color"]

    def mutate(self):
        #Decide how big the mutation will be
        mutation_size = max(1,int(round(random.gauss(15,4))))/100

        #Decide what will be mutated
        mutation_type = random.choice(self.params)

        #Mutate the thing
        if mutation_type == "diameter":
            self.diameter = max(1,random.randint(int(self.diameter*(1-mutation_size)),int(self.diameter*(1+mutation_size))))

        elif mutation_type == "pos":
            x = max(0,random.randint(int(self.pos.x*(1-mutation_size)),int(self.pos.x*(1+mutation_size))))
            y = max(0,random.randint(int(self.pos.y*(1-mutation_size)),int(self.pos.y*(1+mutation_size))))
            self.pos = Point(min(x,self.size[0]),min(y,self.size[1]))

        elif mutation_type == "color":
            r = min(max(0,random.randint(int(self.color.r*(1-mutation_size)),int(self.color.r*(1+mutation_size)))),255)
            g = min(max(0,random.randint(int(self.color.g*(1-mutation_size)),int(self.color.g*(1+mutation_size)))),255)
            b = min(max(0,random.randint(int(self.color.b*(1-mutation_size)),int(self.color.b*(1+mutation_size)))),255)
            self.color = Color(r,g,b)

    def getSave(self):
        """
        Allows us to save an individual gene in case the program is stopped.
        """
        so = {}
        so["size"] = self.size
        so["diameter"] = self.diameter
        so["pos"] = (self.pos.x,self.pos.y)
        so["color"] = (self.color.r,self.color.g,self.color.b)
        return so

    def loadSave(self,so):
        """
        Allows us to load an individual gene in case the program is stopped.
        """
        self.size = so["size"]
        self.diameter = so["diameter"]
        self.pos = Point(so["pos"][0],so["pos"][1])
        self.color = Color(so["color"][0],so["color"][1],so["color"][2])

class Organism:
    """
    The organism consists of a variety of genes that work together for some sort of effect. The main
    effect of the genetic algorithm takes place here, as each step involves mutating some of the
    organisms genes to produce offspring, and the best performing of those offspring carries on.

    This organism contains a bunch of genes that draw circles, working together to draw a picture.
    """
    def __init__(self,size,num):
        self.genes = []
        self.size = size

        #Create random genes up to the number given
        for _ in xrange(num):
            self.genes.append(Gene(size))

    def mutate(self):
        #For small numbers of genes, each one has a random chance of mutating
        if len(self.genes) < 200:
            for g in self.genes:
                if MUTATION_CHANCE < random.random():
                    g.mutate()

        #For large numbers of genes, pick a random sample, this is statistically equivalent and faster
        else:
            for g in random.sample(self.genes,int(len(self.genes)*MUTATION_CHANCE)):
                g.mutate()

        #We also have a chance to add or remove a gene
        if ADD_GENE_CHANCE < random.random():
            self.genes.append(Gene(self.size))
        if len(self.genes) > 0 and REM_GENE_CHANCE < random.random():
            self.genes.remove(random.choice(self.genes))

    def drawImage(self):
        """
        Using the Image module, use the genes to draw the image.
        """
        image = Image.new("RGB",self.size,(255,255,255))
        canvas = ImageDraw.Draw(image)

        for g in self.genes:
            color = (g.color.r,g.color.g,g.color.b)
            canvas.ellipse([g.pos.x-g.diameter,g.pos.y-g.diameter,g.pos.x+g.diameter,g.pos.y+g.diameter],outline=color,fill=color)

        return image

    def getSave(self,generation):
        """
        Allows us to save an individual organism in case the program is stopped.
        """
        so = []
        so.append(generation)
        for g in self.genes:
            so.append(g.getSave())
        return so

    def loadSave(self,so):
        """
        Allows us to load an individual organism in case the program is stopped.
        """
        self.genes = []
        gen = so[0]
        so = so[1:]
        for g in so:
            newGene = Gene(self.size)
            newGene.loadSave(g)
            self.genes.append(newGene)
        return gen

def fitness(im1,im2):
    """
    The fitness function is used by the genetic algorithm to determine how successful a given organism
    is. Usually a genetic algorithm is trying to either minimize or maximize this function.

    This one uses numpy to quickly compute the sum of the differences between the pixels.
    """
    #Convert Image types to numpy arrays
    i1 = numpy.array(im1,numpy.int16)
    i2 = numpy.array(im2,numpy.int16)
    dif = numpy.sum(numpy.abs(i1-i2))
    return (dif / 255.0 * 100) / i1.size

#-------------------------------------------------------------------------------------------------
#Functions to Make Stuff Run
#-------------------------------------------------------------------------------------------------
def run(cores,so=None):
    """
    Contains the loop that creates and tests new generations.
    """
    #Create storage directory in current directory
    if not os.path.exists("results"):
        os.mkdir("results")

    #Create output log file
    f = file(os.path.join("results","log.txt"),'a')

    target = globalTarget

    #Setup the multiprocessing pool
    p = Pool(cores)

    #Create the parent organism (with random genes)
    generation = 1
    parent = Organism(target.size,INITIAL_GENES)

    #Load the save if one is given
    if so != None:
        gen = parent.loadSave(jsonpickle.decode(so))
        generation = int(gen)
    prevScore = 101
    score = fitness(parent.drawImage(),target)

    #Infinite loop (until the process is interrupted)
    while True:
        #Print the current score and write it to the log file
        print "Generation {} - {}".format(generation,score)
        f.write("Generation {} - {}\n".format(generation,score))

        #Save an image of the current best organism to the results directory
        if (generation) % GENERATIONS_PER_IMAGE == 0:
            parent.drawImage().save(os.path.join("results","{}.png".format(generation)))
        generation += 1

        prevScore = score

        #Spawn children
        children = []
        scores = []

        #Keep the best from before in case all mutations are bad
        children.append(parent)
        scores.append(score)

        #Perform the mutations and add to the parent
        results = groupMutate(parent,POP_PER_GENERATION-1,p)
        newScores,newChildren = zip(*results)

        children.extend(newChildren)
        scores.extend(newScores)

        #Find the winner
        score = min(scores)
        parent = children[scores.index(score)]

        #Store a backup to resume running if the program is interrupted
        if generation % 100 == 0:
            sf = file(os.path.join("results","{}.txt".format(generation)),'w')
            sf.write(jsonpickle.encode(parent.getSave(generation)))
            sf.close()

def mutateAndTest(o):
    """
    Given an organism, perform a random mutation on it, and then use the fitness function to
    determine how accurate of a result the mutated offspring draws.
    """
    c = deepcopy(o)
    c.mutate()
    i1 = c.drawImage()
    i2 = globalTarget
    return (fitness(i1,i2),c)

def groupMutate(o,number,p):
    """
    Mutates and tests a number of organisms using the multiprocessing module.
    """
    print o
#    results = p.map(mutateAndTest,[o]*int(number))
    results = [mutateAndTest(i) for i in [o]*int(number)]
    return results

#-------------------------------------------------------------------------------------------------
#Main Function
#-------------------------------------------------------------------------------------------------
if __name__ == "__main__":
    #Check the arguments, options are currents -t (number of threads) and -s (save file)
    if len(sys.argv) > 1:
        args = sys.argv[1:]

        #Set defaults

        cores = max(1,cpu_count()-1)
        so = None
        for i,a in enumerate(args):
            if a == "-t":
                cores = int(args[i+1])
            elif a == "-s":
                with open(args[i+1],'r') as save:
                    so = save.read()

        run(cores,so)
<commit_msg>multiprocess repaired<commit_after>#-------------------------------------------------------------------------------------------------
#MIT License
#If you use this in something awesome, credit would be nice but isn't required
#Written by Keilan
#Website: www.scholtek.com
#Github: https://github.com/Keilan
#-------------------------------------------------------------------------------------------------

#-------------------------------------------------------------------------------------------------
#Imports
#-------------------------------------------------------------------------------------------------
import os
import sys
import random
from copy import deepcopy
import multiprocessing
import jsonpickle

import numpy
import Image, ImageDraw

#-------------------------------------------------------------------------------------------------
#Knobs and Dials
#-------------------------------------------------------------------------------------------------
POP_PER_GENERATION = 50
MUTATION_CHANCE = 0.01
ADD_GENE_CHANCE = 0.3
REM_GENE_CHANCE = 0.2
INITIAL_GENES = 50

#How often to output images and save files
GENERATIONS_PER_IMAGE = 100
GENERATIONS_PER_SAVE = 100

try:
    globalTarget = Image.open("reference.png")
except IOError:
    print "File reference.png must be located in the same directory as poly_artist.py."
    exit()

#-------------------------------------------------------------------------------------------------
#Helper Classes
#-------------------------------------------------------------------------------------------------
class Point:
    """
    A 2D point. You can add them together if you want.
    """
    def __init__(self,x,y):
        self.x = x
        self.y = y

    def __add__(self,o):
        return Point(self.x+o.x,self.y+o.y)

class Color:
    """
    A color. You can shift it by a given value.
    """
    def __init__(self,r,g,b):
        self.r = r
        self.g = g
        self.b = b

    def shift(self,r,g,b):
        self.r = max(0,min(255,self.r+r))
        self.g = max(0,min(255,self.g+g))
        self.b = max(0,min(255,self.b+b))

    def __str__(self):
        return "({},{},{})".format(self.r,self.g,self.b)

#-------------------------------------------------------------------------------------------------
#Genetic Classes
#-------------------------------------------------------------------------------------------------
class Gene:
    """
    A gene is the object that can be mutated. Genetic algorithms work by randomly mutating genes
    and then using some function to determine how "ideal" the resulting organism is.

    This one is basically a circle, with a size, color, and position on the canvas.
    """
    def __init__(self,size):
        self.size = size #The canvas size so we know the maximum position value

        self.diameter = random.randint(5,15)
        self.pos = Point(random.randint(0,size[0]),random.randint(0,size[1]))
        self.color = Color(random.randint(0,255),random.randint(0,255),random.randint(0,255))
        self.params = ["diameter","pos","color"]

    def mutate(self):
        #Decide how big the mutation will be
        mutation_size = max(1,int(round(random.gauss(15,4))))/100

        #Decide what will be mutated
        mutation_type = random.choice(self.params)

        #Mutate the thing
        if mutation_type == "diameter":
            self.diameter = max(1,random.randint(int(self.diameter*(1-mutation_size)),int(self.diameter*(1+mutation_size))))

        elif mutation_type == "pos":
            x = max(0,random.randint(int(self.pos.x*(1-mutation_size)),int(self.pos.x*(1+mutation_size))))
            y = max(0,random.randint(int(self.pos.y*(1-mutation_size)),int(self.pos.y*(1+mutation_size))))
            self.pos = Point(min(x,self.size[0]),min(y,self.size[1]))

        elif mutation_type == "color":
            r = min(max(0,random.randint(int(self.color.r*(1-mutation_size)),int(self.color.r*(1+mutation_size)))),255)
            g = min(max(0,random.randint(int(self.color.g*(1-mutation_size)),int(self.color.g*(1+mutation_size)))),255)
            b = min(max(0,random.randint(int(self.color.b*(1-mutation_size)),int(self.color.b*(1+mutation_size)))),255)
            self.color = Color(r,g,b)

    def getSave(self):
        """
        Allows us to save an individual gene in case the program is stopped.
        """
        so = {}
        so["size"] = self.size
        so["diameter"] = self.diameter
        so["pos"] = (self.pos.x,self.pos.y)
        so["color"] = (self.color.r,self.color.g,self.color.b)
        return so

    def loadSave(self,so):
        """
        Allows us to load an individual gene in case the program is stopped.
        """
        self.size = so["size"]
        self.diameter = so["diameter"]
        self.pos = Point(so["pos"][0],so["pos"][1])
        self.color = Color(so["color"][0],so["color"][1],so["color"][2])

class Organism:
    """
    The organism consists of a variety of genes that work together for some sort of effect. The main
    effect of the genetic algorithm takes place here, as each step involves mutating some of the
    organisms genes to produce offspring, and the best performing of those offspring carries on.

    This organism contains a bunch of genes that draw circles, working together to draw a picture.
    """
    def __init__(self,size,num):
        self.genes = []
        self.size = size

        #Create random genes up to the number given
        for _ in xrange(num):
            self.genes.append(Gene(size))

    def mutate(self):
        #For small numbers of genes, each one has a random chance of mutating
        if len(self.genes) < 200:
            for g in self.genes:
                if MUTATION_CHANCE < random.random():
                    g.mutate()

        #For large numbers of genes, pick a random sample, this is statistically equivalent and faster
        else:
            for g in random.sample(self.genes,int(len(self.genes)*MUTATION_CHANCE)):
                g.mutate()

        #We also have a chance to add or remove a gene
        if ADD_GENE_CHANCE < random.random():
            self.genes.append(Gene(self.size))
        if len(self.genes) > 0 and REM_GENE_CHANCE < random.random():
            self.genes.remove(random.choice(self.genes))

    def drawImage(self):
        """
        Using the Image module, use the genes to draw the image.
        """
        image = Image.new("RGB",self.size,(255,255,255))
        canvas = ImageDraw.Draw(image)

        for g in self.genes:
            color = (g.color.r,g.color.g,g.color.b)
            canvas.ellipse([g.pos.x-g.diameter,g.pos.y-g.diameter,g.pos.x+g.diameter,g.pos.y+g.diameter],outline=color,fill=color)

        return image

    def getSave(self,generation):
        """
        Allows us to save an individual organism in case the program is stopped.
        """
        so = []
        so.append(generation)
        for g in self.genes:
            so.append(g.getSave())
        return so

    def loadSave(self,so):
        """
        Allows us to load an individual organism in case the program is stopped.
        """
        self.genes = []
        gen = so[0]
        so = so[1:]
        for g in so:
            newGene = Gene(self.size)
            newGene.loadSave(g)
            self.genes.append(newGene)
        return gen

def fitness(im1,im2):
    """
    The fitness function is used by the genetic algorithm to determine how successful a given organism
    is. Usually a genetic algorithm is trying to either minimize or maximize this function.

    This one uses numpy to quickly compute the sum of the differences between the pixels.
    """
    #Convert Image types to numpy arrays
    i1 = numpy.array(im1,numpy.int16)
    i2 = numpy.array(im2,numpy.int16)
    dif = numpy.sum(numpy.abs(i1-i2))
    return (dif / 255.0 * 100) / i1.size

#-------------------------------------------------------------------------------------------------
#Functions to Make Stuff Run
#-------------------------------------------------------------------------------------------------
def run(cores,so=None):
    """
    Contains the loop that creates and tests new generations.
    """
    #Create storage directory in current directory
    if not os.path.exists("results"):
        os.mkdir("results")

    #Create output log file
    f = file(os.path.join("results","log.txt"),'a')

    target = globalTarget

    #Setup the multiprocessing pool
#    p = multiprocessing.Pool(cores)

    #Create the parent organism (with random genes)
    generation = 1
    parent = Organism(target.size,INITIAL_GENES)

    #Load the save if one is given
    if so != None:
        gen = parent.loadSave(jsonpickle.decode(so))
        generation = int(gen)
    prevScore = 101
    score = fitness(parent.drawImage(),target)

    #Infinite loop (until the process is interrupted)
    while True:
        p = multiprocessing.Pool(cores)
        #Print the current score and write it to the log file
        print "Generation {} - {}".format(generation,score)
        f.write("Generation {} - {}\n".format(generation,score))

        #Save an image of the current best organism to the results directory
        if (generation) % GENERATIONS_PER_IMAGE == 0:
            parent.drawImage().save(os.path.join("results","{}.png".format(generation)))
        generation += 1

        prevScore = score

        #Spawn children
        children = []
        scores = []

        #Keep the best from before in case all mutations are bad
        children.append(parent)
        scores.append(score)

        #Perform the mutations and add to the parent
        results = groupMutate(parent,POP_PER_GENERATION-1,p)
        newScores,newChildren = zip(*results)
        p.close()

        children.extend(newChildren)
        scores.extend(newScores)

        #Find the winner

        winners = sorted(zip(children,scores),key=lambda x: x[1])

        parent,score = winners[0]

        #Store a backup to resume running if the program is interrupted
        if generation % 100 == 0:
            sf = file(os.path.join("results","{}.txt".format(generation)),'w')
            sf.write(jsonpickle.encode(parent.getSave(generation)))
            sf.close()

def mutateAndTest(o):
    """
    Given an organism, perform a random mutation on it, and then use the fitness function to
    determine how accurate of a result the mutated offspring draws.
    """
    c = deepcopy(o)
    c.mutate()
    i1 = c.drawImage()
    i2 = globalTarget
    return (fitness(i1,i2),c)

def groupMutate(o,number,p):
    """
    Mutates and tests a number of organisms using the multiprocessing module.
    """
    results = p.map(mutateAndTest,[o]*int(number))
#    results = [mutateAndTest(i) for i in [o]*int(number)]
    return results

#-------------------------------------------------------------------------------------------------
#Main Function
#-------------------------------------------------------------------------------------------------
if __name__ == "__main__":
    #Check the arguments, options are currents -t (number of threads) and -s (save file)
    if len(sys.argv) > 1:
        args = sys.argv[1:]

        #Set defaults

        cores = max(1,multiprocessing.cpu_count()-1)
        so = None
        for i,a in enumerate(args):
            if a == "-t":
                cores = int(args[i+1])
            elif a == "-s":
                with open(args[i+1],'r') as save:
                    so = save.read()

        run(cores,so)
<|endoftext|>


<commit_before>import json
import os
import re

from google.cloud.bigquery import Client
from google.cloud.bigquery.job import QueryJobConfig

from pypinfo.fields import AGGREGATES, Downloads

FROM = """\
FROM
  TABLE_DATE_RANGE(
    [the-psf:pypi.downloads],
    DATE_ADD(CURRENT_TIMESTAMP(), {}, "day"),
    DATE_ADD(CURRENT_TIMESTAMP(), {}, "day")
  )
"""
START_DATE = '-31'
END_DATE = '-1'
DEFAULT_LIMIT = '10'

def create_config():
    config = QueryJobConfig()
    config.use_legacy_sql = True
    return config

def normalize(name):
    """https://www.python.org/dev/peps/pep-0503/#normalized-names"""
    return re.sub(r'[-_.]+', '-', name).lower()

def create_client(creds_file=None):
    creds_file = creds_file or os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')

    if creds_file is None:
        raise SystemError('Credentials could not be found.')

    project = json.load(open(creds_file))['project_id']
    return Client.from_service_account_json(creds_file, project=project)

def build_query(project, all_fields, start_date=None, end_date=None,
                days=None, limit=None, where=None, order=None, pip=None):
    project = normalize(project)

    start_date = start_date or START_DATE
    end_date = end_date or END_DATE
    limit = limit or DEFAULT_LIMIT

    if days:
        start_date = str(int(end_date) - int(days))

    if int(start_date) > 0 or int(end_date) > 0:
        raise ValueError('Dates must be in the past (negative).')

    if int(start_date) >= int(end_date):
        raise ValueError('End date must be greater than start date.')

    fields = []
    used_fields = set()
    for f in all_fields:
        if f not in used_fields:
            fields.append(f)
            used_fields.add(f)

    fields.append(Downloads)
    query = 'SELECT\n'

    for field in fields:
        query += '  {} as {},\n'.format(field.data, field.name)

    query += FROM.format(start_date, end_date)

    if where:
        query += 'WHERE\n  {}\n'.format(where)
    else:
        conditions = []
        if project:
            conditions.append('file.project = "{}"\n'.format(project))
        if pip:
            conditions.append('details.installer.name = "pip"\n')
        if conditions:
            query += 'WHERE\n  ' + '  AND '.join(conditions)

    if len(fields) > 1:
        gb = 'GROUP BY\n'
        initial_length = len(gb)

        for field in fields[:-1]:
            if field not in AGGREGATES:
                gb += '  {},\n'.format(field.name)

        if len(gb) > initial_length:
            query += gb

    query += 'ORDER BY\n  {} DESC\n'.format(order or Downloads.name)
    query += 'LIMIT {}'.format(limit)

    return query

def parse_query_result(query_job, query_rows):
    rows = [[field.name for field in query_job.result().schema]]
    rows.extend([str(item) for item in row] for row in query_rows)
    return rows

def add_percentages(rows, include_sign=True):

    headers = rows.pop(0)
    index = headers.index('download_count')
    headers.insert(index, 'percent')
    total_downloads = sum(int(row[index]) for row in rows)
    percent_format = '{:.1%}' if include_sign else '{:.1}'

    for r, row in enumerate(rows):
        percent = percent_format.format(int(row[index]) / total_downloads)
        row.insert(index, percent)

    rows.insert(0, headers)
    return rows

def tabulate(rows, markdown=False):
    column_widths = [0] * len(rows[0])
    right_align = [[False] * len(rows[0])] * len(rows)

    # Get max width of each column
    for r, row in enumerate(rows):
        for i, item in enumerate(row):
            if item.isdigit():
                # Separate the thousands
                rows[r][i] = "{:,}".format(int(item))
                right_align[r][i] = True
            elif item.endswith('%'):
                right_align[r][i] = True
            length = len(item)
            if length > column_widths[i]:
                column_widths[i] = length

    tabulated = '| '

    headers = rows.pop(0)
    for i, item in enumerate(headers):
        tabulated += item + ' ' * (column_widths[i] - len(item) + 1) + '| '

    tabulated = tabulated.rstrip()
    tabulated += '\n| '

    for i, item in enumerate(rows[0]):
        tabulated += '-' * (column_widths[i]-1)
        if right_align[0][i] and markdown:
            tabulated += ': | '
        else:
            tabulated += '- | '

    tabulated = tabulated.rstrip()
    tabulated += '\n'

    for r, row in enumerate(rows):
        for i, item in enumerate(row):
            num_spaces = column_widths[i] - len(item)
            tabulated += '| '
            if right_align[r][i]:
                tabulated += ' ' * num_spaces + item + ' '
            else:
                tabulated += item + ' ' * (num_spaces + 1)
        tabulated += '|\n'

    return tabulated

def format_json(rows, query_info, indent):
    headers, *data = rows
    rows = []

    for d in data:
        item = {}
        for i in range(len(headers)):
            if d[i].isdigit():
                d[i] = int(d[i])
            item[headers[i]] = d[i]
        rows.append(item)

    j = {'rows': rows, 'query': query_info}

    separators = (',', ':') if indent is None else None
    return json.dumps(j, indent=indent, separators=separators, sort_keys=True)
<commit_msg>increase resolution of percentages<commit_after>import json
import os
import re

from google.cloud.bigquery import Client
from google.cloud.bigquery.job import QueryJobConfig

from pypinfo.fields import AGGREGATES, Downloads

FROM = """\
FROM
  TABLE_DATE_RANGE(
    [the-psf:pypi.downloads],
    DATE_ADD(CURRENT_TIMESTAMP(), {}, "day"),
    DATE_ADD(CURRENT_TIMESTAMP(), {}, "day")
  )
"""
START_DATE = '-31'
END_DATE = '-1'
DEFAULT_LIMIT = '10'

def create_config():
    config = QueryJobConfig()
    config.use_legacy_sql = True
    return config

def normalize(name):
    """https://www.python.org/dev/peps/pep-0503/#normalized-names"""
    return re.sub(r'[-_.]+', '-', name).lower()

def create_client(creds_file=None):
    creds_file = creds_file or os.environ.get('GOOGLE_APPLICATION_CREDENTIALS')

    if creds_file is None:
        raise SystemError('Credentials could not be found.')

    project = json.load(open(creds_file))['project_id']
    return Client.from_service_account_json(creds_file, project=project)

def build_query(project, all_fields, start_date=None, end_date=None,
                days=None, limit=None, where=None, order=None, pip=None):
    project = normalize(project)

    start_date = start_date or START_DATE
    end_date = end_date or END_DATE
    limit = limit or DEFAULT_LIMIT

    if days:
        start_date = str(int(end_date) - int(days))

    if int(start_date) > 0 or int(end_date) > 0:
        raise ValueError('Dates must be in the past (negative).')

    if int(start_date) >= int(end_date):
        raise ValueError('End date must be greater than start date.')

    fields = []
    used_fields = set()
    for f in all_fields:
        if f not in used_fields:
            fields.append(f)
            used_fields.add(f)

    fields.append(Downloads)
    query = 'SELECT\n'

    for field in fields:
        query += '  {} as {},\n'.format(field.data, field.name)

    query += FROM.format(start_date, end_date)

    if where:
        query += 'WHERE\n  {}\n'.format(where)
    else:
        conditions = []
        if project:
            conditions.append('file.project = "{}"\n'.format(project))
        if pip:
            conditions.append('details.installer.name = "pip"\n')
        if conditions:
            query += 'WHERE\n  ' + '  AND '.join(conditions)

    if len(fields) > 1:
        gb = 'GROUP BY\n'
        initial_length = len(gb)

        for field in fields[:-1]:
            if field not in AGGREGATES:
                gb += '  {},\n'.format(field.name)

        if len(gb) > initial_length:
            query += gb

    query += 'ORDER BY\n  {} DESC\n'.format(order or Downloads.name)
    query += 'LIMIT {}'.format(limit)

    return query

def parse_query_result(query_job, query_rows):
    rows = [[field.name for field in query_job.result().schema]]
    rows.extend([str(item) for item in row] for row in query_rows)
    return rows

def add_percentages(rows, include_sign=True):

    headers = rows.pop(0)
    index = headers.index('download_count')
    headers.insert(index, 'percent')
    total_downloads = sum(int(row[index]) for row in rows)
    percent_format = '{:.2%}' if include_sign else '{:.2}'

    for r, row in enumerate(rows):
        percent = percent_format.format(int(row[index]) / total_downloads)
        row.insert(index, percent)

    rows.insert(0, headers)
    return rows

def tabulate(rows, markdown=False):
    column_widths = [0] * len(rows[0])
    right_align = [[False] * len(rows[0])] * len(rows)

    # Get max width of each column
    for r, row in enumerate(rows):
        for i, item in enumerate(row):
            if item.isdigit():
                # Separate the thousands
                rows[r][i] = "{:,}".format(int(item))
                right_align[r][i] = True
            elif item.endswith('%'):
                right_align[r][i] = True
            length = len(item)
            if length > column_widths[i]:
                column_widths[i] = length

    tabulated = '| '

    headers = rows.pop(0)
    for i, item in enumerate(headers):
        tabulated += item + ' ' * (column_widths[i] - len(item) + 1) + '| '

    tabulated = tabulated.rstrip()
    tabulated += '\n| '

    for i, item in enumerate(rows[0]):
        tabulated += '-' * (column_widths[i]-1)
        if right_align[0][i] and markdown:
            tabulated += ': | '
        else:
            tabulated += '- | '

    tabulated = tabulated.rstrip()
    tabulated += '\n'

    for r, row in enumerate(rows):
        for i, item in enumerate(row):
            num_spaces = column_widths[i] - len(item)
            tabulated += '| '
            if right_align[r][i]:
                tabulated += ' ' * num_spaces + item + ' '
            else:
                tabulated += item + ' ' * (num_spaces + 1)
        tabulated += '|\n'

    return tabulated

def format_json(rows, query_info, indent):
    headers, *data = rows
    rows = []

    for d in data:
        item = {}
        for i in range(len(headers)):
            if d[i].isdigit():
                d[i] = int(d[i])
            item[headers[i]] = d[i]
        rows.append(item)

    j = {'rows': rows, 'query': query_info}

    separators = (',', ':') if indent is None else None
    return json.dumps(j, indent=indent, separators=separators, sort_keys=True)
<|endoftext|>


<commit_before>__author__ = 'Zhi Deng'

import os
import glob
import shutil
from monty.os import cd
from pymatgen.core.composition import Composition

def _write_input_from_temp(template, jobname, params):
    '''
    Private method as the writer of QE input file.
    '''
    with open('%s.pw.in' % jobname, 'w') as f:
        f.write(template.format(**params))

def get_template(template):
    with open('%s.pw.in.template' % template) as f:
        t = f.read()
    struc_info = template.split('.')
    return t, struc_info

def write_input(temp, params, xkeys, func='pbe', parr=True):
    '''
    Method to write QE input file in scratch dir from template.
    :param temp (list): Returned from get_template() method.
    :param params (dict): Dict of parameters to pass into template. Keys
        are the strings within the placeholders in the template.
    :param xkeys ([str]): Parameters to distinguish between different
        inpur files.temwr
    :param func (str): Functional used. Default to GGA. Refer to the
        format of filename of template files.
    :param parr (bool): Whether use parallel mode (send the calculation
        to computational nodes). Default to True. Note in non-parallel
        (serial, run the calculation on login nodes), you need to cp the
        psp files into scratch dir.
    :return: None
    '''
    t, info = temp
    pl = [str(params[k]) for k in xkeys]
    jobname = "_".join(info+pl)
    if parr:
        c = Composition(info[0])
        elements = [e.symbol for e in c.elements]
        psp = {}
        for e in elements:
            psp[e] = os.path.abspath(glob.glob('../%s.%s*.UPF' % (e, func))[0])
        os.makedirs(jobname)
        with cd(jobname):
            _write_input_from_temp(t, jobname, params)
            for e in psp.keys():
                shutil.copyfile(psp[e],os.path.split(psp[e])[1])
    else:
        _write_input_from_temp(t, jobname, params)

if __name__ == '__main__':
    pass
<commit_msg>add doc to get_template<commit_after>__author__ = 'Zhi Deng'

import os
import glob
import shutil
from monty.os import cd
from pymatgen.core.composition import Composition

def _write_input_from_temp(template, jobname, params):
    '''
    Private method as the writer of QE input file.
    '''
    with open('%s.pw.in' % jobname, 'w') as f:
        f.write(template.format(**params))

def get_template(template):
    '''
    Method to initialize the template from file.
    :param template (str): Template name (e.g., 'Cu.fcc')
    :return: template (str), structural info (list)
    '''
    with open('%s.pw.in.template' % template) as f:
        t = f.read()
    struc_info = template.split('.')
    return t, struc_info

def write_input(temp, params, xkeys, func='pbe', parr=True):
    '''
    Method to write QE input file in scratch dir from template.
    :param temp (list): Returned from get_template() method.
    :param params (dict): Dict of parameters to pass into template. Keys
        are the strings within the placeholders in the template.
    :param xkeys ([str]): Parameters to distinguish between different
        inpur files.temwr
    :param func (str): Functional used. Default to GGA. Refer to the
        format of filename of template files.
    :param parr (bool): Whether use parallel mode (send the calculation
        to computational nodes). Default to True. Note in non-parallel
        (serial, run the calculation on login nodes), you need to cp the
        psp files into scratch dir.
    :return: None
    '''
    t, info = temp
    pl = [str(params[k]) for k in xkeys]
    jobname = "_".join(info+pl)
    if parr:
        c = Composition(info[0])
        elements = [e.symbol for e in c.elements]
        psp = {}
        for e in elements:
            psp[e] = os.path.abspath(glob.glob('../%s.%s*.UPF' % (e, func))[0])
        os.makedirs(jobname)
        with cd(jobname):
            _write_input_from_temp(t, jobname, params)
            for e in psp.keys():
                shutil.copyfile(psp[e],os.path.split(psp[e])[1])
    else:
        _write_input_from_temp(t, jobname, params)

if __name__ == '__main__':
    pass
<|endoftext|>


<commit_before>import numpy as np

required = ['inputs', 'actions', 'tmax', 'n_gradient', 'n_validation']
default  = {
    'Performance':       None,
    'N':                 100,
    'p0':                0.1,
    'lr':                0.004,
    'baseline_lr':       0.004,
    'max_iter':          100000,
    'fix':               [],
    'baseline_fix':      [],
    'target_reward':     np.inf,
    'mode':              'episodic',
    'network_type':      'gru',
    'R_ABORTED':         -1,
    'checkfreq':         50,
    'dt':                10,
    'tau':               100,
    'var_rec':           0.02,
    'baseline_var_rec':  0.01,
    'L2_r':              0,
    'baseline_L2_r':     0,
    'Win':               1,
    'baseline_Win':      1,
    'Win_mask':          None,
    'baseline_Win_mask': None,
    'rho':               2,
    'baseline_rho':      2,
    'L1_Wrec':           0,
    'L2_Wrec':           0,
    'policy_seed':       1,
    'baseline_seed':     2
    }
<commit_msg>Initial commit<commit_after>import numpy as np

required = ['inputs', 'actions', 'tmax', 'n_gradient', 'n_validation']
default  = {
    'Performance':       None,
    'N':                 100,
    'p0':                0.1,
    'lr':                0.004,
    'baseline_lr':       0.004,
    'max_iter':          100000,
    'fix':               [],
    'baseline_fix':      [],
    'target_reward':     np.inf,
    'mode':              'episodic',
    'network_type':      'gru',
    'R_ABORTED':         -1,
    'checkfreq':         50,
    'dt':                10,
    'tau':               100,
    'var_rec':           0.01,
    'baseline_var_rec':  0.01,
    'L2_r':              0,
    'baseline_L2_r':     0,
    'Win':               1,
    'baseline_Win':      1,
    'Win_mask':          None,
    'baseline_Win_mask': None,
    'rho':               2,
    'baseline_rho':      2,
    'L1_Wrec':           0,
    'L2_Wrec':           0,
    'policy_seed':       1,
    'baseline_seed':     2
    }
<|endoftext|>