<commit_before>#!/usr/bin/python -tt

# omega - python client
# https://github.com/jfillmore/Omega-API-Engine
#
# Copyright 2011, Jonathon Fillmore
# Licensed under the MIT license. See LICENSE file.
# http://www.opensource.org/licenses/mit-license.php

"""Omega client for talking to an Omega server."""

import pycurl
import urllib
import httplib
try:
    import json
except:
    import simplejson
    json = simplejson

import re
import sys
import dbg
import StringIO
import os
import tempfile
import base64
import hashlib

from error import Exception
import util

class OmegaClient:
    """Client for talking to an Omega Server."""
    _version = '0.2'
    _http = None
    _hostname = None
    _folder = '/'
    _url = None
    _session_coookie = None
    _cookie_file = None
    _useragent = 'OmegaClient/0.2'

    def __init__(self, url = 'localhost', credentials = None, port = 5800, use_https = True):
        self._cookie_file = os.path.expanduser('~/.omega_cookie') # tempfile.NamedTemporaryFile()
        self.set_https(use_https)
        self.set_credentials(credentials)
        self.set_port(port)
        self.set_url(url)
        # TODO: python 2.7 supports an order tuple object we can use to preserve order :)
        self.encode = json.JSONEncoder().encode
        self.decode = json.JSONDecoder().decode
        # setup cookie jar

    def set_url(self, url):
        if url != '':
            self._url = url
            if url.find('/') == -1:
                self._hostname = url
                self._folder = ''
            else:
                # check for protocol
                if url.lower()[0:7] == 'http://':
                    url = url[8:]
                    self.set_https(false)
                elif url.lower()[0:8] == 'https://':
                    url = url[9:]
                    self.set_https(true)
                else:
                    # some other protocol perhaps?
                    if re.match('^\w+://', url):
                        raise Exception('Only HTTP and HTTPS are supported protocols.')
                # do we have a port?
                match = re.match('^(\w+\.)*\w+(:(\d{0,5}))/?', url)
                if match:
                    self.set_port(match.group(3))
                    # chop out the port
                    url = '/'.join((re.split(':\d+/?', url)))
                (self._hostname, self._folder) = url.split('/', 1)
                # force ourself to end in a slash
                if not self._folder.endswith('/'):
                    self._folder = ''.join((self._folder, '/'))
        else:
            raise Exception('Invalid API service URL: %s.' % url)
        # having set everything, init httplib
        if self._use_https:
            self._http = httplib.HTTPSConnection(self._hostname, self._port)
        else:
            self._http = httplib.HTTPConnection(self._hostname, self._port)
        self._http.cookies = False

    def get_url(self):
        return self._url

    def get_hostname(self):
        return self._hostname

    def get_folder(self):
        return self._folder

    def set_credentials(self, creds):
        if type(creds).__name__ == 'dict':
            # make sure we have a user/pass
            if 'username' in creds and 'password' in creds:
                self._credentials = creds
            elif 'token' in creds:
                self._credentials = creds
            else:
                raise Exception("Invalid credentials. Keys of 'username'/'password' or 'token' expected, but were not found.")
        elif creds == None:
            self._credentials = creds
        else:
            raise Exception("Invalid credentials. Keys of 'username'/'password' or 'token' expected, but were not found.")

    def set_https(self, secure = True):
        if secure:
            self._use_https = True
        else:
            self._use_https = False

    def set_port(self, port):
        if port >= 0 and port <= 65535:
            self._port = port
        else:
            raise Exception('Invalid API service port: %i.' % port)

    # old style API invoker
    def run(self, api, args = (), raw_response = False, full_response = False, get = None, post = None, files = None):
        # check and prep the data
        if api == '':
            raise Exception("Invalid service API: '%s'." %api)
        api = urllib.quote(api)
        curl = pycurl.Curl()
        data = [
            ('OMEGA_ENCODING', (curl.FORM_CONTENTS, 'json')),
            ('OMEGA_API_PARAMS', (curl.FORM_CONTENTS, self.encode(args)))
        ]
        if self._credentials:
            data.append(('OMEGA_CREDENTIALS', (curl.FORM_CONTENTS, self.encode(self._credentials))))
        # include any extra post data
        if post:
            (name, value) = post.split('=', 1)
            data.append((name, (curl.FORM_CONTENTS, value)))
        if files:
            # add in our files to the data
            for name in files:
                data.append((name, (curl.FORM_FILE, files[name])))
        # figure our our URL and get args
        url = self._url
        if self._use_https:
            url = ''.join(('https://', self._hostname))
        else:
            url = ''.join(('http://', self._hostname))
        url = '/'.join((':'.join((url, str(self._port))), self._folder))
        url = '/'.join((url, api))
        if get:
            url = '?'.join((url, get))
        # fire away
        curl.setopt(curl.URL, url)
        curl.setopt(curl.POST, 1)
        curl.setopt(curl.USERAGENT, self._useragent)
        curl.setopt(curl.COOKIEFILE, self._cookie_file)
        curl.setopt(curl.COOKIEJAR, self._cookie_file)
        if self._use_https:
            curl.setopt(curl.SSL_VERIFYPEER, 0) # TODO: don't always assume
            curl.setopt(curl.SSL_VERIFYHOST, 0) # TODO: don't always assume
        if data:
            curl.setopt(curl.HTTPPOST, data)
        else:
            curl.setopt(curl.POSTFIELDS, '&'.join(args))
        response = StringIO.StringIO()
        curl.setopt(curl.WRITEFUNCTION, response.write)
        curl.perform()
        response = response.getvalue()
        http_code = curl.getinfo(curl.HTTP_CODE)
        content_type = curl.getinfo(curl.CONTENT_TYPE) or "";
        if http_code < 200 or http_code >= 300:
            # see if we got json data back
            try:
                if content_type.startswith("application/json"):
                    decoded = self.decode(response)
                    if 'reason' in decoded:
                        error = decoded['reason']
                    else:
                        error = response
                else:
                    error = response
            except:
                error = response
            raise Exception("Server returned HTTP code %s. Response:\n%s" %
                (str(http_code), str(error)))
        curl.close()
        if raw_response:
            return response
        else:
            # decode the response and check whether or not it was successful
            # TODO: check response encoding in header
            try:
                if content_type.startswith("application/json"):
                    response = self.decode(response)
                else:
                    return response
            except:
                raise Exception('Failed to decode API response.', response)
            # check to see if our API call was successful
            if 'result' in response and response['result'] == False:
                if 'reason' in response:
                    if full_response:
                        raise Exception('API "%s" failed.\n%s' %
                            (urllib.unquote(api), dbg.obj2str(response)))
                    else:
                        raise Exception(response['reason'])
                else:
                    raise Exception('API "%s" failed, but did not provide an explanation. Response: %s' % (api, response))
            else:
                if full_response:
                    return response
                else:
                    if 'data' in response:
                        return response['data']
                    else:
                        return None

    def get(self, api, params, opts):
        return self.request('GET', api, params, opts);

    def post(self, api, params, opts = {}):
        return self.request('POST', api, params, opts);

    def put(self, api, params, opts = {}):
        return self.request('PUT', api, params, opts);

    def delete(self, api, params, opts = {}):
        return self.request('DELETE', api, params, opts);

    def request(self, method, api, params = (), raw_response = False, full_response = False, get = None, headers = {}, verbose = False):
        '''New REST-friendly API invoker'''
        # check and prep the data
        if method is None or method == '':
            method = 'GET'
        method = method.upper()
        if api == '':
            raise Exception("Invalid service API: '%s'." %api)
        api = urllib.quote(api)
        if self._credentials:
            creds = self._credentials
            md5 = hashlib.md5();
            md5.update(':'.join(
                [creds['username'], creds['password']]
            ))
            headers['Authentication'] = 'Basic ' + base64.b64encode(
                md5.hexdigest())
        http = self._http
        # figure our our URL and get args
        url = self._url
        headers['Content-type'] = 'application/json'
        headers['Accept'] = 'application/json'
        if http.cookies:
            headers['Cookie'] = http.cookies;

        url = '/'.join(('', self._folder, api))
        if get:
            url = '?'.join((url, get))
        if method == 'GET':
            url = '?'.join((url, '&'.join([
                '='.join((urllib.quote(name), urllib.quote(str(params[name])))) for name in params
            ])))
            data = None
        else:
            data = self.encode(params)
        # fire away
        if verbose:
            if self._use_https:
                proto = 'https'
            else:
                proto = 'http'
            sys.stderr.write(
                '+ %s %s://%s:%d/%s, params: "%s", headers: "%s"\n' %
                ((method, proto, self._hostname, self._port, url, data, str(headers))))
        http.request(method, url, data, headers)
        response = http.getresponse()
        # see if we get a cookie back
        response_headers = str(response.msg).split('\n');
        cookies = [c.split(': ')[1].split('; ')[0] for c in response_headers if c.startswith('Set-Cookie: ')]
        if cookies:
            # note that we ignore the path
            if verbose:
                for cookie in cookies:
                    sys.stderr.write('+ Cookie [%s]\n' % (cookie))
            http.cookies = cookies
        if verbose:
            sys.stderr.write(
                '+ Status [%d], Reason [%s], headers [%s]\n' %
                (response.status, response.reason, response.msg)
            )
        content_type = response.getheader('Content-Type') or '';
        response_data = response.read();
        # handle any errors based on status code
        if response.status < 200 or response.status >= 300:
            if content_type.startswith("application/json"):
                result = self.decode(response_data)
                if 'reason' in result:
                    error = result['reason']
                else:
                    error = 'An unknown error has occurred.'
            else:
                error = response_data
            raise Exception('API "%s" failed: %s' % (api, error))
        # return a raw response if needed; otherwise decode if JSON
        if raw_response or not content_type.startswith("application/json"):
            return response_data
        try:
            result = self.decode(response_data)
        except:
            raise Exception('Failed to decode API result:\n' + response_data)
        # check to see if our API call was successful
        if 'result' in result and result['result'] == False:
            if 'reason' in result:
                if full_response:
                    raise Exception('"%s" failed (%d %s):\n%s' %
                        (urllib.unquote(api), response.status, response.reason, dbg.obj2str(result)))
                else:
                    raise Exception(result['reason'])
            else:
                raise Exception('API "%s" failed: %s' % (api, result))
        else:
            if full_response:
                return result
            else:
                if 'data' in result:
                    return result['data']
                else:
                    return None

if __name__ == '__main__':
    import dbg
    dbg.pretty_print(OmegaClient());
<commit_msg>better response of errors w/ "full response" opt<commit_after>#!/usr/bin/python -tt

# omega - python client
# https://github.com/jfillmore/Omega-API-Engine
#
# Copyright 2011, Jonathon Fillmore
# Licensed under the MIT license. See LICENSE file.
# http://www.opensource.org/licenses/mit-license.php

"""Omega client for talking to an Omega server."""

import pycurl
import urllib
import httplib
try:
    import json
except:
    import simplejson
    json = simplejson

import re
import sys
import dbg
import StringIO
import os
import tempfile
import base64
import hashlib

from error import Exception
import util

class OmegaClient:
    """Client for talking to an Omega Server."""
    _version = '0.2'
    _http = None
    _hostname = None
    _folder = '/'
    _url = None
    _session_coookie = None
    _cookie_file = None
    _useragent = 'OmegaClient/0.2'

    def __init__(self, url = 'localhost', credentials = None, port = 5800, use_https = True):
        self._cookie_file = os.path.expanduser('~/.omega_cookie') # tempfile.NamedTemporaryFile()
        self.set_https(use_https)
        self.set_credentials(credentials)
        self.set_port(port)
        self.set_url(url)
        # TODO: python 2.7 supports an order tuple object we can use to preserve order :)
        self.encode = json.JSONEncoder().encode
        self.decode = json.JSONDecoder().decode
        # setup cookie jar

    def set_url(self, url):
        if url != '':
            self._url = url
            if url.find('/') == -1:
                self._hostname = url
                self._folder = ''
            else:
                # check for protocol
                if url.lower()[0:7] == 'http://':
                    url = url[8:]
                    self.set_https(false)
                elif url.lower()[0:8] == 'https://':
                    url = url[9:]
                    self.set_https(true)
                else:
                    # some other protocol perhaps?
                    if re.match('^\w+://', url):
                        raise Exception('Only HTTP and HTTPS are supported protocols.')
                # do we have a port?
                match = re.match('^(\w+\.)*\w+(:(\d{0,5}))/?', url)
                if match:
                    self.set_port(match.group(3))
                    # chop out the port
                    url = '/'.join((re.split(':\d+/?', url)))
                (self._hostname, self._folder) = url.split('/', 1)
                # force ourself to end in a slash
                if not self._folder.endswith('/'):
                    self._folder = ''.join((self._folder, '/'))
        else:
            raise Exception('Invalid API service URL: %s.' % url)
        # having set everything, init httplib
        if self._use_https:
            self._http = httplib.HTTPSConnection(self._hostname, self._port)
        else:
            self._http = httplib.HTTPConnection(self._hostname, self._port)
        self._http.cookies = False

    def get_url(self):
        return self._url

    def get_hostname(self):
        return self._hostname

    def get_folder(self):
        return self._folder

    def set_credentials(self, creds):
        if type(creds).__name__ == 'dict':
            # make sure we have a user/pass
            if 'username' in creds and 'password' in creds:
                self._credentials = creds
            elif 'token' in creds:
                self._credentials = creds
            else:
                raise Exception("Invalid credentials. Keys of 'username'/'password' or 'token' expected, but were not found.")
        elif creds == None:
            self._credentials = creds
        else:
            raise Exception("Invalid credentials. Keys of 'username'/'password' or 'token' expected, but were not found.")

    def set_https(self, secure = True):
        if secure:
            self._use_https = True
        else:
            self._use_https = False

    def set_port(self, port):
        if port >= 0 and port <= 65535:
            self._port = port
        else:
            raise Exception('Invalid API service port: %i.' % port)

    # old style API invoker
    def run(self, api, args = (), raw_response = False, full_response = False, get = None, post = None, files = None):
        # check and prep the data
        if api == '':
            raise Exception("Invalid service API: '%s'." %api)
        api = urllib.quote(api)
        curl = pycurl.Curl()
        data = [
            ('OMEGA_ENCODING', (curl.FORM_CONTENTS, 'json')),
            ('OMEGA_API_PARAMS', (curl.FORM_CONTENTS, self.encode(args)))
        ]
        if self._credentials:
            data.append(('OMEGA_CREDENTIALS', (curl.FORM_CONTENTS, self.encode(self._credentials))))
        # include any extra post data
        if post:
            (name, value) = post.split('=', 1)
            data.append((name, (curl.FORM_CONTENTS, value)))
        if files:
            # add in our files to the data
            for name in files:
                data.append((name, (curl.FORM_FILE, files[name])))
        # figure our our URL and get args
        url = self._url
        if self._use_https:
            url = ''.join(('https://', self._hostname))
        else:
            url = ''.join(('http://', self._hostname))
        url = '/'.join((':'.join((url, str(self._port))), self._folder))
        url = '/'.join((url, api))
        if get:
            url = '?'.join((url, get))
        # fire away
        curl.setopt(curl.URL, url)
        curl.setopt(curl.POST, 1)
        curl.setopt(curl.USERAGENT, self._useragent)
        curl.setopt(curl.COOKIEFILE, self._cookie_file)
        curl.setopt(curl.COOKIEJAR, self._cookie_file)
        if self._use_https:
            curl.setopt(curl.SSL_VERIFYPEER, 0) # TODO: don't always assume
            curl.setopt(curl.SSL_VERIFYHOST, 0) # TODO: don't always assume
        if data:
            curl.setopt(curl.HTTPPOST, data)
        else:
            curl.setopt(curl.POSTFIELDS, '&'.join(args))
        response = StringIO.StringIO()
        curl.setopt(curl.WRITEFUNCTION, response.write)
        curl.perform()
        response = response.getvalue()
        http_code = curl.getinfo(curl.HTTP_CODE)
        content_type = curl.getinfo(curl.CONTENT_TYPE) or "";
        if http_code < 200 or http_code >= 300:
            # see if we got json data back
            try:
                if content_type.startswith("application/json"):
                    decoded = self.decode(response)
                    if 'reason' in decoded:
                        error = decoded['reason']
                    else:
                        error = response
                else:
                    error = response
            except:
                error = response
            raise Exception("Server returned HTTP code %s. Response:\n%s" %
                (str(http_code), str(error)))
        curl.close()
        if raw_response:
            return response
        else:
            # decode the response and check whether or not it was successful
            # TODO: check response encoding in header
            try:
                if content_type.startswith("application/json"):
                    response = self.decode(response)
                else:
                    return response
            except:
                raise Exception('Failed to decode API response.', response)
            # check to see if our API call was successful
            if 'result' in response and response['result'] == False:
                if 'reason' in response:
                    if full_response:
                        raise Exception('API "%s" failed.\n%s' %
                            (urllib.unquote(api), dbg.obj2str(response)))
                    else:
                        raise Exception(response['reason'])
                else:
                    raise Exception('API "%s" failed, but did not provide an explanation. Response: %s' % (api, response))
            else:
                if full_response:
                    return response
                else:
                    if 'data' in response:
                        return response['data']
                    else:
                        return None

    def get(self, api, params, opts):
        return self.request('GET', api, params, opts);

    def post(self, api, params, opts = {}):
        return self.request('POST', api, params, opts);

    def put(self, api, params, opts = {}):
        return self.request('PUT', api, params, opts);

    def delete(self, api, params, opts = {}):
        return self.request('DELETE', api, params, opts);

    def request(self, method, api, params = (), raw_response = False, full_response = False, get = None, headers = {}, verbose = False):
        '''New REST-friendly API invoker'''
        # check and prep the data
        if method is None or method == '':
            method = 'GET'
        method = method.upper()
        if api == '':
            raise Exception("Invalid service API: '%s'." %api)
        api = urllib.quote(api)
        if self._credentials:
            creds = self._credentials
            md5 = hashlib.md5();
            md5.update(':'.join(
                [creds['username'], creds['password']]
            ))
            headers['Authentication'] = 'Basic ' + base64.b64encode(
                md5.hexdigest())
        http = self._http
        # figure our our URL and get args
        url = self._url
        headers['Content-type'] = 'application/json'
        headers['Accept'] = 'application/json'
        if http.cookies:
            headers['Cookie'] = http.cookies;

        url = '/'.join(('', self._folder, api))
        if get:
            url = '?'.join((url, get))
        if method == 'GET':
            url = '?'.join((url, '&'.join([
                '='.join((urllib.quote(name), urllib.quote(str(params[name])))) for name in params
            ])))
            data = None
        else:
            data = self.encode(params)
        # fire away
        if verbose:
            if self._use_https:
                proto = 'https'
            else:
                proto = 'http'
            sys.stderr.write(
                '+ %s %s://%s:%d/%s, params: "%s", headers: "%s"\n' %
                ((method, proto, self._hostname, self._port, url, data, str(headers))))
        http.request(method, url, data, headers)
        response = http.getresponse()
        # see if we get a cookie back
        response_headers = str(response.msg).split('\n');
        cookies = [c.split(': ')[1].split('; ')[0] for c in response_headers if c.startswith('Set-Cookie: ')]
        if cookies:
            # note that we ignore the path
            if verbose:
                for cookie in cookies:
                    sys.stderr.write('+ Cookie [%s]\n' % (cookie))
            http.cookies = cookies
        if verbose:
            sys.stderr.write(
                '+ Status [%d], Reason [%s], headers [%s]\n' %
                (response.status, response.reason, response.msg)
            )
        content_type = response.getheader('Content-Type') or '';
        response_data = response.read();
        # handle any errors based on status code
        if response.status < 200 or response.status >= 300:
            if content_type.startswith("application/json"):
                result = self.decode(response_data)
                if 'reason' in result:
                    error = result['reason']
                else:
                    error = 'An unknown error has occurred.'
            else:
                error = response_data
            if full_response:
                raise Exception('API "%s" failed (%d %s):\n%s' %
                    (urllib.unquote(api), response.status, response.reason, dbg.obj2str(result)))
            else:
                raise Exception('API "%s" failed: %s' % (api, error))
        # return a raw response if needed; otherwise decode if JSON
        if raw_response or not content_type.startswith("application/json"):
            return response_data
        try:
            result = self.decode(response_data)
        except:
            raise Exception('Failed to decode API result:\n' + response_data)
        # check to see if our API call was successful
        if 'result' in result and result['result'] == False:
            if 'reason' in result:
                if full_response:
                    raise Exception('"%s" failed (%d %s):\n%s' %
                        (urllib.unquote(api), response.status, response.reason, dbg.obj2str(result)))
                else:
                    raise Exception(result['reason'])
            else:
                raise Exception('API "%s" failed: %s' % (api, result))
        else:
            if full_response:
                return result
            else:
                if 'data' in result:
                    return result['data']
                else:
                    return None

if __name__ == '__main__':
    import dbg
    dbg.pretty_print(OmegaClient());
<|endoftext|>


<commit_before>
class Sphere():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawSphere('+str(self.mv)+','+self.color+');'
        return return_string

class Plane():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawPlane('+str(self.mv)+','+self.color+');'
        return return_string

class Line():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawLine('+str(self.mv)+','+self.color+');'
        return return_string

class Circle():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawCircle('+str(self.mv)+','+self.color+');'
        return return_string

class PointPair():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawPointPair('+str(self.mv)+','+self.color+');'
        return return_string

class EucPoint():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawEucPoint('+str(self.mv)+','+self.color+');'
        return return_string

class GAScene():
    def __init__(self):
        self.objects = []
    def add(self,object_added):
        self.objects.append(object_added)
    def concat(self,other):
        if isinstance(other,GAScene):
            for obj in other.objects:
                self.add(obj)
        else:
            raise ValueError('Other object must be a GAScene')
    def clear(self):
        self.objects = []
    def __repr__(self):
        s = ''
        for obj in self.objects:
            s += obj.gen_string() + '\n'
        s = s[:-1]
        return s
    def add_euc_point(self,mv,color):
        self.add(EucPoint(mv,color))
    def add_circle(self,mv,color):
        self.add(Circle(mv,color))
    def add_line(self,mv,color):
        self.add(Line(mv,color))
    def add_plane(self,mv,color):
        self.add(Plane(mv,color))
    def add_sphere(self,mv,color):
        self.add(Sphere(mv,color))
    def add_point_pair(self,mv,color):
        self.add(PointPair(mv,color))
    def add_object(self,mv,mv_type,color):
        if mv_type == 'line':
            self.add_line(mv, color)
        elif mv_type == 'point_pair':
            self.add_point_pair(mv, color)
        elif mv_type == 'circle':
            self.add_circle(mv, color)
        elif mv_type == 'plane':
            self.add_plane(mv, color)
        elif mv_type == 'sphere':
            self.add_sphere(mv, color)
        else:
            raise ValueError(str(mv_type) + ' is not a valid mv_type. You must specify a valid mv_type.')
    def save_to_file(self,filename):
        with open(filename,'w') as fobj:
            print(self,file=fobj)

class GAAnimation():
    def __init__(self):
        self.scenes = []
    def add_scene(self,scene):
        self.scenes.append(scene)
    def __repr__(self):
        s = ''
        for obj in self.scenes:
            s += str(obj).replace("\n", "#") + '\n'
        s = s[:-1]
        return s
    def save_to_file(self,filename):
        with open(filename,'w') as fobj:
            print(self,file=fobj)<commit_msg>print statement incompatibility python 2 to 3 attempted fix<commit_after>
from __future__ import print_function

class Sphere():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawSphere('+str(self.mv)+','+self.color+');'
        return return_string

class Plane():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawPlane('+str(self.mv)+','+self.color+');'
        return return_string

class Line():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawLine('+str(self.mv)+','+self.color+');'
        return return_string

class Circle():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawCircle('+str(self.mv)+','+self.color+');'
        return return_string

class PointPair():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawPointPair('+str(self.mv)+','+self.color+');'
        return return_string

class EucPoint():
    def __init__(self,mv,color):
        self.mv = mv
        self.color = color
    def gen_string(self):
        return_string = 'DrawEucPoint('+str(self.mv)+','+self.color+');'
        return return_string

class GAScene():
    def __init__(self):
        self.objects = []
    def add(self,object_added):
        self.objects.append(object_added)
    def concat(self,other):
        if isinstance(other,GAScene):
            for obj in other.objects:
                self.add(obj)
        else:
            raise ValueError('Other object must be a GAScene')
    def clear(self):
        self.objects = []
    def __repr__(self):
        s = ''
        for obj in self.objects:
            s += obj.gen_string() + '\n'
        s = s[:-1]
        return s
    def add_euc_point(self,mv,color):
        self.add(EucPoint(mv,color))
    def add_circle(self,mv,color):
        self.add(Circle(mv,color))
    def add_line(self,mv,color):
        self.add(Line(mv,color))
    def add_plane(self,mv,color):
        self.add(Plane(mv,color))
    def add_sphere(self,mv,color):
        self.add(Sphere(mv,color))
    def add_point_pair(self,mv,color):
        self.add(PointPair(mv,color))
    def add_object(self,mv,mv_type,color):
        if mv_type == 'line':
            self.add_line(mv, color)
        elif mv_type == 'point_pair':
            self.add_point_pair(mv, color)
        elif mv_type == 'circle':
            self.add_circle(mv, color)
        elif mv_type == 'plane':
            self.add_plane(mv, color)
        elif mv_type == 'sphere':
            self.add_sphere(mv, color)
        else:
            raise ValueError(str(mv_type) + ' is not a valid mv_type. You must specify a valid mv_type.')
    def save_to_file(self,filename):
        with open(filename,'w') as fobj:
            print(self,file=fobj)

class GAAnimation():
    def __init__(self):
        self.scenes = []
    def add_scene(self,scene):
        self.scenes.append(scene)
    def __repr__(self):
        s = ''
        for obj in self.scenes:
            s += str(obj).replace("\n", "#") + '\n'
        s = s[:-1]
        return s
    def save_to_file(self,filename):
        with open(filename,'w') as fobj:
            print(self,file=fobj)<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
"""
File: cloudforest_mpi.py
Author: Brant Faircloth

Created by Brant Faircloth on 12 May 2012 10:05 PDT (-0700)
Copyright (c) 2012 Brant C. Faircloth. All rights reserved.

Description: Compute genetrees and bootstrap replicates from a
directory of phylip files input on the CLI using MPI.

"""

import os
import sys
import glob
import argparse
import numpy as np
from core import Phyml, is_dir, is_file, FullPaths, phylip_to_oneliner, split_oneliner, oneliner_to_array, get_bootstraps, array_to_oneliner, oneliner_to_phylip

#import pdb

def get_args():
    """Get arguments from CLI"""
    parser = argparse.ArgumentParser(
            description="""Program description""")
    parser.add_argument(
            "input",
            type=is_dir,
            action=FullPaths,
            help="""The path to a directory of PHYLIP-formatted alignments."""
        )
    parser.add_argument(
            "output",
            type=is_dir,
            action=FullPaths,
            help="""The path to a directory in which to store the output.""",
        )
    parser.add_argument(
            "run",
            choices=['genetrees', 'bootstraps', 'both'],
            default=None,
            help="""The type of analysis to run.""",
        )
    parser.add_argument(
            "phyml",
            action=FullPaths,
            help="""The path (including executable name) to PhyML on your platform.""",
        )
    parser.add_argument(
            "--genetrees",
            action=FullPaths,
            type=is_file,
            default=None,
            help="""The path to the genetrees output by a previous run.  Needed only when bootstrapping a prior run.""",
        )
    parser.add_argument(
            "--bootreps",
            type=int,
            default=100,
            help="""The number of bootstrap replicates to run.""",
        )
    args = parser.parse_args()

    if args.run == 'bootstraps' and args.genetrees is None:
        sys.exit("\nIf runnning only boostraps, you must pass a genetrees file")

    return args

def make_tree_name(args_dict):
    """Converts dictionary of argument into a sorted string with the
    format:  """
    # join is faster and cleaner than concatenate
    # list comp is faster than for
    name = ','.join(['='.join([pair[0], pair[1]]) for pair in sorted(args_dict.items())])
    return name

def generate_bootreps(bootreps, phyml, oneliners):
    """Replicate the data set bootrep numer of times, prior to bootstrapping"""
    # keep bootrep numbers indexed by one
    for i in xrange(1, bootreps + 1):
        yield i, phyml, oneliners

def get_models_from_genetrees(genetrees):
    """Iterate over a set of precomputed genetrees and return the models in metadata"""
    # iterate over genetrees to grab models
    models = {}
    for line, tree in enumerate(open(genetrees, 'rU')):
        ts = tree.split(' ')[1].strip("'")
        for item in ts.split(','):
            if item.startswith('chrm'):
                locus = item.split('=')[1]
            elif item.startswith('model'):
                model = item.split('=')[1]
        models[locus] = model
    assert len(models) == line + 1, "Not all loci have a model"
    return models

def get_bootstrap_replicates(multilocus_bstrap):
    """Boostrap the bases of all alignments in a resampled population of alignments"""
    for locus in multilocus_bstrap:
        # split keys from alignments
        args_dict, locus = split_oneliner(locus)
        # convert alignments to arrays
        taxa, numpy_alignment = oneliner_to_array(locus.rstrip(';'))
        # transpose so we bootstrap by columns
        bases_by_col = np.column_stack(numpy_alignment)
        # bootstrap by columns
        shuffled = get_bootstraps(bases_by_col)
        # transpose back to rows of sequences
        shuffled = np.column_stack(shuffled)
        # convert back to oneliner
        oneliner = array_to_oneliner(taxa, shuffled)
        oneliner = "%s:%s" % (make_tree_name(args_dict), oneliner)
        yield oneliner

def genetree_worker(params):
    """Worker function to compute genetrees for individual loci"""
    locus, fullpth = params
    name, phylip = locus
    pth, exe = os.path.split(fullpth)
    args_dict = {}
    phyml = Phyml(phylip, pth=pth, exe=exe)
    model, tree = phyml.best_aicc_model_and_tree()
    args_dict['chrm'] = name
    args_dict['model'] = model
    tree = "tree '%s' = [&U] %s" % (make_tree_name(args_dict), tree)
    return (name, model, tree)

def bootstrap_worker(params):
    """Worker function to compute boostrap replicates of datasets and indiv. loci"""
    bootstrap_trees = []
    rep, fullpth, oneliners = params
    pth, exe = os.path.split(fullpth)
    # first, resample w/ replacement/bootstrap across loci
    multilocus_bstrap = get_bootstraps(oneliners)
    # Resample w/ replacement/boostrap bases within loci
    bootstraps = get_bootstrap_replicates(multilocus_bstrap)
    for oneliner in bootstraps:
        args_dict, locus = split_oneliner(oneliner, default_model=True)
        phylip = oneliner_to_phylip(locus)
        phyml = Phyml(phylip, pth=pth, exe=exe)
        # run phyml.  if no model, defaults to GTR
        # TOOD: Why do we need LnL?
        args_dict['lnL'], tree = phyml.run(args_dict['model'])
        #bootstrap_trees.append("tree '%s' = [&U] %s" % (make_tree_name(args_dict), tree))
        bootstrap_trees.append('''%s\t"%s"''' % (rep, tree))
    return bootstrap_trees

def boostrap_all_loci(args, models, alns):
    """Compute trees from bootstrap replicates of a dataset"""
    oneliners = [phylip_to_oneliner(phylip, locus, models[locus]) \
                for locus, phylip in alns.iteritems()]
    # for every rep in boostraps, map loci onto worker that will
    # bootstrap, run phyml, and return bootstrap trees
    params = generate_bootreps(args.bootreps, args.phyml, oneliners)
    bootreps = mmap(bootstrap_worker, params)
    # write
    outname = "%s-bootreps.tree" % (args.bootreps)
    outf = open(os.path.join(args.output, outname), 'w')
    for bootrep in bootreps:
        for tree in bootrep:
            outf.write("%s\n" % (tree))

def main():
    args = get_args()
    # read in phylip files
    alns = {}
    for f in glob.glob(os.path.join(args.input, '*.phy*')):
        alns[os.path.splitext(os.path.basename(f))[0]] = open(f, 'rU').read()

    # replicate our options for passing to map()
    opts = [args.phyml for i in range(len(alns))]
    # compute genetrees
    if args.run == 'genetrees' or args.run == 'both':
        params = zip(alns.items(), opts)
        genetrees = mmap(genetree_worker, params)
        # write genetrees to output file
        outf = open(os.path.join(args.output, 'genetrees.tre'), 'w')
        for tree in genetrees:
            outf.write("%s\n" % (tree[2]))
        outf.close()
    # compute bootreps on genetrees from above
    if args.run == 'both':
        # get models for each locus based on genetrees in-memory
        models = dict([[tree[0], tree[1]] for tree in genetrees])
        boostrap_all_loci(args, models, alns)
    # compute bootreps on genetrees from a file
    if args.run == 'bootstraps':
        # get models for each locus based on genetrees in genetree file
        models = get_models_from_genetrees(args.genetrees)
        boostrap_all_loci(args, models, alns)

if __name__ == '__main__':
    from deap.dtm import map as mmap
    from deap.dtm import start
    start(main)

    #from multiprocessing import Pool
    #pool = Pool(4)
    #mmap = pool.map

    #mmap = map
    #main()

<commit_msg>adding diff. parallelism and stdout notices<commit_after>#!/usr/bin/env python
# encoding: utf-8
"""
File: cloudforest_mpi.py
Author: Brant Faircloth

Created by Brant Faircloth on 12 May 2012 10:05 PDT (-0700)
Copyright (c) 2012 Brant C. Faircloth. All rights reserved.

Description: Compute genetrees and bootstrap replicates from a
directory of phylip files input on the CLI using MPI.

"""

import os
import sys
import glob
import argparse
import numpy as np
from core import Phyml, is_dir, is_file, FullPaths, phylip_to_oneliner, split_oneliner, oneliner_to_array, get_bootstraps, array_to_oneliner, oneliner_to_phylip

#import pdb

def get_args():
    """Get arguments from CLI"""
    parser = argparse.ArgumentParser(
            description="""Program description""")
    parser.add_argument(
            "input",
            type=is_dir,
            action=FullPaths,
            help="""The path to a directory of PHYLIP-formatted alignments."""
        )
    parser.add_argument(
            "output",
            type=is_dir,
            action=FullPaths,
            help="""The path to a directory in which to store the output.""",
        )
    parser.add_argument(
            "run",
            choices=['genetrees', 'bootstraps', 'both'],
            default=None,
            help="""The type of analysis to run.""",
        )
    parser.add_argument(
            "phyml",
            action=FullPaths,
            help="""The path (including executable name) to PhyML on your platform.""",
        )
    parser.add_argument(
            "--genetrees",
            action=FullPaths,
            type=is_file,
            default=None,
            help="""The path to the genetrees output by a previous run.  Needed only when bootstrapping a prior run.""",
        )
    parser.add_argument(
            "--bootreps",
            type=int,
            default=100,
            help="""The number of bootstrap replicates to run.""",
        )
    parser.add_argument(
            "--parallelism",
            choices=['mpi', 'multiprocessing', 'single'],
            default='mpi',
            help="""The type of parallelism to use.""",
        )

    args = parser.parse_args()

    if args.run == 'bootstraps' and args.genetrees is None:
        sys.exit("\nIf runnning only boostraps, you must pass a genetrees file")

    return args

def make_tree_name(args_dict):
    """Converts dictionary of argument into a sorted string with the
    format:  """
    # join is faster and cleaner than concatenate
    # list comp is faster than for
    name = ','.join(['='.join([pair[0], pair[1]]) for pair in sorted(args_dict.items())])
    return name

def generate_bootreps(bootreps, phyml, oneliners):
    """Replicate the data set bootrep numer of times, prior to bootstrapping"""
    # keep bootrep numbers indexed by one
    for i in xrange(1, bootreps + 1):
        yield i, phyml, oneliners

def get_models_from_genetrees(genetrees):
    """Iterate over a set of precomputed genetrees and return the models in metadata"""
    # iterate over genetrees to grab models
    models = {}
    for line, tree in enumerate(open(genetrees, 'rU')):
        ts = tree.split(' ')[1].strip("'")
        for item in ts.split(','):
            if item.startswith('chrm'):
                locus = item.split('=')[1]
            elif item.startswith('model'):
                model = item.split('=')[1]
        models[locus] = model
    assert len(models) == line + 1, "Not all loci have a model"
    return models

def get_bootstrap_replicates(multilocus_bstrap):
    """Boostrap the bases of all alignments in a resampled population of alignments"""
    for locus in multilocus_bstrap:
        # split keys from alignments
        args_dict, locus = split_oneliner(locus)
        # convert alignments to arrays
        taxa, numpy_alignment = oneliner_to_array(locus.rstrip(';'))
        # transpose so we bootstrap by columns
        bases_by_col = np.column_stack(numpy_alignment)
        # bootstrap by columns
        shuffled = get_bootstraps(bases_by_col)
        # transpose back to rows of sequences
        shuffled = np.column_stack(shuffled)
        # convert back to oneliner
        oneliner = array_to_oneliner(taxa, shuffled)
        oneliner = "%s:%s" % (make_tree_name(args_dict), oneliner)
        yield oneliner

def genetree_worker(params):
    """Worker function to compute genetrees for individual loci"""
    locus, fullpth = params
    name, phylip = locus
    pth, exe = os.path.split(fullpth)
    args_dict = {}
    phyml = Phyml(phylip, pth=pth, exe=exe)
    model, tree = phyml.best_aicc_model_and_tree()
    args_dict['chrm'] = name
    args_dict['model'] = model
    tree = "tree '%s' = [&U] %s" % (make_tree_name(args_dict), tree)
    return (name, model, tree)

def bootstrap_worker(params):
    """Worker function to compute boostrap replicates of datasets and indiv. loci"""
    bootstrap_trees = []
    rep, fullpth, oneliners = params
    pth, exe = os.path.split(fullpth)
    # first, resample w/ replacement/bootstrap across loci
    multilocus_bstrap = get_bootstraps(oneliners)
    # Resample w/ replacement/boostrap bases within loci
    bootstraps = get_bootstrap_replicates(multilocus_bstrap)
    for oneliner in bootstraps:
        args_dict, locus = split_oneliner(oneliner, default_model=True)
        phylip = oneliner_to_phylip(locus)
        phyml = Phyml(phylip, pth=pth, exe=exe)
        # run phyml.  if no model, defaults to GTR
        # TOOD: Why do we need LnL?
        args_dict['lnL'], tree = phyml.run(args_dict['model'])
        #bootstrap_trees.append("tree '%s' = [&U] %s" % (make_tree_name(args_dict), tree))
        bootstrap_trees.append('''%s\t"%s"''' % (rep, tree))
    return bootstrap_trees

def boostrap_all_loci(args, models, alns):
    """Compute trees from bootstrap replicates of a dataset"""
    oneliners = [phylip_to_oneliner(phylip, locus, models[locus]) \
                for locus, phylip in alns.iteritems()]
    # for every rep in boostraps, map loci onto worker that will
    # bootstrap, run phyml, and return bootstrap trees
    params = generate_bootreps(args.bootreps, args.phyml, oneliners)
    bootreps = mmap(bootstrap_worker, params)
    # write
    outname = "%s-bootreps.tree" % (args.bootreps)
    outf = open(os.path.join(args.output, outname), 'w')
    for bootrep in bootreps:
        for tree in bootrep:
            outf.write("%s\n" % (tree))

def main():
    args = get_args()
    # read in phylip files
    alns = {}
    for f in glob.glob(os.path.join(args.input, '*.phy*')):
        alns[os.path.splitext(os.path.basename(f))[0]] = open(f, 'rU').read()
    # replicate our options for passing to map()
    opts = [args.phyml for i in range(len(alns))]
    # compute genetrees
    if args.run == 'genetrees' or args.run == 'both':
        sys.stdout.write("Running genetrees...\n")
        params = zip(alns.items(), opts)
        genetrees = mmap(genetree_worker, params)
        # write genetrees to output file
        outf = open(os.path.join(args.output, 'genetrees.tre'), 'w')
        for tree in genetrees:
            outf.write("%s\n" % (tree[2]))
        outf.close()
    # compute bootreps on genetrees from above
    if args.run == 'both':
        sys.stdout.write("Running boostraps of genetrees...\n")
        # get models for each locus based on genetrees in-memory
        models = dict([[tree[0], tree[1]] for tree in genetrees])
        boostrap_all_loci(args, models, alns)
    # compute bootreps on genetrees from a file
    if args.run == 'bootstraps':
        sys.stdout.write("Running boostraps...\n")
        # get models for each locus based on genetrees in genetree file
        models = get_models_from_genetrees(args.genetrees)
        boostrap_all_loci(args, models, alns)

if __name__ == '__main__':
    args = get_args()
    if args.parallelism == 'mpi':
        from deap.dtm import map as mmap
        from deap.dtm import start
        start(main)
    elif args.parallelism == 'multiprocessing':
        from multiprocessing import Pool
        pool = Pool(7)
        mmap = pool.map
        main()
    elif args.parallelism == 'single':
        mmap = map
        main()
<|endoftext|>


<commit_before>from cloudmesh_install import config_file
from cloudmesh_install.util import path_expand
from cloudmesh_common.logger import LOGGER
from cloudmesh.util.config import read_yaml_config
from cloudmesh.util.config import ordered_dump
from collections import OrderedDict
import simplejson
from pprint import pprint
import yaml
import json
import os
import stat
import sys

log = LOGGER(__file__)
package_dir = os.path.dirname(os.path.abspath(__file__))

class OrderedJsonEncoder(simplejson.JSONEncoder):

    indent = 4

    def encode(self, o, depth=0):
        if isinstance(o, OrderedDict):
            return "{" + (",\n ").join([self.encode(k) + ":" + \
                                        self.encode(v, depth + 1) \
                                        for (k, v) in o.iteritems()]) + "}\n"
        else:
            return simplejson.JSONEncoder.encode(self, o)

def custom_print(data_structure, indent):
    for key, value in data_structure.items():
        print "\n%s%s:" % ('    ' * indent, str(key)),
        if isinstance(value, OrderedDict):
            custom_print(value, indent + 1)
        elif isinstance(value, dict):
            custom_print(value, indent + 1)
        else:
            print "%s" % (str(value)),

class ConfigDict (OrderedDict):

    def _set_filename(self, filename):
        self['filename'] = filename
        self['location'] = path_expand(self["filename"])

    def __init__(self, *args, **kwargs):
        OrderedDict.__init__(self, *args, **kwargs)

        if 'filename' in kwargs:
            self._set_filename(kwargs['filename'])
        else:
            log.error("filename not specified")
            #sys.exit()

        self.load(self['location'])

        for attribute in ['prefix']:
            if attribute in kwargs:
                self[attribute] = kwargs[attribute]
            else:
                self[attribute] = None

        self._update_meta()

    def _update_meta(self):
        for v in ["filename", "location", "prefix"]:
            self["meta"][v] = self[v]
            del self[v]

    def read(self, filename):
        """does the same as load"""
        self.load(filename)

    def load(self, filename):
        self._set_filename(filename)
        #d = OrderedDict(read_yaml_config(self['location'], check=True))
        d = read_yaml_config(self['location'], check=True)
        self.update(d)

    def write(self, filename=None, format="dict"):
        """write the dict"""

        if filename is not None:
            location = path_expand(filename)
        else:
            location = self['meta']['location']

        # with open('data.yml', 'w') as outfile:
            #    outfile.write( yaml.dump(data, default_flow_style=True) )

        f = os.open(location, os.O_CREAT | os.O_TRUNC |
                    os.O_WRONLY, stat.S_IRUSR | stat.S_IWUSR)
        if format == "json":
            os.write(f, self.json())
        elif format in ['yml', 'yaml']:
            #d = dict(self)
            #os.write(f, yaml.dump(d, default_flow_style=False))
            os.write(f, ordered_dump(OrderedDict(self), Dumper=yaml.SafeDumper,
                                     default_flow_style=False))
        elif format == "print":
            os.write(f, custom_print(self, 4))
        else:
            os.write(f, self.dump())
        os.close(f)

    def error_keys_not_found(self, keys):
        log.error("Filename: {0}".format(self['meta']['location']))
        log.error("Key '{0}' does not exist".format('.'.join(keys)))
        indent = ""
        last_index = len(keys) - 1
        for i, k in enumerate(keys):
            if i == last_index:
                log.error(indent + k + ": <- this value is missing")
            else:
                log.error(indent + k + ":")
            indent = indent + "    "

    def __str__(self):
        return self.json()

    def json(self):
        return json.dumps(self, indent=4)

    def dump(self):
        orderedPrinter = OrderedJsonEncoder()
        return orderedPrinter.encode(self)

    def pprint(self):
        print custom_print(self, 4)

    """
    def __getitem__(self, *mykeys):
        try:
            item = self.get(mykeys[0])
        except:
            self._notify_of_error(mykeys)
            sys.exit()
        return item
    """

    def get(self, *keys):
        """
        returns the dict of the information as read from the yaml
        file. To access the file safely, you can use the keys in the
        order of the access.  Example: get("provisiner","policy") will
        return the value of config["provisiner"]["policy"] from the
        yaml file if it does not exists an error will be printing that
        the value does not exists.  Alternatively you can
        use the . notation e.g. get("provisiner.policy")
        """
        if keys is None:
            return self

        if "." in keys[0]:
            keys = keys[0].split('.')
        element = self
        for v in keys:
            try:
                element = element[v]
            except KeyError:
                self.error_keys_not_found(keys)
                #sys.exit()
        return element

    def set(self, value, *keys):
        """
        Sets the dict of the information as read from the yaml file. To access
        the file safely, you can use the keys in the order of the access.
        Example: set("{'project':{'fg82':[i0-i10]}}", "provisiner","policy")
        will set the value of config["provisiner"]["policy"] in the yaml
        file if it does not exists an error will be printing that the value does
        not exists.  Alternatively you can use the . notation e.g.
        set("{'project':{'fg82':[i0-i10]}}", "provisiner.policy")
        """
        element = self

        if keys is None:
            return self

        if '.' in keys[0]:
            keys = keys[0].split(".")

        nested_str = ''.join([ "['{0}']".format(x) for x in keys ])
        exec("self" + nested_str + "=" + "'" + str(value) + "'")
        return element

    def _update(self, keys, value):
        """Updates the selected key with the value

        Args:
            keys (str): key names e.g. cloudmesh.server.loglevel
            value (str): value to set
        """
        return self.set(value, keys)

    def attribute(self, keys):
        if self['meta']['prefix'] is None:
            k = keys
        else:
            k = self['meta']['prefix'] + "." + keys
        return self.get(k)

if __name__ == "__main__":
    config = ConfigDict({"a": "1", "b": {"c": 3}},
                        prefix="cloudmesh.server",
                        filename="./etc/cloudmesh_server.yaml")

    print "PPRINT"
    print 70 * "="
    pprint(config)

    print "PRINT"
    print 70 * "="
    print config.pprint()
    print config.json()

    print 70 * "="
    print "A =", config["a"]
    config.write(config_file("/d.yaml"), format="dict")
    config.write(config_file("/j.yaml"), format="json")
    config.write(config_file("/y.yaml"), format="yaml")

    # this does not work
    # config.write(config_file("/print.yaml"), format="print")

    print "mongo.path GET =", config.get("cloudmesh.server.mongo.path")
    print "mongo.path ATTIRBUTE =", config.attribute("mongo.path")

    print "get A =", config.get("a")

    print "wrong mongo.path ATTRIBUTE =", config.attribute("mongo.path.wrong")
    print "wrong mongo.path GET =", config.get("cloudmesh.server.mongo.path.wrong")

    # print config["dummy"]
    # config["x"] = "2"
    # print config["x"]
    # print config.x
<commit_msg>Make a copy of yaml for restore<commit_after>from cloudmesh_install import config_file
from cloudmesh_install.util import path_expand
from cloudmesh_common.logger import LOGGER
from cloudmesh_common.util import backup_name
from cloudmesh.util.config import read_yaml_config
from cloudmesh.util.config import ordered_dump
from collections import OrderedDict
import simplejson
from pprint import pprint
import yaml
import json
import os
import stat
import sys

log = LOGGER(__file__)
package_dir = os.path.dirname(os.path.abspath(__file__))

class OrderedJsonEncoder(simplejson.JSONEncoder):

    indent = 4

    def encode(self, o, depth=0):
        if isinstance(o, OrderedDict):
            return "{" + (",\n ").join([self.encode(k) + ":" + \
                                        self.encode(v, depth + 1) \
                                        for (k, v) in o.iteritems()]) + "}\n"
        else:
            return simplejson.JSONEncoder.encode(self, o)

def custom_print(data_structure, indent):
    for key, value in data_structure.items():
        print "\n%s%s:" % ('    ' * indent, str(key)),
        if isinstance(value, OrderedDict):
            custom_print(value, indent + 1)
        elif isinstance(value, dict):
            custom_print(value, indent + 1)
        else:
            print "%s" % (str(value)),

class ConfigDict (OrderedDict):

    def _set_filename(self, filename):
        self['filename'] = filename
        self['location'] = path_expand(self["filename"])

    def __init__(self, *args, **kwargs):
        OrderedDict.__init__(self, *args, **kwargs)

        if 'filename' in kwargs:
            self._set_filename(kwargs['filename'])
        else:
            log.error("filename not specified")
            #sys.exit()

        self.load(self['location'])

        for attribute in ['prefix']:
            if attribute in kwargs:
                self[attribute] = kwargs[attribute]
            else:
                self[attribute] = None

        self._update_meta()

    def _update_meta(self):
        for v in ["filename", "location", "prefix"]:
            self["meta"][v] = self[v]
            del self[v]

    def read(self, filename):
        """does the same as load"""
        self.load(filename)

    def load(self, filename):
        self._set_filename(filename)
        #d = OrderedDict(read_yaml_config(self['location'], check=True))
        d = read_yaml_config(self['location'], check=True)
        self.update(d)

    def make_a_copy(self, location=None):
        import shutil
        dest = backup_name(location)
        shutil.copyfile(location, dest)

    def write(self, filename=None, format="dict"):
        """write the dict"""

        if filename is not None:
            location = path_expand(filename)
        else:
            location = self['meta']['location']

        # with open('data.yml', 'w') as outfile:
            #    outfile.write( yaml.dump(data, default_flow_style=True) )

        # Make a backup
        self.make_a_copy(location)

        f = os.open(location, os.O_CREAT | os.O_TRUNC |
                    os.O_WRONLY, stat.S_IRUSR | stat.S_IWUSR)
        if format == "json":
            os.write(f, self.json())
        elif format in ['yml', 'yaml']:
            #d = dict(self)
            #os.write(f, yaml.dump(d, default_flow_style=False))
            os.write(f, ordered_dump(OrderedDict(self), Dumper=yaml.SafeDumper,
                                     default_flow_style=False))
        elif format == "print":
            os.write(f, custom_print(self, 4))
        else:
            os.write(f, self.dump())
        os.close(f)

    def error_keys_not_found(self, keys):
        log.error("Filename: {0}".format(self['meta']['location']))
        log.error("Key '{0}' does not exist".format('.'.join(keys)))
        indent = ""
        last_index = len(keys) - 1
        for i, k in enumerate(keys):
            if i == last_index:
                log.error(indent + k + ": <- this value is missing")
            else:
                log.error(indent + k + ":")
            indent = indent + "    "

    def __str__(self):
        return self.json()

    def json(self):
        return json.dumps(self, indent=4)

    def dump(self):
        orderedPrinter = OrderedJsonEncoder()
        return orderedPrinter.encode(self)

    def pprint(self):
        print custom_print(self, 4)

    """
    def __getitem__(self, *mykeys):
        try:
            item = self.get(mykeys[0])
        except:
            self._notify_of_error(mykeys)
            sys.exit()
        return item
    """

    def get(self, *keys):
        """
        returns the dict of the information as read from the yaml
        file. To access the file safely, you can use the keys in the
        order of the access.  Example: get("provisiner","policy") will
        return the value of config["provisiner"]["policy"] from the
        yaml file if it does not exists an error will be printing that
        the value does not exists.  Alternatively you can
        use the . notation e.g. get("provisiner.policy")
        """
        if keys is None:
            return self

        if "." in keys[0]:
            keys = keys[0].split('.')
        element = self
        for v in keys:
            try:
                element = element[v]
            except KeyError:
                self.error_keys_not_found(keys)
                #sys.exit()
        return element

    def set(self, value, *keys):
        """
        Sets the dict of the information as read from the yaml file. To access
        the file safely, you can use the keys in the order of the access.
        Example: set("{'project':{'fg82':[i0-i10]}}", "provisiner","policy")
        will set the value of config["provisiner"]["policy"] in the yaml
        file if it does not exists an error will be printing that the value does
        not exists.  Alternatively you can use the . notation e.g.
        set("{'project':{'fg82':[i0-i10]}}", "provisiner.policy")
        """
        element = self

        if keys is None:
            return self

        if '.' in keys[0]:
            keys = keys[0].split(".")

        nested_str = ''.join([ "['{0}']".format(x) for x in keys ])
        exec("self" + nested_str + "=" + "'" + str(value) + "'")
        return element

    def _update(self, keys, value):
        """Updates the selected key with the value

        Args:
            keys (str): key names e.g. cloudmesh.server.loglevel
            value (str): value to set
        """
        return self.set(value, keys)

    def attribute(self, keys):
        if self['meta']['prefix'] is None:
            k = keys
        else:
            k = self['meta']['prefix'] + "." + keys
        return self.get(k)

if __name__ == "__main__":
    config = ConfigDict({"a": "1", "b": {"c": 3}},
                        prefix="cloudmesh.server",
                        filename="./etc/cloudmesh_server.yaml")

    print "PPRINT"
    print 70 * "="
    pprint(config)

    print "PRINT"
    print 70 * "="
    print config.pprint()
    print config.json()

    print 70 * "="
    print "A =", config["a"]
    config.write(config_file("/d.yaml"), format="dict")
    config.write(config_file("/j.yaml"), format="json")
    config.write(config_file("/y.yaml"), format="yaml")

    # this does not work
    # config.write(config_file("/print.yaml"), format="print")

    print "mongo.path GET =", config.get("cloudmesh.server.mongo.path")
    print "mongo.path ATTIRBUTE =", config.attribute("mongo.path")

    print "get A =", config.get("a")

    print "wrong mongo.path ATTRIBUTE =", config.attribute("mongo.path.wrong")
    print "wrong mongo.path GET =", config.get("cloudmesh.server.mongo.path.wrong")

    # print config["dummy"]
    # config["x"] = "2"
    # print config["x"]
    # print config.x
<|endoftext|>


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

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from builtins import object

import numpy as np
import pkg_resources
import fsps
import os
import linecache
#grid: 2 files: line, cont
#columns: wavelengths
#rows: models
#header
#wavelength grid
#age, Z, logU
#flux1
#flux2
#flux3

###
# reads from ZAU***.out_lines, ZAU***.out_cont
# produces ZAU_**.lines, ZAU_**.cont
###
class writeFormattedOutput(object):
    #sp = fsps.StellarPopulation()
    #fsps_lam = sp.wavelengths
    def __init__(self, dir_, mod_prefix, mod_suffix,
                 use_extended_lines=False, more_info=False,**kwargs):
        '''
        writeFormattedOutput(mod_dir, 'ZAU', 'BPASS')
        '''
        self.dir_, self.mod_prefix = dir_, mod_prefix
        self.file_pr = dir_ + mod_prefix
        if mod_suffix is None:
            self.out_pr = dir_ + mod_prefix
        else:
            self.out_pr = dir_ + mod_prefix + mod_suffix
        # each model's final info will be in prefix00.lines, prefix00.cont
        self.line_out = self.out_pr + ".lines"
        self.cont_out = self.out_pr + ".cont"
        # load each model's parameters from prefix.pars
        self.loadModInfo()
        # print ordered emission line wavelengths + fluxes
        self.doLineOut(use_extended_lines=use_extended_lines,
                       more_info=more_info)
        # interp and print neb cont onto FSPS wavelenth arr
        self.doContOut()
        return
    def loadModInfo(self, **kwargs):
        '''
        reads model parameters from "ZAU.pars"
        '''
        name_keys = ["mod_num", "logZ", "Age", "logU", "logR", "logQ", "nH", "efrac"]
        data = np.genfromtxt(self.file_pr+".pars", unpack=True)
        ddata = {}
        for i,key in enumerate(name_keys):
            ddata[key] = data[i]
            self.__setattr__(key, data[i])
        self.__setattr__("modpars", ddata)
        self.NZ = len(np.unique(self.logZ))
        self.NA = len(np.unique(self.Age))
        self.NU = len(np.unique(self.logU))
        return
    def doLineOut(self, more_info=False, **kwargs):
        '''
        prints line fluxes to prefix00.lines file
        '''
        f = open(self.line_out, "w")
        self.printLineLam(f, **kwargs)
        for n in self.mod_num:
            self.printLineFlu(f, n.astype(int), more_info=more_info)
        f.close()
        print("lines: {0:.0f} models to file {1}".format(self.mod_num[-1], self.line_out))
        return
    def printLineLam(self, f, use_extended_lines=False):
        '''
        prints the wavelength array for the emission lines as the second
        line in the output file. converts all wavelengths to vacuum.
        '''
        #read in file containing wavelength info
        if use_extended_lines:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLinesEXT.dat")
        else:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLines.dat")
        data_vac = np.genfromtxt(linefile)
        #data_vac = air_to_vac(data) # new file is already in vac
        nlines = len(data_vac)
        nmods = np.max(self.mod_num)
        #print header to file
        head_str = "#{0} cols {1:.0f} rows {2} logZ {3} Age {4} logU".format(nlines, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(dat) for dat in data_vac])
        f.write(p_str+"\n")
        return
    def printLineFlu(self, f, n, more_info=False):
        #write model parameters
        if more_info:
            tstr = linecache.getline(self.file_pr+'.pars', n)
        else:
            tstr = "{0:2.4e} {1:2.4e} {2:2.4e}\n".format(self.logZ[n-1], self.Age[n-1], self.logU[n-1])
        f.write(tstr)
        #read in and print emission line intensities (Lsun/Q)
        nst = "{0}".format(n)
        filename = self.file_pr+nst+".out_lines"
        lam, flu = np.genfromtxt(filename, unpack=True)
        I_str = ["{0:1.4e}".format(s) for s in flu]
        tstr = " ".join(I_str)
        f.write(tstr+"\n")
        return
    def doContOut(self, **kwargs):
        f = open(self.cont_out, "w")
        self.printContLam(f)
        for num in self.modpars['mod_num']:
            iind = num.astype(int)
            pars = dict(logZ=self.logZ[iind-1],
                        Age=self.Age[iind-1],
                        nH=self.nH[iind-1],
                        logQ=self.logQ[iind-1],
                        logU=self.logU[iind-1],
                        logR=self.logR[iind-1],
                        mod_num=iind,
                        efrac=self.efrac[iind-1])
            self.printContFlu(f, pars)
        f.close()
        print("cont: {0:.0f} models to file {1}".format(self.modpars['mod_num'][-1], self.cont_out))
        return
    def printContFlu(self, f, pars):
        #write model parameters
        tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(pars["logZ"],
                                                   pars["Age"],
                                                   pars["logU"])
        f.write(tstr+"\n")
        #read in and print emission line intensities
        nst = "{0}".format(pars["mod_num"])
        filename = self.file_pr+nst+".out_cont"
        # read lambda, diffuse cont
        mdata = np.genfromtxt(filename)
        lam = mdata[:,0]
        flu = mdata[:,1]
        y_str = " ".join(["{0:1.4}".format(y) for y in flu])
        f.write(y_str+"\n")
        return
    def printContLam(self, f):
        '''
        prints header in first line of file, and the FSPS wavelength array
        (for the nebular continuum) as the second line in the output file.
        # nlam cols nmod rows logZ Age logU
        # fsps_lam_1 fsps_lam_2 .... fsps_lam_n
        '''
        #grab fsps wavelength info
        lamfile = pkg_resources.resource_filename(__name__, "data/FSPSlam.dat")
        fsps_lam = np.genfromtxt(lamfile)
        self.__setattr__("fsps_lam", fsps_lam)
        nlam = len(fsps_lam)
        nmods = np.max(self.mod_num).astype(int)
        #print header to file
        head_str = "#{0} cols {1} rows {2} logZ {3} Age {4} logU".format(nlam, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(lam) for lam in fsps_lam])
        f.write(p_str+"\n")
        return

#------------------------------------------------

class writeAltFormattedOutput(object):
    #sp = fsps.StellarPopulation()
    #fsps_lam = sp.wavelengths
    def __init__(self, dir_, mod_prefix, mod_suffix,
                 use_extended_lines=False, more_info=False,**kwargs):
        '''
        writeFormattedOutput(mod_dir, 'ZAU', 'BPASS')
        '''
        self.dir_, self.mod_prefix = dir_, mod_prefix
        self.file_pr = dir_ + mod_prefix
        if mod_suffix is None:
            self.out_pr = dir_ + mod_prefix
        else:
            self.out_pr = dir_ + mod_prefix + mod_suffix
        # each model's final info will be in prefix00.lines, prefix00.cont
        self.line_out = self.out_pr + ".lines"
        self.cont_out = self.out_pr + ".cont"
        self.loadModInfo() # load each model's parameters from prefix.pars
        # print ordered emission line wavelengths + fluxes
        self.doLineOut(use_extended_lines=use_extended_lines,
                       more_info=more_info)
        self.doContOut() # interp and print neb cont onto FSPS wavelenth arr
        return
    def loadModInfo(self, **kwargs):
        '''
        reads model parameters from "ZAU.pars"
        '''
        name_keys = ["mod_num", "logZ", "Age", "logU", "logR", "logQ", "nH", "efrac", "zmet"]
        data = np.genfromtxt(self.file_pr+".pars", unpack=True)
        ddata = {}
        for i,key in enumerate(name_keys):
            ddata[key] = data[i]
            self.__setattr__(key, data[i])
        self.__setattr__("modpars", ddata)
        self.NZ = len(np.unique(self.zmet))
        self.NA = len(np.unique(self.Age))
        self.NU = len(np.unique(self.logU))
        return
    def doLineOut(self, more_info=False, **kwargs):
        '''
        prints line fluxes to prefix00.lines file
        '''
        f = open(self.line_out, "w")
        self.printLineLam(f, **kwargs)
        for n in self.mod_num:
            self.printLineFlu(f, n.astype(int), more_info=more_info)
        f.close()
        print("lines: {0:.0f} models to file {1}".format(self.mod_num[-1], self.line_out))
        return
    def printLineLam(self, f, use_extended_lines=False):
        '''
        prints the wavelength array for the emission lines as the second
        line in the output file. converts all wavelengths to vacuum.
        '''
        #read in file containing wavelength info
        if use_extended_lines:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLinesEXT.dat")
        else:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLines.dat")
        data_vac = np.genfromtxt(linefile)
        #data_vac = air_to_vac(data) # new file is already in vac
        nlines = len(data_vac)
        nmods = np.max(self.mod_num)
        #print header to file
        head_str = "#{0} cols {1:.0f} rows {2} logZ {3} Age {4} logU".format(nlines, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(dat) for dat in data_vac])
        f.write(p_str+"\n")
        return
    def printLineFlu(self, f, n, more_info=False):
        #write model parameters
        if more_info:
            tstr = linecache.getline(self.file_pr+'.pars', n)
            f.write(tstr)
        else:
            tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(self.zmet[n-1], self.Age[n-1], self.logU[n-1])
            f.write(tstr+"\n")
        #read in and print emission line intensities (Lsun/Q)
        nst = "{0}".format(n)
        filename = self.file_pr+nst+".out_lines"
        lam, flu = np.genfromtxt(filename, unpack=True)
        I_str = ["{0:1.4e}".format(s) for s in flu]
        tstr = " ".join(I_str)
        f.write(tstr+"\n")
        return
    def doContOut(self, **kwargs):
        f = open(self.cont_out, "w")
        self.printContLam(f)
        for num in self.modpars['mod_num']:
            iind = num.astype(int)
            pars = dict(logZ=self.logZ[iind-1],
                        Age=self.Age[iind-1],
                        nH=self.nH[iind-1],
                        logQ=self.logQ[iind-1],
                        logU=self.logU[iind-1],
                        logR=self.logR[iind-1],
                        mod_num=iind,
                        zmet=self.zmet[iind-1])
            self.printContFlu(f, pars)
        f.close()
        print("cont: {0:.0f} models to file {1}".format(self.modpars['mod_num'][-1], self.cont_out))
        return
    def printContFlu(self, f, pars):
        #write model parameters
        tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(pars["zmet"],
                                                   pars["Age"],
                                                   pars["logU"])
        f.write(tstr+"\n")
        #read in and print emission line intensities
        nst = "{0}".format(pars["mod_num"])
        filename = self.file_pr+nst+".out_cont"
        # read lambda, diffuse cont
        mdata = np.genfromtxt(filename)
        lam = mdata[:,0]
        flu = mdata[:,1]
        y_str = " ".join(["{0:1.4}".format(y) for y in flu])
        f.write(y_str+"\n")
        return
    def printContLam(self, f):
        '''
        prints header in first line of file, and the FSPS wavelength array
        (for the nebular continuum) as the second line in the output file.
        # nlam cols nmod rows logZ Age logU
        # fsps_lam_1 fsps_lam_2 .... fsps_lam_n
        '''
        #grab fsps wavelength info
        lamfile = pkg_resources.resource_filename(__name__, "data/FSPSlam.dat")
        fsps_lam = np.genfromtxt(lamfile)
        self.__setattr__("fsps_lam", fsps_lam)
        nlam = len(fsps_lam)
        nmods = np.max(self.mod_num).astype(int)
        #print header to file
        head_str = "#{0} cols {1} rows {2} logZ {3} Age {4} logU".format(nlam, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(lam) for lam in fsps_lam])
        f.write(p_str+"\n")
        return
<commit_msg>updating for rotating zmet<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from builtins import object

import numpy as np
import pkg_resources
import fsps
import os
import linecache
#grid: 2 files: line, cont
#columns: wavelengths
#rows: models
#header
#wavelength grid
#age, Z, logU
#flux1
#flux2
#flux3

###
# reads from ZAU***.out_lines, ZAU***.out_cont
# produces ZAU_**.lines, ZAU_**.cont
###
class writeFormattedOutput(object):
    #sp = fsps.StellarPopulation()
    #fsps_lam = sp.wavelengths
    def __init__(self, dir_, mod_prefix, mod_suffix,
                 use_extended_lines=False, more_info=False,**kwargs):
        '''
        writeFormattedOutput(mod_dir, 'ZAU', 'BPASS')
        '''
        self.dir_, self.mod_prefix = dir_, mod_prefix
        self.file_pr = dir_ + mod_prefix
        if mod_suffix is None:
            self.out_pr = dir_ + mod_prefix
        else:
            self.out_pr = dir_ + mod_prefix + mod_suffix
        # each model's final info will be in prefix00.lines, prefix00.cont
        self.line_out = self.out_pr + ".lines"
        self.cont_out = self.out_pr + ".cont"
        # load each model's parameters from prefix.pars
        self.loadModInfo()
        # print ordered emission line wavelengths + fluxes
        self.doLineOut(use_extended_lines=use_extended_lines,
                       more_info=more_info)
        # interp and print neb cont onto FSPS wavelenth arr
        self.doContOut()
        return
    def loadModInfo(self, **kwargs):
        '''
        reads model parameters from "ZAU.pars"
        '''
        name_keys = ["mod_num", "logZ", "Age", "logU", "logR", "logQ", "nH", "efrac"]
        data = np.genfromtxt(self.file_pr+".pars", unpack=True)
        ddata = {}
        for i,key in enumerate(name_keys):
            ddata[key] = data[i]
            self.__setattr__(key, data[i])
        self.__setattr__("modpars", ddata)
        self.NZ = len(np.unique(self.logZ))
        self.NA = len(np.unique(self.Age))
        self.NU = len(np.unique(self.logU))
        return
    def doLineOut(self, more_info=False, **kwargs):
        '''
        prints line fluxes to prefix00.lines file
        '''
        f = open(self.line_out, "w")
        self.printLineLam(f, **kwargs)
        for n in self.mod_num:
            self.printLineFlu(f, n.astype(int), more_info=more_info)
        f.close()
        print("lines: {0:.0f} models to file {1}".format(self.mod_num[-1], self.line_out))
        return
    def printLineLam(self, f, use_extended_lines=False):
        '''
        prints the wavelength array for the emission lines as the second
        line in the output file. converts all wavelengths to vacuum.
        '''
        #read in file containing wavelength info
        if use_extended_lines:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLinesEXT.dat")
        else:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLines.dat")
        data_vac = np.genfromtxt(linefile)
        #data_vac = air_to_vac(data) # new file is already in vac
        nlines = len(data_vac)
        nmods = np.max(self.mod_num)
        #print header to file
        head_str = "#{0} cols {1:.0f} rows {2} logZ {3} Age {4} logU".format(nlines, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(dat) for dat in data_vac])
        f.write(p_str+"\n")
        return
    def printLineFlu(self, f, n, more_info=False):
        #write model parameters
        if more_info:
            tstr = linecache.getline(self.file_pr+'.pars', n)
        else:
            tstr = "{0:2.4e} {1:2.4e} {2:2.4e}\n".format(self.logZ[n-1], self.Age[n-1], self.logU[n-1])
        f.write(tstr)
        #read in and print emission line intensities (Lsun/Q)
        nst = "{0}".format(n)
        filename = self.file_pr+nst+".out_lines"
        lam, flu = np.genfromtxt(filename, unpack=True)
        I_str = ["{0:1.4e}".format(s) for s in flu]
        tstr = " ".join(I_str)
        f.write(tstr+"\n")
        return
    def doContOut(self, **kwargs):
        f = open(self.cont_out, "w")
        self.printContLam(f)
        for num in self.modpars['mod_num']:
            iind = num.astype(int)
            pars = dict(logZ=self.logZ[iind-1],
                        Age=self.Age[iind-1],
                        nH=self.nH[iind-1],
                        logQ=self.logQ[iind-1],
                        logU=self.logU[iind-1],
                        logR=self.logR[iind-1],
                        mod_num=iind,
                        efrac=self.efrac[iind-1])
            self.printContFlu(f, pars)
        f.close()
        print("cont: {0:.0f} models to file {1}".format(self.modpars['mod_num'][-1], self.cont_out))
        return
    def printContFlu(self, f, pars):
        #write model parameters
        tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(pars["logZ"],
                                                   pars["Age"],
                                                   pars["logU"])
        f.write(tstr+"\n")
        #read in and print emission line intensities
        nst = "{0}".format(pars["mod_num"])
        filename = self.file_pr+nst+".out_cont"
        # read lambda, diffuse cont
        mdata = np.genfromtxt(filename)
        lam = mdata[:,0]
        flu = mdata[:,1]
        y_str = " ".join(["{0:1.4}".format(y) for y in flu])
        f.write(y_str+"\n")
        return
    def printContLam(self, f):
        '''
        prints header in first line of file, and the FSPS wavelength array
        (for the nebular continuum) as the second line in the output file.
        # nlam cols nmod rows logZ Age logU
        # fsps_lam_1 fsps_lam_2 .... fsps_lam_n
        '''
        #grab fsps wavelength info
        lamfile = pkg_resources.resource_filename(__name__, "data/FSPSlam.dat")
        fsps_lam = np.genfromtxt(lamfile)
        self.__setattr__("fsps_lam", fsps_lam)
        nlam = len(fsps_lam)
        nmods = np.max(self.mod_num).astype(int)
        #print header to file
        head_str = "#{0} cols {1} rows {2} logZ {3} Age {4} logU".format(nlam, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(lam) for lam in fsps_lam])
        f.write(p_str+"\n")
        return

#------------------------------------------------
def zmet_to_nuZ(zmet):
    '''
    zmet = 1,2,3,4,5,6,7,8
        ...corresponding to...
    m1.00, m1.02, m1.04, m1.06, p0.00, p0.02, p0.04, p0.06
    '''
    zmets = np.array([1,2,3,4,5,6,7,8])
    zs = np.array([-1.00, -1.02, -1.04, -1.06, 0.00, 0.02, 0.04, 0.06])
    if type(zmet) is float:
        ii, = np.where(zmets == zmet)[0]
        out_val = zs[ii]
    else:
        out_val = np.zeros_like(zmet, dtype='float')
        for i in range(len(zmet)):
            ii, = np.where(zmets == zmet[i])[0]
            out_val[i] = zs[ii]
    return out_val

class writeAltFormattedOutput(object):
    #sp = fsps.StellarPopulation()
    #fsps_lam = sp.wavelengths
    def __init__(self, dir_, mod_prefix, mod_suffix,
                 use_extended_lines=False, more_info=False,**kwargs):
        '''
        writeFormattedOutput(mod_dir, 'ZAU', 'BPASS')
        '''
        self.dir_, self.mod_prefix = dir_, mod_prefix
        self.file_pr = dir_ + mod_prefix
        if mod_suffix is None:
            self.out_pr = dir_ + mod_prefix
        else:
            self.out_pr = dir_ + mod_prefix + mod_suffix
        # each model's final info will be in prefix00.lines, prefix00.cont
        self.line_out = self.out_pr + ".lines"
        self.cont_out = self.out_pr + ".cont"
        self.loadModInfo() # load each model's parameters from prefix.pars
        # print ordered emission line wavelengths + fluxes
        self.doLineOut(use_extended_lines=use_extended_lines,
                       more_info=more_info)
        self.doContOut() # interp and print neb cont onto FSPS wavelenth arr
        return
    def loadModInfo(self, **kwargs):
        '''
        reads model parameters from "ZAU.pars"
        '''
        name_keys = ["mod_num", "logZ", "Age", "logU", "logR", "logQ", "nH", "efrac", "zmet"]
        data = np.genfromtxt(self.file_pr+".pars", unpack=True)
        ddata = {}
        for i,key in enumerate(name_keys):
            if key == 'zmet':
                temp_arr = zmet_to_nuZ(data[i])
                ddata[key] = temp_arr
                self.__setattr__(key, temp_arr)
            else:
                ddata[key] = data[i]
                self.__setattr__(key, data[i])
        self.__setattr__("modpars", ddata)
        self.NZ = len(np.unique(self.zmet))
        self.NA = len(np.unique(self.Age))
        self.NU = len(np.unique(self.logU))
        return
    def doLineOut(self, more_info=False, **kwargs):
        '''
        prints line fluxes to prefix00.lines file
        '''
        f = open(self.line_out, "w")
        self.printLineLam(f, **kwargs)
        for n in self.mod_num:
            self.printLineFlu(f, n.astype(int), more_info=more_info)
        f.close()
        print("lines: {0:.0f} models to file {1}".format(self.mod_num[-1], self.line_out))
        return
    def printLineLam(self, f, use_extended_lines=False):
        '''
        prints the wavelength array for the emission lines as the second
        line in the output file. converts all wavelengths to vacuum.
        '''
        #read in file containing wavelength info
        if use_extended_lines:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLinesEXT.dat")
        else:
            linefile = pkg_resources.resource_filename(__name__, "data/orderedLines.dat")
        data_vac = np.genfromtxt(linefile)
        #data_vac = air_to_vac(data) # new file is already in vac
        nlines = len(data_vac)
        nmods = np.max(self.mod_num)
        #print header to file
        head_str = "#{0} cols {1:.0f} rows {2} logZ {3} Age {4} logU".format(nlines, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(dat) for dat in data_vac])
        f.write(p_str+"\n")
        return
    def printLineFlu(self, f, n, more_info=False):
        #write model parameters
        if more_info:
            tstr = linecache.getline(self.file_pr+'.pars', n)
            f.write(tstr)
        else:
            tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(self.zmet[n-1], self.Age[n-1], self.logU[n-1])
            f.write(tstr+"\n")
        #read in and print emission line intensities (Lsun/Q)
        nst = "{0}".format(n)
        filename = self.file_pr+nst+".out_lines"
        lam, flu = np.genfromtxt(filename, unpack=True)
        I_str = ["{0:1.4e}".format(s) for s in flu]
        tstr = " ".join(I_str)
        f.write(tstr+"\n")
        return
    def doContOut(self, **kwargs):
        f = open(self.cont_out, "w")
        self.printContLam(f)
        for num in self.modpars['mod_num']:
            iind = num.astype(int)
            pars = dict(logZ=self.logZ[iind-1],
                        Age=self.Age[iind-1],
                        nH=self.nH[iind-1],
                        logQ=self.logQ[iind-1],
                        logU=self.logU[iind-1],
                        logR=self.logR[iind-1],
                        mod_num=iind,
                        zmet=self.zmet[iind-1])
            self.printContFlu(f, pars)
        f.close()
        print("cont: {0:.0f} models to file {1}".format(self.modpars['mod_num'][-1], self.cont_out))
        return
    def printContFlu(self, f, pars):
        #write model parameters
        tstr = "{0:2.4e} {1:2.4e} {2:2.4e}".format(pars["zmet"],
                                                   pars["Age"],
                                                   pars["logU"])
        f.write(tstr+"\n")
        #read in and print emission line intensities
        nst = "{0}".format(pars["mod_num"])
        filename = self.file_pr+nst+".out_cont"
        # read lambda, diffuse cont
        mdata = np.genfromtxt(filename)
        lam = mdata[:,0]
        flu = mdata[:,1]
        y_str = " ".join(["{0:1.4}".format(y) for y in flu])
        f.write(y_str+"\n")
        return
    def printContLam(self, f):
        '''
        prints header in first line of file, and the FSPS wavelength array
        (for the nebular continuum) as the second line in the output file.
        # nlam cols nmod rows logZ Age logU
        # fsps_lam_1 fsps_lam_2 .... fsps_lam_n
        '''
        #grab fsps wavelength info
        lamfile = pkg_resources.resource_filename(__name__, "data/FSPSlam.dat")
        fsps_lam = np.genfromtxt(lamfile)
        self.__setattr__("fsps_lam", fsps_lam)
        nlam = len(fsps_lam)
        nmods = np.max(self.mod_num).astype(int)
        #print header to file
        head_str = "#{0} cols {1} rows {2} logZ {3} Age {4} logU".format(nlam, nmods, self.NZ, self.NA, self.NU)
        f.write(head_str+"\n")
        #print lambda array
        p_str = " ".join(["{0:1.6e}".format(lam) for lam in fsps_lam])
        f.write(p_str+"\n")
        return
<|endoftext|>