<commit_before>"""Generates input files for fspsq."""
import socket
import os
import subprocess
import multiprocessing
import datetime
import cPickle as pickle

import pytz
import pymongo
import numpy as np

from bson.binary import Binary
from pymongo.son_manipulator import SONManipulator

class FSPSLibrary(object):
    """For building/accessing a library of FSPS stellar population models.
    
    FSPSLibrary is itself an interface for the PyMongo database that stores
    model parameters and realizations from the FSPS code.
    """
    def __init__(self, libname, dbname="fsps"):
        super(FSPSLibrary, self).__init__()
        self.libname = libname
        self.dbname = dbname
        
        connection = pymongo.Connection()
        self.host = connection.host
        self.port = connection.port
        self.db = connection[self.dbname]
        self.db.add_son_manipulator(NumpySONManipulator())
        self.collection = self.db[self.libname]
    
    def register_pset(self, pset):
        """Adds a ParameterSet to the collection of model definitions."""
        self.remove_existing_model(pset.name)
        doc = {"_id": pset.name, "pset": pset.get_doc(),
            "compute_started": False, "compute_complete": False}
        self.collection.insert(doc)
    
    def compute_models(self, nThreads=1, maxN=10):
        """Fire up FSPS to generate SP models.
        
        Parameters
        ----------
        
        nThreads : int, optional
            Number of processes to be spawned, for multiprocessing. If set
            to `1`, multiprocessing is not used; useful for debugging.
        maxN : int, optional
            Maximum number of jobs that can be given to a single
            `fspsq` process.
        """
        fspsqPath = os.path.join(os.getcwd(), "fspsq") # echo for debug
        # fspsqPath = "echo"
        args = []
        for i in range(nThreads):
            processorName = "processor-%i" % i
            commandPath = processorName+".txt"
            args.append((self.libname, self.dbname, self.host, self.port,
                maxN, fspsqPath, commandPath))
        print args
        if nThreads > 1:
            pool = multiprocessing.Pool(processes=nThreads)
            pool.map(run_fspsq, args)
        else:
            map(run_fspsq, args)
    
    def reset(self):
        """Drop (delete) the library.
        
        The collection member is reinstantiated. This allows the collection
        to be freshly re-populated with models.
        """
        self.db.drop_collection(self.libname)
        self.collection = self.db[self.libname]
    
    def remove_existing_model(self, modelName):
        """Delete a model, by name, if it exists."""
        if self.collection.find_one({"_id": modelName}) is not None:
            self.collection.remove({"_id": modelName})

def run_fspsq(args):
    """A process for running fsps jobs."""
    # print "hello"
    libname, dbname, host, port, maxN, fspsqPath, commandPath = args
    thisHost = socket.gethostname() # hostname of current process
    
    # Connect to the library in MongoDB
    connection = pymongo.Connection(host=host, port=port)
    db = connection[dbname]
    db.add_son_manipulator(NumpySONManipulator())
    collection = db[libname]
    
    commonVarSets = _make_common_var_sets(collection)
    print "commonVarSets:", commonVarSets
    for varSet in commonVarSets:
        while True:
            psets = []
            modelNames = []
            now = datetime.datetime.utcnow()
            now = now.replace(tzinfo=pytz.utc)
            while len(psets) <= maxN:
                q = {"compute_complete": False, "compute_started": False}
                q.update(varSet)
                print "q", q
                doc = collection.find_and_modify(query=q,
                    update={"$set": {"compute_started": True,
                                     "compute_date": now,
                                     "compute_host": thisHost}},)
                print "doc", doc
                if doc is None: break # no available models
                modelName = str(doc['_id'])
                pset = ParameterSet(modelName, **doc['pset'])
                psets.append(pset)
                modelNames.append(pset.name)
            if len(psets) == 0: break # empty job queue
            # Startup a computation: write command file and start fspsq
            cmdTxt = "\n".join([p.command() for p in psets])+"\n"
            if os.path.exists(commandPath): os.remove(commandPath)
            f = open(commandPath, 'w')
            f.write(cmdTxt)
            f.close()
            shellCmd = _make_shell_command(fspsqPath, commandPath, varSet)
            print "cmd::", shellCmd
            subprocess.call(shellCmd, shell=True)
            # Get output data from FSPS
            _gather_fsps_outputs(collection, modelNames)

def _make_common_var_sets(c):
    """Make a list of common variable setups.
    
    .. note:: This merely lists all *possible* combinations of common variable
       configurations. There is not guarantee that there are models needing
       computation for each configuration set.
    
    Parameters
    ----------
    
    c : pymongo.Collection instance
    
    Returns
    -------
    
    List of dictionaries. Each dictionary has keys representing common
    variables: sfh, zmet, dust_type, imf_type, compute_vega_mags, redshift_colors
    """
    params = ['sfh', 'zmet', 'dust_type', 'imf_type',
        'compute_vega_mags', 'redshift_colors']
    possibleValues = {}
    for param in params:
        possibleValues[param] = c.distinct("pset."+param)
    groups = []
    for isfh in possibleValues['sfh']:
        for izmet in possibleValues['zmet']:
            for idust_type in possibleValues['dust_type']:
                for iimf_type in possibleValues['imf_type']:
                    for ivega in possibleValues['compute_vega_mags']:
                        for iredshift in possibleValues['redshift_colors']:
                            groups.append({"pset.sfh":isfh,"pset.zmet":izmet,
                                "pset.dust_type":idust_type,"pset.imf_type":iimf_type,
                                "pset.compute_vega_mags":ivega,
                                "pset.redshift_colors":iredshift})
    return groups

def _make_shell_command(fspsqPath, commandFilePath, varSet):
    """Crafts the command for running `fspsq`, returning a string."""
    params = ['sfh', 'zmet', 'dust_type', 'imf_type',
        'compute_vega_mags', 'redshift_colors']
    cmd = "%s %s %i %i %i %i %i %i" % (fspsqPath, commandFilePath,
        varSet['pset.sfh'], varSet['pset.zmet'], varSet['pset.dust_type'],
        varSet['pset.imf_type'], varSet['pset.compute_vega_mags'],
        varSet['pset.redshift_colors'])
    return cmd

def _gather_fsps_outputs(c, modelNames):
    """Adds the .mag and .spec outputs for each model to MongoDB.
    
    .. note:: Need to propagate the spec type through here...
    """
    fspsDir = os.environ['SPS_HOME'] # environment variable for FSPS
    outputDir = os.path.join(fspsDir, "OUTPUTS")
    print modelNames
    for modelName in modelNames:
        magPath = os.path.join(outputDir, modelName+".out.mags")
        magSpec = os.path.join(outputDir, modelName+".out.spec")
        npdata = ingest_output(magPath, magSpec)
        print c.find_one({"_id": modelName})
        # c.update({"_id": modelName}, {"$set": {"np_data": npdata, "compute_complete": True}})
        print "update complete!"

class ParameterSet(object):
    """An input parameter set for a FSPS model run."""
    def __init__(self, name, **kwargs):
        self.name = name # name of this model
        # Default values
        self.p = {"compute_vega_mags":0, "dust_type":0, "imf_type":0,
                "isoc_type":'pdva', "redshift_colors":0, "time_res_incr":2,
                "zred":0., "zmet":1, "sfh":0, "tau":1., "const":0.,
                "sf_start":0.,"tage":0., "fburst":0., "tburst":0., "imf1":1.3,
                "imf2":2.3, "imf3":2.3, "vdmc":0.08, "mdave":0.5,
                "dust_tesc":7., "dust1":0., "dust2":0., "dust_clumps":-99.,
                "frac_nodust":0., "dust_index":-0.7, "mwr":3.1,
                "uvb":1., "wgp1":1, "wgp2":1, "wgp3":0, "dell":0.,
                "delt":0., "sbss":0., "fbhb":0, "pagb":1.}
        self.knownParameters = self.p.keys()
        # Update values with user's arguments
        for k, v in kwargs.iteritems():
            if k in self.knownParameters:
                self.p[k] = v
    
    def command(self):
        """Write the string for this paramter set."""
        # These are pset variables, (aside from sfh and zmet)
        dt = [("zred","%.2f"),("tau","%.10f"),("const","%.4f"),<commit_msg>fspsq; only trying to insert one line of output to mongodb<commit_after>"""Generates input files for fspsq."""
import socket
import os
import subprocess
import multiprocessing
import datetime
import cPickle as pickle

import pytz
import pymongo
import numpy as np

from bson.binary import Binary
from pymongo.son_manipulator import SONManipulator

class FSPSLibrary(object):
    """For building/accessing a library of FSPS stellar population models.
    
    FSPSLibrary is itself an interface for the PyMongo database that stores
    model parameters and realizations from the FSPS code.
    """
    def __init__(self, libname, dbname="fsps"):
        super(FSPSLibrary, self).__init__()
        self.libname = libname
        self.dbname = dbname
        
        connection = pymongo.Connection()
        self.host = connection.host
        self.port = connection.port
        self.db = connection[self.dbname]
        self.db.add_son_manipulator(NumpySONManipulator())
        self.collection = self.db[self.libname]
    
    def register_pset(self, pset):
        """Adds a ParameterSet to the collection of model definitions."""
        self.remove_existing_model(pset.name)
        doc = {"_id": pset.name, "pset": pset.get_doc(),
            "compute_started": False, "compute_complete": False}
        self.collection.insert(doc)
    
    def compute_models(self, nThreads=1, maxN=10):
        """Fire up FSPS to generate SP models.
        
        Parameters
        ----------
        
        nThreads : int, optional
            Number of processes to be spawned, for multiprocessing. If set
            to `1`, multiprocessing is not used; useful for debugging.
        maxN : int, optional
            Maximum number of jobs that can be given to a single
            `fspsq` process.
        """
        fspsqPath = os.path.join(os.getcwd(), "fspsq") # echo for debug
        # fspsqPath = "echo"
        args = []
        for i in range(nThreads):
            processorName = "processor-%i" % i
            commandPath = processorName+".txt"
            args.append((self.libname, self.dbname, self.host, self.port,
                maxN, fspsqPath, commandPath))
        print args
        if nThreads > 1:
            pool = multiprocessing.Pool(processes=nThreads)
            pool.map(run_fspsq, args)
        else:
            map(run_fspsq, args)
    
    def reset(self):
        """Drop (delete) the library.
        
        The collection member is reinstantiated. This allows the collection
        to be freshly re-populated with models.
        """
        self.db.drop_collection(self.libname)
        self.collection = self.db[self.libname]
    
    def remove_existing_model(self, modelName):
        """Delete a model, by name, if it exists."""
        if self.collection.find_one({"_id": modelName}) is not None:
            self.collection.remove({"_id": modelName})

def run_fspsq(args):
    """A process for running fsps jobs."""
    # print "hello"
    libname, dbname, host, port, maxN, fspsqPath, commandPath = args
    thisHost = socket.gethostname() # hostname of current process
    
    # Connect to the library in MongoDB
    connection = pymongo.Connection(host=host, port=port)
    db = connection[dbname]
    db.add_son_manipulator(NumpySONManipulator())
    collection = db[libname]
    
    commonVarSets = _make_common_var_sets(collection)
    print "commonVarSets:", commonVarSets
    for varSet in commonVarSets:
        while True:
            psets = []
            modelNames = []
            now = datetime.datetime.utcnow()
            now = now.replace(tzinfo=pytz.utc)
            while len(psets) <= maxN:
                q = {"compute_complete": False, "compute_started": False}
                q.update(varSet)
                print "q", q
                doc = collection.find_and_modify(query=q,
                    update={"$set": {"compute_started": True,
                                     "queue_date": now,
                                     "compute_host": thisHost}},)
                print "doc", doc
                if doc is None: break # no available models
                modelName = str(doc['_id'])
                pset = ParameterSet(modelName, **doc['pset'])
                psets.append(pset)
                modelNames.append(pset.name)
            if len(psets) == 0: break # empty job queue
            # Startup a computation: write command file and start fspsq
            cmdTxt = "\n".join([p.command() for p in psets])+"\n"
            if os.path.exists(commandPath): os.remove(commandPath)
            f = open(commandPath, 'w')
            f.write(cmdTxt)
            f.close()
            shellCmd = _make_shell_command(fspsqPath, commandPath, varSet)
            print "cmd::", shellCmd
            subprocess.call(shellCmd, shell=True)
            # Get output data from FSPS
            _gather_fsps_outputs(collection, modelNames)

def _make_common_var_sets(c):
    """Make a list of common variable setups.
    
    .. note:: This merely lists all *possible* combinations of common variable
       configurations. There is not guarantee that there are models needing
       computation for each configuration set.
    
    Parameters
    ----------
    
    c : pymongo.Collection instance
    
    Returns
    -------
    
    List of dictionaries. Each dictionary has keys representing common
    variables: sfh, zmet, dust_type, imf_type, compute_vega_mags, redshift_colors
    """
    params = ['sfh', 'zmet', 'dust_type', 'imf_type',
        'compute_vega_mags', 'redshift_colors']
    possibleValues = {}
    for param in params:
        possibleValues[param] = c.distinct("pset."+param)
    groups = []
    for isfh in possibleValues['sfh']:
        for izmet in possibleValues['zmet']:
            for idust_type in possibleValues['dust_type']:
                for iimf_type in possibleValues['imf_type']:
                    for ivega in possibleValues['compute_vega_mags']:
                        for iredshift in possibleValues['redshift_colors']:
                            groups.append({"pset.sfh":isfh,"pset.zmet":izmet,
                                "pset.dust_type":idust_type,"pset.imf_type":iimf_type,
                                "pset.compute_vega_mags":ivega,
                                "pset.redshift_colors":iredshift})
    return groups

def _make_shell_command(fspsqPath, commandFilePath, varSet):
    """Crafts the command for running `fspsq`, returning a string."""
    params = ['sfh', 'zmet', 'dust_type', 'imf_type',
        'compute_vega_mags', 'redshift_colors']
    cmd = "%s %s %i %i %i %i %i %i" % (fspsqPath, commandFilePath,
        varSet['pset.sfh'], varSet['pset.zmet'], varSet['pset.dust_type'],
        varSet['pset.imf_type'], varSet['pset.compute_vega_mags'],
        varSet['pset.redshift_colors'])
    return cmd

def _gather_fsps_outputs(c, modelNames):
    """Adds the .mag and .spec outputs for each model to MongoDB.
    
    .. note:: Need to propagate the spec type through here...
    """
    fspsDir = os.environ['SPS_HOME'] # environment variable for FSPS
    outputDir = os.path.join(fspsDir, "OUTPUTS")
    print modelNames
    for modelName in modelNames:
        magPath = os.path.join(outputDir, modelName+".out.mags")
        magSpec = os.path.join(outputDir, modelName+".out.spec")
        npdata = ingest_output(magPath, magSpec)
        npDataSmall = npdata[:1]
        c.update({"_id": modelName}, {"$set": {"compute_complete": True}})
        print c.find_one({"_id": modelName})
        print "type:", type(npDataSmall)
        c.update({"_id": modelName}, {"$set": {"np_data": npDataSmall}})
        print c.find_one({"_id": modelName})['np_data']
        print "update complete!"

class ParameterSet(object):
    """An input parameter set for a FSPS model run."""
    def __init__(self, name, **kwargs):
        self.name = name # name of this model
        # Default values
        self.p = {"compute_vega_mags":0, "dust_type":0, "imf_type":0,
                "isoc_type":'pdva', "redshift_colors":0, "time_res_incr":2,
                "zred":0., "zmet":1, "sfh":0, "tau":1., "const":0.,
                "sf_start":0.,"tage":0., "fburst":0., "tburst":0., "imf1":1.3,
                "imf2":2.3, "imf3":2.3, "vdmc":0.08, "mdave":0.5,
                "dust_tesc":7., "dust1":0., "dust2":0., "dust_clumps":-99.,
                "frac_nodust":0., "dust_index":-0.7, "mwr":3.1,
                "uvb":1., "wgp1":1, "wgp2":1, "wgp3":0, "dell":0.,
                "delt":0., "sbss":0., "fbhb":0, "pagb":1.}
        self.knownParameters = self.p.keys()
        # Update values with user's arguments
        for k, v in kwargs.iteritems():
            if k in self.knownParameters:
                self.p[k] = v
    
    def command(self):
        """Write the string for this paramter set."""
        # These are pset variables, (aside from sfh and zmet)
        dt = [("zred","%.2f"),("tau","%.10f"),("const","%.4f"),
{'fraction_non_alphanumeric': 0.12250182964589315, 'fraction_numerical': 0.007037099589033384, 'mean_word_length': 2.243973703433163, 'pattern_counts': {'":': 121, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '150286', 'n_tokens_mistral': 5537, 'n_tokens_neox': 5393, 'n_words': 1370}