<commit_before>import os
import json
import ibmiotf.application
import iso8601
import base64
from bottle import Bottle, template

import requests
from requests_futures.sessions import FuturesSession

import urllib
import argparse
import logging
from logging.handlers import RotatingFileHandler

class Server():

	def __init__(self, args):
		# Setup logging - Generate a default rotating file log handler and stream handler
		logFileName = 'connector-messagehub.log'
		fhFormatter = logging.Formatter('%(asctime)-25s %(name)-30s ' + ' %(levelname)-7s %(message)s')
		rfh = RotatingFileHandler(logFileName, mode='a', maxBytes=26214400 , backupCount=2, encoding=None, delay=True)
		rfh.setFormatter(fhFormatter)

		self.logger = logging.getLogger("server")
		self.logger.addHandler(rfh)
		self.logger.setLevel(logging.DEBUG)

		self.port = int(os.getenv('VCAP_APP_PORT', '8000'))
		self.host = str(os.getenv('VCAP_APP_HOST', 'localhost'))

		if args.bluemix == True:
			# Bluemix VCAP lookups
			application = json.loads(os.getenv('VCAP_APPLICATION'))
			service = json.loads(os.getenv('VCAP_SERVICES'))

			# IoTF
			self.options = ibmiotf.application.ParseConfigFromBluemixVCAP()

			# Cloudant
			self.messagehubApikey = service['messagehub'][0]['credentials']['api_key']
			self.messagehubUrl = service['messagehub'][0]['credentials']['kafka_rest_url']
		else:
			self.options = ibmiotf.application.ParseConfigFile(args.config)
			self.messagehubApikey = args.messagehubApikey
			self.messagehubUrl = args.messagehubUrl

		# Bottle
		self._app = Bottle()
		self._route()

		# Init MessageHub Topic & requests session
		self._createTopic()
		self.session = FuturesSession()

		# Init IOTF client
		self.client = ibmiotf.application.Client(self.options, logHandlers=[rfh])

	def _route(self):
		self._app.route('/', method="GET", callback=self._status)

	def myEventCallback(self, evt):
		#self.logger.info("%-33s%-30s%s" % (evt.timestamp.isoformat(), evt.device, evt.event + ": " + json.dumps(evt.data)))
		#self.logger.info(evt.data)

		headers = {
			"Content-Type": "application/json",
			"X-Auth-Token": self.messagehubApikey
		}

		message = {
			'typeId': evt.deviceType,
			'deviceId': evt.deviceId,
			'eventId': evt.event,
			'timestamp': evt.timestamp.isoformat(),
			'data': evt.data,
			'format': evt.format,
			'payload': base64.encodestring(evt.payload).decode('ascii')
		}

		payload = {
			'records': [{"value": json.dumps(message)}]
		}

		future = self.session.post(self.messagehubUrl + "/topics/events", data=json.dumps(payload), headers=headers)
		future.add_done_callback(self._eventRecordedCallback)

	def _eventRecordedCallback(self, future):
		response = future.result()
		if response.status_code not in [200]:
			self.logger.info("%s - %s: %s" % (response.url, response.status_code, response.text))

	def start(self):
		self.client.connect()
		self.client.deviceEventCallback = self.myEventCallback
		self.client.subscribeToDeviceEvents()
		self.logger.info("Serving at %s:%s" % (self.host, self.port))
		self._app.run(host=self.host, port=self.port)

	def stop(self):
		self.client.disconnect()

	def _status(self):
		return template('status', env_options=os.environ)

	# =============================================================================
	# MessageHub methods
	# =============================================================================
	def _createTopic(self):
		headers = {
			"Content-Type": "application/json",
			"X-Auth-Token": self.messagehubApikey
		}
		payload = {"name": "events"}

		r = requests.post(self.messagehubUrl + "/admin/topics", data=json.dumps(payload), headers=headers)
		print(r.text)

# Initialize the properties we need
parser = argparse.ArgumentParser()
parser.add_argument('-b', '--bluemix', required=False, action='store_true')
parser.add_argument('-c', '--config', required=False)
parser.add_argument('-a', '--messagehubApikey', required=False)
parser.add_argument('-u', '--messagehubUrl', required=False)

args, unknown = parser.parse_known_args()

server = Server(args)
server.start()<commit_msg>Update Kafka topic to use "orgId-events"<commit_after>import os
import json
import ibmiotf.application
import iso8601
import base64
from bottle import Bottle, template

import requests
from requests_futures.sessions import FuturesSession

import urllib
import argparse
import logging
from logging.handlers import RotatingFileHandler

class Server():

	def __init__(self, args):
		# Setup logging - Generate a default rotating file log handler and stream handler
		logFileName = 'connector-messagehub.log'
		fhFormatter = logging.Formatter('%(asctime)-25s %(name)-30s ' + ' %(levelname)-7s %(message)s')
		rfh = RotatingFileHandler(logFileName, mode='a', maxBytes=26214400 , backupCount=2, encoding=None, delay=True)
		rfh.setFormatter(fhFormatter)

		self.logger = logging.getLogger("server")
		self.logger.addHandler(rfh)
		self.logger.setLevel(logging.DEBUG)

		self.port = int(os.getenv('VCAP_APP_PORT', '8000'))
		self.host = str(os.getenv('VCAP_APP_HOST', 'localhost'))

		if args.bluemix == True:
			# Bluemix VCAP lookups
			application = json.loads(os.getenv('VCAP_APPLICATION'))
			service = json.loads(os.getenv('VCAP_SERVICES'))

			# IoTF
			self.options = ibmiotf.application.ParseConfigFromBluemixVCAP()

			# Cloudant
			self.messagehubApikey = service['messagehub'][0]['credentials']['api_key']
			self.messagehubUrl = service['messagehub'][0]['credentials']['kafka_rest_url']
		else:
			self.options = ibmiotf.application.ParseConfigFile(args.config)
			self.messagehubApikey = args.messagehubApikey
			self.messagehubUrl = args.messagehubUrl

		# Bottle
		self._app = Bottle()
		self._route()

		# Init MessageHub Topic & requests session
		self.topicName = self.options['org'] + "-events"
		self._createTopic()
		self.session = FuturesSession()

		# Init IOTF client
		self.client = ibmiotf.application.Client(self.options, logHandlers=[rfh])

	def _route(self):
		self._app.route('/', method="GET", callback=self._status)

	def myEventCallback(self, evt):
		#self.logger.info("%-33s%-30s%s" % (evt.timestamp.isoformat(), evt.device, evt.event + ": " + json.dumps(evt.data)))
		#self.logger.info(evt.data)

		headers = {
			"Content-Type": "application/json",
			"X-Auth-Token": self.messagehubApikey
		}

		message = {
			'typeId': evt.deviceType,
			'deviceId': evt.deviceId,
			'eventId': evt.event,
			'timestamp': evt.timestamp.isoformat(),
			'data': evt.data,
			'format': evt.format,
			'payload': base64.encodestring(evt.payload).decode('ascii')
		}

		payload = {
			'records': [{"value": json.dumps(message)}]
		}

		future = self.session.post(self.messagehubUrl + "/topics/" + self.topicName, data=json.dumps(payload), headers=headers)
		future.add_done_callback(self._eventRecordedCallback)

	def _eventRecordedCallback(self, future):
		response = future.result()
		if response.status_code not in [200]:
			self.logger.info("%s - %s: %s" % (response.url, response.status_code, response.text))

	def start(self):
		self.client.connect()
		self.client.deviceEventCallback = self.myEventCallback
		self.client.subscribeToDeviceEvents()
		self.logger.info("Serving at %s:%s" % (self.host, self.port))
		self._app.run(host=self.host, port=self.port)

	def stop(self):
		self.client.disconnect()

	def _status(self):
		return template('status', env_options=os.environ)

	# =============================================================================
	# MessageHub methods
	# =============================================================================
	def _createTopic(self):
		headers = {
			"Content-Type": "application/json",
			"X-Auth-Token": self.messagehubApikey
		}
		payload = {"name": self.topicName}

		r = requests.post(self.messagehubUrl + "/admin/topics", data=json.dumps(payload), headers=headers)
		print(r.text)

# Initialize the properties we need
parser = argparse.ArgumentParser()
parser.add_argument('-b', '--bluemix', required=False, action='store_true')
parser.add_argument('-c', '--config', required=False)
parser.add_argument('-a', '--messagehubApikey', required=False)
parser.add_argument('-u', '--messagehubUrl', required=False)

args, unknown = parser.parse_known_args()

server = Server(args)
server.start()<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Copyright 2012-2013 Universidad Politécnica de Madrid

# This file is part of Wirecloud.

# Wirecloud is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# Wirecloud is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with Wirecloud.  If not, see <http://www.gnu.org/licenses/>.

import json
from cStringIO import StringIO
import os

from django.conf import settings
from django.db import IntegrityError
from django.db.models import Q
from django.http import HttpResponse
from django.shortcuts import get_object_or_404
from django.utils import simplejson
from django.utils.encoding import smart_str
from django.utils.translation import ugettext as _
from django.views.static import serve

from wirecloud.catalogue.models import CatalogueResource
import wirecloud.catalogue.utils as catalogue_utils
from wirecloud.catalogue.views import iframe_error
from wirecloud.commons.baseviews import Resource
from wirecloud.commons.utils import downloader
from wirecloud.commons.utils.http import authentication_required, build_error_response, get_content_type, supported_request_mime_types
from wirecloud.commons.utils.template import TemplateParseException
from wirecloud.commons.utils.transaction import commit_on_http_success
from wirecloud.commons.utils.wgt import WgtFile
from wirecloud.platform.get_data import get_widget_data
from wirecloud.platform.markets.utils import get_market_managers
from wirecloud.platform.models import Widget, IWidget
from wirecloud.platform.localcatalogue.utils import install_resource_to_user, get_or_add_resource_from_available_marketplaces
from wirecloud.platform.widget.utils import get_or_add_widget_from_catalogue

class ResourceCollection(Resource):

    @authentication_required
    def read(self, request):

        resources = {}
        for resource in CatalogueResource.objects.filter(Q(public=True) | Q(users=request.user) | Q(groups=request.user.groups.all())):
            if resource.resource_type() == 'widget':
                try:
                    widget = resource.widget

                    resources[resource.local_uri_part] = get_widget_data(widget, request)
                    resources[resource.local_uri_part]['type'] = 'widget'

                except Widget.DoesNotExist:
                    pass

            else:
                options = json.loads(resource.json_description)
                resources[resource.local_uri_part] = options

        return HttpResponse(json.dumps(resources), mimetype='application/json; chatset=UTF-8')

    @authentication_required
    @iframe_error
    @supported_request_mime_types(('application/x-www-form-urlencoded', 'application/json', 'multipart/form-data', 'application/octet-stream'))
    @commit_on_http_success
    def create(self, request):

        force_create = False
        templateURL = None
        file_contents = None
        content_type = get_content_type(request)[0]
        if content_type == 'multipart/form-data':
            packaged = True
            force_create = request.POST.get('force_create', False) == 'true'
            if not 'file' in request.FILES:
                return build_error_response(request, 400, _('Missing file to upload'))

            downloaded_file = request.FILES['file']
            try:
                file_contents = WgtFile(downloaded_file)
            except:
                return build_error_response(request, 400, _('Bad resource file'))

        elif content_type == 'application/octet-stream':

            packaged = True
            downloaded_file = StringIO(request.raw_post_data)
            try:
                file_contents = WgtFile(downloaded_file)
            except:
                return build_error_response(request, 400, _('Bad resource file'))
        else:

            market_endpoint = None

            if content_type == 'application/json':
                try:
                    data = simplejson.loads(request.raw_post_data)
                except Exception, e:
                    msg = _("malformed json data: %s") % unicode(e)
                    return build_error_response(request, 400, msg)

                force_create = data.get('force_create', False)
                packaged = data.get('packaged', False)
                templateURL = data.get('template_uri')
                market_endpoint = data.get('market_endpoint', None)

            else:
                force_create = request.POST.get('force_create', False) == 'true'
                packaged = request.POST.get('packaged', False) == 'true'
                if 'url' in request.POST:
                    templateURL = request.POST['url']
                elif 'template_uri' in request.POST:
                    templateURL = request.POST['template_uri']

            if market_endpoint is not None:

                if 'name' not in market_endpoint:
                    msg = _('Missing market name')
                    return build_error_response(request, 400, msg)

                market_managers = get_market_managers(request.user)
                market_manager = market_managers[market_endpoint['name']]
                import ipdb; ipdb.set_trace()
                downloaded_file = market_manager.download_resource(request.user, templateURL, market_endpoint)

            else:

                try:
                    downloaded_file = downloader.download_http_content(templateURL)
                except:
                    return build_error_response(request, 409, _('Content cannot be downloaded'))

            if packaged:
                downloaded_file = StringIO(downloaded_file)
                file_contents = WgtFile(downloaded_file)
            else:
                file_contents = downloaded_file

        # TODO for now, install dependencies if force_create is true
        install_dep = force_create
        try:
            resource = install_resource_to_user(request.user, file_contents=file_contents, templateURL=templateURL, packaged=packaged, raise_conflicts=force_create)

        except TemplateParseException, e:

            return build_error_response(request, 400, unicode(e.msg))

        except IntegrityError:

            return build_error_response(request, 409, _('Resource already exists'))

        if resource.resource_type() == 'widget':

            data = get_widget_data(resource.widget, request)
            data['type'] = 'widget'
            if install_dep:
                return HttpResponse(simplejson.dumps((data,)), status=201, mimetype='application/json; charset=UTF-8')
            else:
                return HttpResponse(simplejson.dumps(data), status=201, mimetype='application/json; charset=UTF-8')

        elif install_dep and resource.resource_type() == 'mashup':
            resources = [json.loads(resource.json_description)]
            workspace_info = json.loads(resource.json_description)
            for tab_entry in workspace_info['tabs']:
                for resource in tab_entry['resources']:
                    widget = get_or_add_widget_from_catalogue(resource.get('vendor'), resource.get('name'), resource.get('version'), request.user)
                    widget_data = get_widget_data(widget, request)
                    widget_data['type'] = 'widget'
                    resources.append(widget_data)

            for id_, op in workspace_info['wiring']['operators'].iteritems():
                op_id_args = op['name'].split('/')
                op_id_args.append(request.user)
                operator = get_or_add_resource_from_available_marketplaces(*op_id_args)
                resources.append(json.loads(operator.json_description))

            return HttpResponse(simplejson.dumps(resources), status=201, mimetype='application/json; charset=UTF-8')

        elif install_dep:
            return HttpResponse('[' + resource.json_description + ']', status=201, mimetype='application/json; charset=UTF-8')
        else:
            return HttpResponse(resource.json_description, status=201, mimetype='application/json; charset=UTF-8')

class ResourceEntry(Resource):

    @authentication_required
    def read(self, request, vendor, name, version):

        file_name = '_'.join((vendor, name, version)) + '.wgt'
        base_dir = catalogue_utils.wgt_deployer.get_base_dir(vendor, name, version)
        local_path = os.path.normpath(os.path.join(base_dir, file_name))

        if not os.path.isfile(local_path):
            return HttpResponse(status=404)

        if not getattr(settings, 'USE_XSENDFILE', False):
            return serve(request, local_path, document_root='/')
        else:
            response = HttpResponse()
            response['X-Sendfile'] = smart_str(local_path)
            return response

    @authentication_required
    @commit_on_http_success
    def delete(self, request, vendor, name, version):

        resource = get_object_or_404(CatalogueResource, vendor=vendor, short_name=name, version=version)
        resource.users.remove(request.user)

        if resource.resource_type() == 'widget':
            result = {'removedIWidgets': []}

            query = Widget.objects.filter(resource=resource)
            if query.exists():

                widget = query.get()

                # Remove all iwidgets that matches the resource
                iwidgets = IWidget.objects.filter(widget=widget, tab__workspace__creator=request.user)
                for iwidget in iwidgets:
                    result['removedIWidgets'].append(iwidget.id)
                    iwidget.delete()

            if request.GET.get('affected', 'false').lower() == 'true':
                return HttpResponse(simplejson.dumps(result), mimetype='application/json; charset=UTF-8')

        return HttpResponse(status=204)
<commit_msg>Remove forgotten ipdb<commit_after># -*- coding: utf-8 -*-

# Copyright 2012-2013 Universidad Politécnica de Madrid

# This file is part of Wirecloud.

# Wirecloud is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# Wirecloud is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with Wirecloud.  If not, see <http://www.gnu.org/licenses/>.

import json
from cStringIO import StringIO
import os

from django.conf import settings
from django.db import IntegrityError
from django.db.models import Q
from django.http import HttpResponse
from django.shortcuts import get_object_or_404
from django.utils import simplejson
from django.utils.encoding import smart_str
from django.utils.translation import ugettext as _
from django.views.static import serve

from wirecloud.catalogue.models import CatalogueResource
import wirecloud.catalogue.utils as catalogue_utils
from wirecloud.catalogue.views import iframe_error
from wirecloud.commons.baseviews import Resource
from wirecloud.commons.utils import downloader
from wirecloud.commons.utils.http import authentication_required, build_error_response, get_content_type, supported_request_mime_types
from wirecloud.commons.utils.template import TemplateParseException
from wirecloud.commons.utils.transaction import commit_on_http_success
from wirecloud.commons.utils.wgt import WgtFile
from wirecloud.platform.get_data import get_widget_data
from wirecloud.platform.markets.utils import get_market_managers
from wirecloud.platform.models import Widget, IWidget
from wirecloud.platform.localcatalogue.utils import install_resource_to_user, get_or_add_resource_from_available_marketplaces
from wirecloud.platform.widget.utils import get_or_add_widget_from_catalogue

class ResourceCollection(Resource):

    @authentication_required
    def read(self, request):

        resources = {}
        for resource in CatalogueResource.objects.filter(Q(public=True) | Q(users=request.user) | Q(groups=request.user.groups.all())):
            if resource.resource_type() == 'widget':
                try:
                    widget = resource.widget

                    resources[resource.local_uri_part] = get_widget_data(widget, request)
                    resources[resource.local_uri_part]['type'] = 'widget'

                except Widget.DoesNotExist:
                    pass

            else:
                options = json.loads(resource.json_description)
                resources[resource.local_uri_part] = options

        return HttpResponse(json.dumps(resources), mimetype='application/json; chatset=UTF-8')

    @authentication_required
    @iframe_error
    @supported_request_mime_types(('application/x-www-form-urlencoded', 'application/json', 'multipart/form-data', 'application/octet-stream'))
    @commit_on_http_success
    def create(self, request):

        force_create = False
        templateURL = None
        file_contents = None
        content_type = get_content_type(request)[0]
        if content_type == 'multipart/form-data':
            packaged = True
            force_create = request.POST.get('force_create', False) == 'true'
            if not 'file' in request.FILES:
                return build_error_response(request, 400, _('Missing file to upload'))

            downloaded_file = request.FILES['file']
            try:
                file_contents = WgtFile(downloaded_file)
            except:
                return build_error_response(request, 400, _('Bad resource file'))

        elif content_type == 'application/octet-stream':

            packaged = True
            downloaded_file = StringIO(request.raw_post_data)
            try:
                file_contents = WgtFile(downloaded_file)
            except:
                return build_error_response(request, 400, _('Bad resource file'))
        else:

            market_endpoint = None

            if content_type == 'application/json':
                try:
                    data = simplejson.loads(request.raw_post_data)
                except Exception, e:
                    msg = _("malformed json data: %s") % unicode(e)
                    return build_error_response(request, 400, msg)

                force_create = data.get('force_create', False)
                packaged = data.get('packaged', False)
                templateURL = data.get('template_uri')
                market_endpoint = data.get('market_endpoint', None)

            else:
                force_create = request.POST.get('force_create', False) == 'true'
                packaged = request.POST.get('packaged', False) == 'true'
                if 'url' in request.POST:
                    templateURL = request.POST['url']
                elif 'template_uri' in request.POST:
                    templateURL = request.POST['template_uri']

            if market_endpoint is not None:

                if 'name' not in market_endpoint:
                    msg = _('Missing market name')
                    return build_error_response(request, 400, msg)

                market_managers = get_market_managers(request.user)
                market_manager = market_managers[market_endpoint['name']]
                downloaded_file = market_manager.download_resource(request.user, templateURL, market_endpoint)

            else:

                try:
                    downloaded_file = downloader.download_http_content(templateURL)
                except:
                    return build_error_response(request, 409, _('Content cannot be downloaded'))

            if packaged:
                downloaded_file = StringIO(downloaded_file)
                file_contents = WgtFile(downloaded_file)
            else:
                file_contents = downloaded_file

        # TODO for now, install dependencies if force_create is true
        install_dep = force_create
        try:
            resource = install_resource_to_user(request.user, file_contents=file_contents, templateURL=templateURL, packaged=packaged, raise_conflicts=force_create)

        except TemplateParseException, e:

            return build_error_response(request, 400, unicode(e.msg))

        except IntegrityError:

            return build_error_response(request, 409, _('Resource already exists'))

        if resource.resource_type() == 'widget':

            data = get_widget_data(resource.widget, request)
            data['type'] = 'widget'
            if install_dep:
                return HttpResponse(simplejson.dumps((data,)), status=201, mimetype='application/json; charset=UTF-8')
            else:
                return HttpResponse(simplejson.dumps(data), status=201, mimetype='application/json; charset=UTF-8')

        elif install_dep and resource.resource_type() == 'mashup':
            resources = [json.loads(resource.json_description)]
            workspace_info = json.loads(resource.json_description)
            for tab_entry in workspace_info['tabs']:
                for resource in tab_entry['resources']:
                    widget = get_or_add_widget_from_catalogue(resource.get('vendor'), resource.get('name'), resource.get('version'), request.user)
                    widget_data = get_widget_data(widget, request)
                    widget_data['type'] = 'widget'
                    resources.append(widget_data)

            for id_, op in workspace_info['wiring']['operators'].iteritems():
                op_id_args = op['name'].split('/')
                op_id_args.append(request.user)
                operator = get_or_add_resource_from_available_marketplaces(*op_id_args)
                resources.append(json.loads(operator.json_description))

            return HttpResponse(simplejson.dumps(resources), status=201, mimetype='application/json; charset=UTF-8')

        elif install_dep:
            return HttpResponse('[' + resource.json_description + ']', status=201, mimetype='application/json; charset=UTF-8')
        else:
            return HttpResponse(resource.json_description, status=201, mimetype='application/json; charset=UTF-8')

class ResourceEntry(Resource):

    @authentication_required
    def read(self, request, vendor, name, version):

        file_name = '_'.join((vendor, name, version)) + '.wgt'
        base_dir = catalogue_utils.wgt_deployer.get_base_dir(vendor, name, version)
        local_path = os.path.normpath(os.path.join(base_dir, file_name))

        if not os.path.isfile(local_path):
            return HttpResponse(status=404)

        if not getattr(settings, 'USE_XSENDFILE', False):
            return serve(request, local_path, document_root='/')
        else:
            response = HttpResponse()
            response['X-Sendfile'] = smart_str(local_path)
            return response

    @authentication_required
    @commit_on_http_success
    def delete(self, request, vendor, name, version):

        resource = get_object_or_404(CatalogueResource, vendor=vendor, short_name=name, version=version)
        resource.users.remove(request.user)

        if resource.resource_type() == 'widget':
            result = {'removedIWidgets': []}

            query = Widget.objects.filter(resource=resource)
            if query.exists():

                widget = query.get()

                # Remove all iwidgets that matches the resource
                iwidgets = IWidget.objects.filter(widget=widget, tab__workspace__creator=request.user)
                for iwidget in iwidgets:
                    result['removedIWidgets'].append(iwidget.id)
                    iwidget.delete()

            if request.GET.get('affected', 'false').lower() == 'true':
                return HttpResponse(simplejson.dumps(result), mimetype='application/json; charset=UTF-8')

        return HttpResponse(status=204)
<|endoftext|>


<commit_before>"""
Statistical tests to be used in conjunction with the models

Notes
-----
These functions haven't been formally tested.
"""

from scipy import stats
import numpy as np

# TODO: these are pretty straightforward but they should be tested
def durbin_watson(resids, axis=0):
    """
    Calculates the Durbin-Watson statistic

    Parameters
    -----------
    resids : array-like

    Returns
    --------
    Durbin Watson statistic.  This is defined as
    sum_(t=2)^(T)((e_t - e_(t-1))^(2))/sum_(t=1)^(T)e_t^(2)
    """
    resids = np.asarray(resids)
    diff_resids = np.diff(resids, 1, axis=axis)
    dw = np.sum(diff_resids**2, axis=axis) / np.sum(resids**2, axis=axis)
    return dw

def omni_normtest(resids, axis=0):
    """
    Omnibus test for normality

    Parameters
    -----------
    resid : array-like
    axis : int, optional
        Default is 0

    Returns
    -------
    Chi^2 score, two-tail probability
    """
    #TODO: change to exception in summary branch and catch in summary()
    #behavior changed between scipy 0.9 and 0.10
    resids = np.asarray(resids)
    n = resids.shape[axis]
    if n < 8:
        from warnings import warn

        warn("omni_normtest is not valid with less than 8 observations; %i "
             "samples were given." % int(n))
        return np.nan, np.nan

    return stats.normaltest(resids, axis=axis)

def jarque_bera(resids, axis=0):
    """
    Calculate residual skewness, kurtosis, and do the JB test for normality

    Parameters
    -----------
    resids : array-like
    axis : int, optional
        Default is 0

    Returns
    -------
    JB, JBpv, skew, kurtosis

    JB = n/6*(S^2 + (K-3)^2/4)

    JBpv is the Chi^2 two-tail probability value

    skew is the measure of skewness

    kurtosis is the measure of kurtosis

    """
    resids = np.asarray(resids)
    # Calculate residual skewness and kurtosis
    skew = stats.skew(resids, axis=axis)
    kurtosis = 3 + stats.kurtosis(resids, axis=axis)

    # Calculate the Jarque-Bera test for normality
    n = resids.shape[axis]
    jb = (n / 6.) * (skew**2 + (1 / 4.) * (kurtosis - 3)**2)
    jb_pv = stats.chi2.sf(jb, 2)

    return jb, jb_pv, skew, kurtosis

def robust_skewness(y, axis=0):
    """
    Calculates the four skewness measures in Kim & White

    Parameters
    ----------
    y : array-like

    axis : int or None, optional
        Axis along which the skewness measures are computed.  If `None`, the
        entire array is used.

    Returns
    -------
    sk1 : ndarray
          The standard skewness estimator.
    sk2 : ndarray
          Skewness estimator based on quartiles.
    sk3 : ndarray
          Skewness estimator based on mean-median difference, standardized by
          absolute deviation.
    sk4 : ndarray
          Skewness estimator based on mean-median difference, standardized by
          standard deviation.

    Notes
    -----
    The robust skewness measures are defined

    .. math ::

        SK_{2}=\\frac{\\left(q_{.75}-q_{.5}\\right)
        -\\left(q_{.5}-q_{.25}\\right)}{q_{.75}-q_{.25}}

    .. math ::

        SK_{3}=\\frac{\\mu-\\hat{q}_{0.5}}
        {\\hat{E}\\left[\\left|y-\\hat{\\mu}\\right|\\right]}

    .. math ::

        SK_{4}=\\frac{\\mu-\\hat{q}_{0.5}}{\\hat{\\sigma}}

    .. [1] Tae-Hwan Kim and Halbert White, "On more robust estimation of
    skewness and kurtosis," Finance Research Letters, vol. 1, pp. 56-73,
    March 2004.
    """

    if axis is None:
        y = y.ravel()
        axis = 0

    y = np.sort(y, axis)

    q1, q2, q3 = np.percentile(y, [25.0, 50.0, 75.0], axis=axis)

    mu = y.mean(axis)
    shape = (y.size,)
    if axis is not None:
        shape = list(mu.shape)
        shape.insert(axis, 1)
        shape = tuple(shape)

    mu_b = np.reshape(mu, shape)
    q2_b = np.reshape(q2, shape)

    sigma = np.mean(((y - mu_b)**2), axis)

    sk1 = stats.skew(y, axis=axis)
    sk2 = (q1 + q3 - 2.0 * q2) / (q3 - q1)
    sk3 = (mu - q2) / np.mean(abs(y - q2_b), axis=axis)
    sk4 = (mu - q2) / sigma

    return sk1, sk2, sk3, sk4

def _kr3(y, alpha=5.0, beta=50.0):
    """
    KR3 estimator from Kim & White

    Parameters
    ----------
    y : array-like, 1-d
    alpha : float, optional
        Lower cut-off for measuring expectation in tail.
    beta :  float, optional
        Lower cut-off for measuring expectation in center.

    Returns
    -------
    kr3 : float
        Robust kurtosis estimator based on standardized lower- and upper-tail
        expected values

    Notes
    -----
    .. [1] Tae-Hwan Kim and Halbert White, "On more robust estimation of
    skewness and kurtosis," Finance Research Letters, vol. 1, pp. 56-73,
    March 2004.
    """
    perc = (alpha, 100.0 - alpha, beta, 100.0 - beta)
    lower_alpha, upper_alpha, lower_beta, upper_beta = np.percentile(y, perc)
    l_alpha = np.mean(y[y < lower_alpha])
    u_alpha = np.mean(y[y > upper_alpha])

    l_beta = np.mean(y[y < lower_beta])
    u_beta = np.mean(y[y > upper_beta])

    return (u_alpha - l_alpha) / (u_beta - l_beta)

def expected_robust_kurtosis(ab=(5.0, 50.0), dg=(2.5, 25.0)):
    """
    Calculates the expected value of the robust kurtosis measures in Kim and
    White assuming the data are normally distributed.

    Parameters
    ----------
    ab: iterable, optional
        Contains 100*(alpha, beta) in the kr3 measure where alpha is the tail
        quantile cut-off for measuring the extreme tail and beta is the central
        quantile cutoff for the standardization of the measure
    db: iterable, optional
        Contains 100*(delta, gamma) in the kr4 measure where delta is the tail
        quantile for measuring extreme values and gamma is the central quantile
        used in the the standardization of the measure

    Returns
    -------
    ekr: array, 4-element
        Contains the expected values of the 4 robust kurtosis measures

    Notes
    -----
    See `robust_kurtosis` for definitions of the robust kurtosis measures
    """

    alpha, beta = ab
    delta, gamma = dg
    expected_value = np.zeros(4)
    ppf = stats.norm.ppf
    pdf = stats.norm.pdf
    q1, q2, q3, q5, q6, q7 = ppf(np.array((1.0, 2.0, 3.0, 5.0, 6.0, 7.0)) / 8)
    expected_value[0] = 3

    expected_value[1] = ((q7 - q5) + (q3 - q1)) / (q6 - q2)

    q_alpha, q_beta = ppf(np.array((alpha / 100.0, beta / 100.0)))
    expected_value[2] = (2 * pdf(q_alpha) / alpha) / (2 * pdf(q_beta) / beta)

    q_delta, q_gamma = ppf(np.array((delta / 100.0, gamma / 100.0)))
    expected_value[3] = (-2.0 * q_delta) / (-2.0 * q_gamma)

    return expected_value

def robust_kurtosis(y, axis=0, ab=(5.0, 50.0), dg=(2.5, 25.0), excess=True):
    """
    Calculates the four kurtosis measures in Kim & White

    Parameters
    ----------
    y : array-like
    axis : int or None, optional
        Axis along which the kurtoses are computed.  If `None`, the
        entire array is used.
    ab: iterable, optional
        Contains 100*(alpha, beta) in the kr3 measure where alpha is the tail
        quantile cut-off for measuring the extreme tail and beta is the central
        quantile cutoff for the standardization of the measure
    db: iterable, optional
        Contains 100*(delta, gamma) in the kr4 measure where delta is the tail
        quantile for measuring extreme values and gamma is the central quantile
        used in the the standardization of the measure
    excess : bool, optional
        If true (default), computed values are excess of those for a standard
        normal distribution.

    Returns
    -------
    kr1 : ndarray
          The standard kurtosis estimator.
    kr2 : ndarray
          Kurtosis estimator based on octiles.
    kr3 : ndarray
          Kurtosis estimators based on exceedence expectations.
    kr4 : ndarray
          Kurtosis measure based on the spread between high and low quantiles.

    Notes
    -----
    The robust kurtosis measures are defined

    .. math::

        KR_{2}=\\frac{\\left(\\hat{q}_{.875}-\\hat{q}_{.625}\\right)
        +\\left(\\hat{q}_{.375}-\\hat{q}_{.125}\\right)}
        {\\hat{q}_{.75}-\\hat{q}_{.25}}

    .. math::

        KR_{3}=\\frac{\\hat{E}\\left(y|y>\\hat{q}_{1-\\alpha}\\right)
        -\\hat{E}\\left(y|y<\\hat{q}_{\\alpha}\\right)}
        {\\hat{E}\\left(y|y>\\hat{q}_{1-\\beta}\\right)
        -\\hat{E}\\left(y|y<\\hat{q}_{\\beta}\\right)}

    .. math::

        KR_{4}=\\frac{\\hat{q}_{1-\\delta}-\\hat{q}_{\\delta}}
        {\\hat{q}_{1-\\gamma}-\\hat{q}_{\\gamma}}

    where :math:`\\hat{q}_{p}` is the estimated quantile at :math:`p`.

    .. [1] Tae-Hwan Kim and Halbert White, "On more robust estimation of
    skewness and kurtosis," Finance Research Letters, vol. 1, pp. 56-73,
    March 2004.
    """
    if (axis is None or
            (y.squeeze().ndim == 1 and y.ndim != 1)):
        y = y.ravel()
        axis = 0

    alpha, beta = ab
    delta, gamma = dg

    perc = (12.5, 25.0, 37.5, 62.5, 75.0, 87.5,
            delta, 100.0 - delta, gamma, 100.0 - gamma)
    e1, e2, e3, e5, e6, e7, fd, f1md, fg, f1mg = np.percentile(y, perc,
                                                               axis=axis)

    expected_value = expected_robust_kurtosis(ab, dg) if excess else np.zeros(4)

    kr1 = stats.kurtosis(y, axis, False) - expected_value[0]
    kr2 = ((e7 - e5) + (e3 - e1)) / (e6 - e2) - expected_value[1]
    if y.ndim == 1:
        kr3 = _kr3(y, alpha, beta)
    else:
        kr3 = np.apply_along_axis(_kr3, axis, y, alpha, beta)
    kr3 -= expected_value[2]
    kr4 = (f1md - fd) / (f1mg - fg) - expected_value[3]
    return kr1, kr2, kr3, kr4

def _medcouple_1d(y):
    """
    Calculates the medcouple robust measure of skew.

    Parameters
    ----------
    y : array-like, 1-d

    Returns
    -------
    mc : float
        The medcouple statistic

    Notes
    -----
    The current algorithm requires a O(N**2) memory allocations, and so may
    not work for very large arrays (N>10000).

    .. [1] M. Huberta and E. Vandervierenb, "An adjusted boxplot for skewed
    distributions" Computational Statistics & Data Analysis, vol. 52,
    pp. 5186-5201, August 2008.
    """

    # Parameter changes the algorithm to the slower for large n

    y = np.squeeze(np.asarray(y))
    if y.ndim != 1:
        raise ValueError("y must be squeezable to a 1-d array")

    y = np.sort(y)

    n = y.shape[0]
    if n % 2 == 0:
        mf = (y[n // 2 - 1] + y[n // 2]) / 2
    else:
        mf = y[(n - 1) // 2]

    z = y - mf
    lower = z[z <= 0.0]
    upper = z[z >= 0.0]
    upper = upper[:, None]
    standardization = upper - lower
    is_zero = np.logical_and(lower == 0.0, upper == 0.0)
    standardization[is_zero] = np.inf
    spread = upper + lower
    return np.median(spread / standardization)

def medcouple(y, axis=0):
    """
    Calculates the medcouple robust measure of skew.

    Parameters
    ----------
    y : array-like
    axis : int or None, optional
        Axis along which the medcouple statistic is computed.  If `None`, the
        entire array is used.

    Returns
    -------
    mc : ndarray
        The medcouple statistic with the same shape as `y`, with the specified
        axis removed.

    Notes
    -----
    The current algorithm requires a O(N**2) memory allocations, and so may
    not work for very large arrays (N>10000).

    .. [1] M. Huberta and E. Vandervierenb, "An adjusted boxplot for skewed
    distributions" Computational Statistics & Data Analysis, vol. 52,
    pp. 5186-5201, August 2008.
    """
    if axis is None:
        return _medcouple_1d(y.ravel())

    return np.apply_along_axis(_medcouple_1d, axis, y)<commit_msg>DOC: Expand documentation notes<commit_after>"""
Statistical tests to be used in conjunction with the models

Notes
-----
These functions haven't been formally tested.
"""

from scipy import stats
import numpy as np

# TODO: these are pretty straightforward but they should be tested
def durbin_watson(resids, axis=0):
    """
    Calculates the Durbin-Watson statistic

    Parameters
    -----------
    resids : array-like

    Returns
    --------
    dw : float, array-like

    The Durbin-Watson statistic.

    Notes
    -----
    The null hypothesis of the test is that there is no serial correlation.
    The Durbin-Watson test statistics is defined as:

    .. math::

       \sum_{t=2}^T((e_t - e_{t-1})^2)/\sum_{t=1}^Te_t^2

    The test statistic is approximately equal to 2*(1-r) where ``r`` is the
    sample autocorrelation of the residuals. Thus, for r == 0, indicating no
    serial correlation, the test statistic equals 2. This statistic will
    always be between 0 and 4. The closer to 0 the statistic, the more
    evidence for positive serial correlation. The closer to 4, the more
    evidence for negative serial correlation.
    """
    resids = np.asarray(resids)
    diff_resids = np.diff(resids, 1, axis=axis)
    dw = np.sum(diff_resids**2, axis=axis) / np.sum(resids**2, axis=axis)
    return dw

def omni_normtest(resids, axis=0):
    """
    Omnibus test for normality

    Parameters
    -----------
    resid : array-like
    axis : int, optional
        Default is 0

    Returns
    -------
    Chi^2 score, two-tail probability
    """
    #TODO: change to exception in summary branch and catch in summary()
    #behavior changed between scipy 0.9 and 0.10
    resids = np.asarray(resids)
    n = resids.shape[axis]
    if n < 8:
        from warnings import warn

        warn("omni_normtest is not valid with less than 8 observations; %i "
             "samples were given." % int(n))
        return np.nan, np.nan

    return stats.normaltest(resids, axis=axis)

def jarque_bera(resids, axis=0):
    """
    Calculate residual skewness, kurtosis, and do the JB test for normality

    Parameters
    -----------
    resids : array-like
    axis : int, optional
        Default is 0

    Returns
    -------
    JB, JBpv, skew, kurtosis

    JB = n/6*(S^2 + (K-3)^2/4)

    JBpv is the Chi^2 two-tail probability value

    skew is the measure of skewness

    kurtosis is the measure of kurtosis

    """
    resids = np.asarray(resids)
    # Calculate residual skewness and kurtosis
    skew = stats.skew(resids, axis=axis)
    kurtosis = 3 + stats.kurtosis(resids, axis=axis)

    # Calculate the Jarque-Bera test for normality
    n = resids.shape[axis]
    jb = (n / 6.) * (skew**2 + (1 / 4.) * (kurtosis - 3)**2)
    jb_pv = stats.chi2.sf(jb, 2)

    return jb, jb_pv, skew, kurtosis

def robust_skewness(y, axis=0):
    """
    Calculates the four skewness measures in Kim & White

    Parameters
    ----------
    y : array-like

    axis : int or None, optional
        Axis along which the skewness measures are computed.  If `None`, the
        entire array is used.

    Returns
    -------
    sk1 : ndarray
          The standard skewness estimator.
    sk2 : ndarray
          Skewness estimator based on quartiles.
    sk3 : ndarray
          Skewness estimator based on mean-median difference, standardized by
          absolute deviation.
    sk4 : ndarray
          Skewness estimator based on mean-median difference, standardized by
          standard deviation.

    Notes
    -----
    The robust skewness measures are defined

    .. math ::

        SK_{2}=\\frac{\\left(q_{.75}-q_{.5}\\right)
        -\\left(q_{.5}-q_{.25}\\right)}{q_{.75}-q_{.25}}

    .. math ::

        SK_{3}=\\frac{\\mu-\\hat{q}_{0.5}}
        {\\hat{E}\\left[\\left|y-\\hat{\\mu}\\right|\\right]}

    .. math ::

        SK_{4}=\\frac{\\mu-\\hat{q}_{0.5}}{\\hat{\\sigma}}

    .. [1] Tae-Hwan Kim and Halbert White, "On more robust estimation of
    skewness and kurtosis," Finance Research Letters, vol. 1, pp. 56-73,
    March 2004.
    """

    if axis is None:
        y = y.ravel()
        axis = 0

    y = np.sort(y, axis)

    q1, q2, q3 = np.percentile(y, [25.0, 50.0, 75.0], axis=axis)

    mu = y.mean(axis)
    shape = (y.size,)
    if axis is not None:
        shape = list(mu.shape)
        shape.insert(axis, 1)
        shape = tuple(shape)

    mu_b = np.reshape(mu, shape)
    q2_b = np.reshape(q2, shape)

    sigma = np.mean(((y - mu_b)**2), axis)

    sk1 = stats.skew(y, axis=axis)
    sk2 = (q1 + q3 - 2.0 * q2) / (q3 - q1)
    sk3 = (mu - q2) / np.mean(abs(y - q2_b), axis=axis)
    sk4 = (mu - q2) / sigma

    return sk1, sk2, sk3, sk4

def _kr3(y, alpha=5.0, beta=50.0):
    """
    KR3 estimator from Kim & White

    Parameters
    ----------
    y : array-like, 1-d
    alpha : float, optional
        Lower cut-off for measuring expectation in tail.
    beta :  float, optional
        Lower cut-off for measuring expectation in center.

    Returns
    -------
    kr3 : float
        Robust kurtosis estimator based on standardized lower- and upper-tail
        expected values

    Notes
    -----
    .. [1] Tae-Hwan Kim and Halbert White, "On more robust estimation of
    skewness and kurtosis," Finance Research Letters, vol. 1, pp. 56-73,
    March 2004.
    """
    perc = (alpha, 100.0 - alpha, beta, 100.0 - beta)
    lower_alpha, upper_alpha, lower_beta, upper_beta = np.percentile(y, perc)
    l_alpha = np.mean(y[y < lower_alpha])
    u_alpha = np.mean(y[y > upper_alpha])

    l_beta = np.mean(y[y < lower_beta])
    u_beta = np.mean(y[y > upper_beta])

    return (u_alpha - l_alpha) / (u_beta - l_beta)

def expected_robust_kurtosis(ab=(5.0, 50.0), dg=(2.5, 25.0)):
    """
    Calculates the expected value of the robust kurtosis measures in Kim and
    White assuming the data are normally distributed.

    Parameters
    ----------
    ab: iterable, optional
        Contains 100*(alpha, beta) in the kr3 measure where alpha is the tail
        quantile cut-off for measuring the extreme tail and beta is the central
        quantile cutoff for the standardization of the measure
    db: iterable, optional
        Contains 100*(delta, gamma) in the kr4 measure where delta is the tail
        quantile for measuring extreme values and gamma is the central quantile
        used in the the standardization of the measure

    Returns
    -------
    ekr: array, 4-element
        Contains the expected values of the 4 robust kurtosis measures

    Notes
    -----
    See `robust_kurtosis` for definitions of the robust kurtosis measures
    """

    alpha, beta = ab
    delta, gamma = dg
    expected_value = np.zeros(4)
    ppf = stats.norm.ppf
    pdf = stats.norm.pdf
    q1, q2, q3, q5, q6, q7 = ppf(np.array((1.0, 2.0, 3.0, 5.0, 6.0, 7.0)) / 8)
    expected_value[0] = 3

    expected_value[1] = ((q7 - q5) + (q3 - q1)) / (q6 - q2)

    q_alpha, q_beta = ppf(np.array((alpha / 100.0, beta / 100.0)))
    expected_value[2] = (2 * pdf(q_alpha) / alpha) / (2 * pdf(q_beta) / beta)

    q_delta, q_gamma = ppf(np.array((delta / 100.0, gamma / 100.0)))
    expected_value[3] = (-2.0 * q_delta) / (-2.0 * q_gamma)

    return expected_value

def robust_kurtosis(y, axis=0, ab=(5.0, 50.0), dg=(2.5, 25.0), excess=True):
    """
    Calculates the four kurtosis measures in Kim & White

    Parameters
    ----------
    y : array-like
    axis : int or None, optional
        Axis along which the kurtoses are computed.  If `None`, the
        entire array is used.
    ab: iterable, optional
        Contains 100*(alpha, beta) in the kr3 measure where alpha is the tail
        quantile cut-off for measuring the extreme tail and beta is the central
        quantile cutoff for the standardization of the measure
    db: iterable, optional
        Contains 100*(delta, gamma) in the kr4 measure where delta is the tail
        quantile for measuring extreme values and gamma is the central quantile
        used in the the standardization of the measure
    excess : bool, optional
        If true (default), computed values are excess of those for a standard
        normal distribution.

    Returns
    -------
    kr1 : ndarray
          The standard kurtosis estimator.
    kr2 : ndarray
          Kurtosis estimator based on octiles.
    kr3 : ndarray
          Kurtosis estimators based on exceedence expectations.
    kr4 : ndarray
          Kurtosis measure based on the spread between high and low quantiles.

    Notes
    -----
    The robust kurtosis measures are defined

    .. math::

        KR_{2}=\\frac{\\left(\\hat{q}_{.875}-\\hat{q}_{.625}\\right)
        +\\left(\\hat{q}_{.375}-\\hat{q}_{.125}\\right)}
        {\\hat{q}_{.75}-\\hat{q}_{.25}}

    .. math::

        KR_{3}=\\frac{\\hat{E}\\left(y|y>\\hat{q}_{1-\\alpha}\\right)
        -\\hat{E}\\left(y|y<\\hat{q}_{\\alpha}\\right)}
        {\\hat{E}\\left(y|y>\\hat{q}_{1-\\beta}\\right)
        -\\hat{E}\\left(y|y<\\hat{q}_{\\beta}\\right)}

    .. math::

        KR_{4}=\\frac{\\hat{q}_{1-\\delta}-\\hat{q}_{\\delta}}
        {\\hat{q}_{1-\\gamma}-\\hat{q}_{\\gamma}}

    where :math:`\\hat{q}_{p}` is the estimated quantile at :math:`p`.

    .. [1] Tae-Hwan Kim and Halbert White, "On more robust estimation of
    skewness and kurtosis," Finance Research Letters, vol. 1, pp. 56-73,
    March 2004.
    """
    if (axis is None or
            (y.squeeze().ndim == 1 and y.ndim != 1)):
        y = y.ravel()
        axis = 0

    alpha, beta = ab
    delta, gamma = dg

    perc = (12.5, 25.0, 37.5, 62.5, 75.0, 87.5,
            delta, 100.0 - delta, gamma, 100.0 - gamma)
    e1, e2, e3, e5, e6, e7, fd, f1md, fg, f1mg = np.percentile(y, perc,
                                                               axis=axis)

    expected_value = expected_robust_kurtosis(ab, dg) if excess else np.zeros(4)

    kr1 = stats.kurtosis(y, axis, False) - expected_value[0]
    kr2 = ((e7 - e5) + (e3 - e1)) / (e6 - e2) - expected_value[1]
    if y.ndim == 1:
        kr3 = _kr3(y, alpha, beta)
    else:
        kr3 = np.apply_along_axis(_kr3, axis, y, alpha, beta)
    kr3 -= expected_value[2]
    kr4 = (f1md - fd) / (f1mg - fg) - expected_value[3]
    return kr1, kr2, kr3, kr4

def _medcouple_1d(y):
    """
    Calculates the medcouple robust measure of skew.

    Parameters
    ----------
    y : array-like, 1-d

    Returns
    -------
    mc : float
        The medcouple statistic

    Notes
    -----
    The current algorithm requires a O(N**2) memory allocations, and so may
    not work for very large arrays (N>10000).

    .. [1] M. Huberta and E. Vandervierenb, "An adjusted boxplot for skewed
    distributions" Computational Statistics & Data Analysis, vol. 52,
    pp. 5186-5201, August 2008.
    """

    # Parameter changes the algorithm to the slower for large n

    y = np.squeeze(np.asarray(y))
    if y.ndim != 1:
        raise ValueError("y must be squeezable to a 1-d array")

    y = np.sort(y)

    n = y.shape[0]
    if n % 2 == 0:
        mf = (y[n // 2 - 1] + y[n // 2]) / 2
    else:
        mf = y[(n - 1) // 2]

    z = y - mf
    lower = z[z <= 0.0]
    upper = z[z >= 0.0]
    upper = upper[:, None]
    standardization = upper - lower
    is_zero = np.logical_and(lower == 0.0, upper == 0.0)
    standardization[is_zero] = np.inf
    spread = upper + lower
    return np.median(spread / standardization)

def medcouple(y, axis=0):
    """
    Calculates the medcouple robust measure of skew.

    Parameters
    ----------
    y : array-like
    axis : int or None, optional
        Axis along which the medcouple statistic is computed.  If `None`, the
        entire array is used.

    Returns
    -------
    mc : ndarray
        The medcouple statistic with the same shape as `y`, with the specified
        axis removed.

    Notes
    -----
    The current algorithm requires a O(N**2) memory allocations, and so may
    not work for very large arrays (N>10000).

    .. [1] M. Huberta and E. Vandervierenb, "An adjusted boxplot for skewed
    distributions" Computational Statistics & Data Analysis, vol. 52,
    pp. 5186-5201, August 2008.
    """
    if axis is None:
        return _medcouple_1d(y.ravel())

    return np.apply_along_axis(_medcouple_1d, axis, y)
<|endoftext|>


<commit_before>import re

from re import _pattern_type as PatternType
from typing import List, Union

__all__ = [
    'Lexeme', 'Indent', 'NewLine', 'Number', 'Name', 'Class', 'Operator',
    'LexerError', 'Lexer',
    'lex_file', 'lex_line',
]

# FIXME: Ambiguity between names with symbol endings and operators with those same symbols.
# FIXME: Ambiguity between function call and postfix Operator `()`
# TODO: Handle strings.
# TODO: Determine whether parentheses should be special or should be considered individual operators.
# TODO: Investigate using built-in regex magic parser:
#         http://code.activestate.com/recipes/457664-hidden-scanner-functionality-in-re-module/

# Regular expression patterns.
INDENT_SIZE = 4
RE_LEADING_INDENT = re.compile(fr'^((?: {{{INDENT_SIZE}}})*)(.*)$')

RE_COMMA = re.compile(r',')
RE_NUMBER = re.compile(r'(?:\d+)'                           # 42
                       r'|'
                       r'(?:\.\d+(?:[eE][+-]?\d+)?)'        # .42 | .42e-8
                       r'|'
                       r'(?:\d+[eE][+-]?\d+)'               # 42e3
                       r'|'
                       r'(?:\d+\.\d*(?:[eE][+-]?\d+)?)')    # 42.7e2 | 42.e9 | 42. | 42.3e-8
RE_NAME = re.compile(r'_+|(?:_*[a-z][_a-zA-Z0-9]*(?:-[_a-zA-Z0-9]+)*[!@$%^&*?]?)')
RE_CLASS = re.compile(r'[A-Z][_a-zA-Z0-9]*(?:-[_a-zA-Z0-9]+)*')
RE_OPERATOR = re.compile(r'[!@$%^&*()\-=+|:/?<>\[\]{}~.]+')

def make_infix_re(pattern: PatternType, group_name: str) -> PatternType:
    return re.compile(r'(?P<left_val>.*)(?P<' + group_name + '>' + pattern + r')(?P<right_val>.*)')

RE_INFIX_COMMA = make_infix_re(RE_COMMA.pattern, 'comma')
RE_INFIX_OP = make_infix_re(RE_OPERATOR.pattern, 'op')

# Regular expression magic class for making if/else matching simpler. Idea from:
#   http://code.activestate.com/recipes/456151-using-rematch-research-and-regroup-in-if-elif-elif/
class RegexMatcher:
    def __init__(self, token):
        self._token = token
        self._match = None

    def fullmatch(self, pattern: PatternType):
        self._match = pattern.fullmatch(self._token)
        return self._match

    def group(self, grp: Union[int, str]):
        return self._match.group(grp)

# Tokens
class Lexeme:
    def __init__(self, text, repl_with_text=True):
        self.text = text
        self._repl_with_text = repl_with_text

    def __repr__(self):
        if self._repl_with_text:
            return f'{type(self).__name__}({self.text})'
        else:
            return type(self).__name__

    def __str__(self):
        return self.text

    def __eq__(self, other):
        if not isinstance(other, type(self)):
            return False
        return self.text == other.text

    def __ne__(self, other):
        return not self.__eq__(other)

class Indent(Lexeme):
    def __init__(self):
        super().__init__(' ' * INDENT_SIZE, False)

class NewLine(Lexeme):
    def __init__(self):
        super().__init__('\n', False)

class Period(Lexeme):
    def __init__(self):
        super().__init__('.', False)

class Comma(Lexeme):
    def __init__(self):
        super().__init__(',', False)

class OpenParen(Lexeme):
    def __init__(self):
        super().__init__('(', False)

class CloseParen(Lexeme):
    def __init__(self):
        super().__init__(')', False)

class Colon(Lexeme):
    def __init__(self):
        super().__init__(':', False)

class Arrow(Lexeme):
    def __init__(self):
        super().__init__('->', False)

class Number(Lexeme):
    pass

class Name(Lexeme):
    pass

class Class(Lexeme):
    pass

class Operator(Lexeme):
    pass

INDENT = Indent()
NEWLINE = NewLine()
PERIOD = Period()
COMMA = Comma()
OPEN_PAREN = OpenParen()
CLOSE_PAREN = CloseParen()
COLON = Colon()
ARROW = Arrow()

# Lexer errors.
class LexerError(Exception):
    pass

# Lexer implementation.
class Lexer:
    @classmethod
    def lex(cls, text: str) -> List[Lexeme]:
        lexemes = []
        for line in text.splitlines():
            lexemes.extend(cls.lex_line(line))
            lexemes.append(NEWLINE)
        return lexemes

    @classmethod
    def lex_file(cls, file: str) -> List[Lexeme]:
        lexemes = []
        with open(file) as f:
            for line in f:
                lexemes.extend(cls.lex_line(line))
                lexemes.append(NEWLINE)
        return lexemes

    @classmethod
    def lex_line(cls, line: str) -> List[Lexeme]:
        lexemes = []
        if not line:
            return lexemes
        # Remove properly-indented leading whitespace.
        match = RE_LEADING_INDENT.match(line)
        if match is None:  # pragma: no cover
            raise LexerError(f"invalid line given: '{line}'")
        indentation, rest = match.groups()
        for indentation_level in range(len(indentation) // INDENT_SIZE):
            lexemes.append(Indent())
        for token in rest.split():
            for lexeme in cls.lex_token(token):
                lexemes.append(lexeme)
        return lexemes

    @classmethod
    def lex_token(cls, token: str) -> List[Lexeme]:
        matcher = RegexMatcher(token)
        lexemes = []
        if not token:
            return lexemes
        elif matcher.fullmatch(RE_COMMA):
            lexemes.append(COMMA)
        elif matcher.fullmatch(RE_NUMBER):
            lexemes.append(Number(matcher.group(0)))
        elif matcher.fullmatch(RE_NAME):
            lexemes.append(Name(matcher.group(0)))
        elif matcher.fullmatch(RE_CLASS):
            lexemes.append(Class(matcher.group(0)))
        elif matcher.fullmatch(RE_OPERATOR):
            symbol = matcher.group(0)
            if symbol == PERIOD.text:
                lexemes.append(PERIOD)
            elif symbol == OPEN_PAREN.text:
                lexemes.append(OPEN_PAREN)
            elif symbol == CLOSE_PAREN.text:
                lexemes.append(CLOSE_PAREN)
            elif symbol == COLON.text:
                lexemes.append(COLON)
            elif symbol == ARROW.text:
                lexemes.append(ARROW)
            else:
                lexemes.append(Operator(symbol))
        elif matcher.fullmatch(RE_INFIX_COMMA):
            lexemes.extend(cls.lex_token(matcher.group('left_val')))
            lexemes.append(COMMA)
            lexemes.extend(cls.lex_token(matcher.group('right_val')))
        elif matcher.fullmatch(RE_INFIX_OP):
            lexemes.extend(cls.lex_token(matcher.group('left_val')))
            lexemes.extend(cls.lex_token(matcher.group('op')))
            lexemes.extend(cls.lex_token(matcher.group('right_val')))
        else:
            raise LexerError(f"invalid token: '{token}'")
        return lexemes

lex_file = Lexer.lex_file
lex_line = Lexer.lex_line
<commit_msg>Add Dedent to lexer<commit_after>import re

from re import _pattern_type as PatternType
from typing import List, Union

__all__ = [
    'Lexeme', 'GenericLexeme',
    'Indent', 'Dedent', 'NewLine', 'Period', 'Comma', 'OpenParen', 'CloseParen', 'Colon', 'Arrow',
    'Number', 'Name', 'Class', 'Operator',
    'INDENT', 'DEDENT', 'NEWLINE', 'PERIOD', 'COMMA', 'OPEN_PAREN', 'CLOSE_PAREN', 'COLON', 'ARROW',
    'LexerError', 'Lexer',
    'lex_file', 'lex_line',
]

# FIXME: Ambiguity between names with symbol endings and operators with those same symbols.
# FIXME: Ambiguity between function call and postfix Operator `()`
# TODO: Handle strings.
# TODO: Determine whether parentheses should be special or should be considered individual operators.
# TODO: Investigate using built-in regex magic parser:
#         http://code.activestate.com/recipes/457664-hidden-scanner-functionality-in-re-module/

# Regular expression patterns.
INDENT_SIZE = 4
RE_LEADING_INDENT = re.compile(fr'^((?: {{{INDENT_SIZE}}})*)(.*)$')

RE_COMMA = re.compile(r',')
RE_NUMBER = re.compile(r'(?:\d+)'                           # 42
                       r'|'
                       r'(?:\.\d+(?:[eE][+-]?\d+)?)'        # .42 | .42e-8
                       r'|'
                       r'(?:\d+[eE][+-]?\d+)'               # 42e3
                       r'|'
                       r'(?:\d+\.\d*(?:[eE][+-]?\d+)?)')    # 42.7e2 | 42.e9 | 42. | 42.3e-8
RE_NAME = re.compile(r'_+|(?:_*[a-z][_a-zA-Z0-9]*(?:-[_a-zA-Z0-9]+)*[!@$%^&*?]?)')
RE_CLASS = re.compile(r'[A-Z][_a-zA-Z0-9]*(?:-[_a-zA-Z0-9]+)*')
RE_OPERATOR = re.compile(r'[!@$%^&*()\-=+|:/?<>\[\]{}~.]+')

def make_infix_re(pattern: PatternType, group_name: str) -> PatternType:
    return re.compile(r'(?P<left_val>.*)(?P<' + group_name + '>' + pattern + r')(?P<right_val>.*)')

RE_INFIX_COMMA = make_infix_re(RE_COMMA.pattern, 'comma')
RE_INFIX_OP = make_infix_re(RE_OPERATOR.pattern, 'op')

# Regular expression magic class for making if/else matching simpler. Idea from:
#   http://code.activestate.com/recipes/456151-using-rematch-research-and-regroup-in-if-elif-elif/
class RegexMatcher:
    def __init__(self, token):
        self._token = token
        self._match = None

    def fullmatch(self, pattern: PatternType):
        self._match = pattern.fullmatch(self._token)
        return self._match

    def group(self, grp: Union[int, str]):
        return self._match.group(grp)

# Tokens
class Lexeme:
    def __init__(self, text, repl_with_text=True):
        self.text = text
        self._repl_with_text = repl_with_text

    def __repr__(self):
        if self._repl_with_text:
            return f'{type(self).__name__}({self.text})'
        else:
            return type(self).__name__

    def __str__(self):
        return self.text

    def __eq__(self, other):
        if not isinstance(other, type(self)):
            return False
        return self.text == other.text

    def __ne__(self, other):
        return not self.__eq__(other)

class Indent(Lexeme):
    def __init__(self):
        super().__init__(' ' * INDENT_SIZE, False)

class Dedent(Lexeme):
    def __init__(self):
        super().__init__('', False)

class NewLine(Lexeme):
    def __init__(self):
        super().__init__('\n', False)

class Period(Lexeme):
    def __init__(self):
        super().__init__('.', False)

class Comma(Lexeme):
    def __init__(self):
        super().__init__(',', False)

class OpenParen(Lexeme):
    def __init__(self):
        super().__init__('(', False)

class CloseParen(Lexeme):
    def __init__(self):
        super().__init__(')', False)

class Colon(Lexeme):
    def __init__(self):
        super().__init__(':', False)

class Arrow(Lexeme):
    def __init__(self):
        super().__init__('->', False)

class Number(Lexeme):
    pass

class Name(Lexeme):
    pass

class Class(Lexeme):
    pass

class Operator(Lexeme):
    pass

INDENT = Indent()
DEDENT = Dedent()
NEWLINE = NewLine()
PERIOD = Period()
COMMA = Comma()
OPEN_PAREN = OpenParen()
CLOSE_PAREN = CloseParen()
COLON = Colon()
ARROW = Arrow()

# Lexer errors.
class LexerError(Exception):
    pass

# Lexer implementation.
class Lexer:
    @classmethod
    def lex(cls, text: str) -> List[Lexeme]:
        lexemes = []
        for line in text.splitlines():
            lexemes.extend(cls.lex_line(line))
            lexemes.append(NEWLINE)
        return lexemes

    @classmethod
    def lex_file(cls, file: str) -> List[Lexeme]:
        lexemes = []
        with open(file) as f:
            for line in f:
                lexemes.extend(cls.lex_line(line))
                lexemes.append(NEWLINE)
        return lexemes

    @classmethod
    def lex_line(cls, line: str) -> List[Lexeme]:
        lexemes = []
        if not line:
            return lexemes
        # Remove properly-indented leading whitespace.
        match = RE_LEADING_INDENT.match(line)
        if match is None:  # pragma: no cover
            raise LexerError(f"invalid line given: '{line}'")
        indentation, rest = match.groups()
        for indentation_level in range(len(indentation) // INDENT_SIZE):
            lexemes.append(Indent())
        for token in rest.split():
            for lexeme in cls.lex_token(token):
                lexemes.append(lexeme)
        return lexemes

    @classmethod
    def lex_token(cls, token: str) -> List[Lexeme]:
        matcher = RegexMatcher(token)
        lexemes = []
        if not token:
            return lexemes
        elif matcher.fullmatch(RE_COMMA):
            lexemes.append(COMMA)
        elif matcher.fullmatch(RE_NUMBER):
            lexemes.append(Number(matcher.group(0)))
        elif matcher.fullmatch(RE_NAME):
            lexemes.append(Name(matcher.group(0)))
        elif matcher.fullmatch(RE_CLASS):
            lexemes.append(Class(matcher.group(0)))
        elif matcher.fullmatch(RE_OPERATOR):
            symbol = matcher.group(0)
            if symbol == PERIOD.text:
                lexemes.append(PERIOD)
            elif symbol == OPEN_PAREN.text:
                lexemes.append(OPEN_PAREN)
            elif symbol == CLOSE_PAREN.text:
                lexemes.append(CLOSE_PAREN)
            elif symbol == COLON.text:
                lexemes.append(COLON)
            elif symbol == ARROW.text:
                lexemes.append(ARROW)
            else:
                lexemes.append(Operator(symbol))
        elif matcher.fullmatch(RE_INFIX_COMMA):
            lexemes.extend(cls.lex_token(matcher.group('left_val')))
            lexemes.append(COMMA)
            lexemes.extend(cls.lex_token(matcher.group('right_val')))
        elif matcher.fullmatch(RE_INFIX_OP):
            lexemes.extend(cls.lex_token(matcher.group('left_val')))
            lexemes.extend(cls.lex_token(matcher.group('op')))
            lexemes.extend(cls.lex_token(matcher.group('right_val')))
        else:
            raise LexerError(f"invalid token: '{token}'")
        return lexemes

lex_file = Lexer.lex_file
lex_line = Lexer.lex_line
<|endoftext|>


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

import glob
import os
import re
import traceback
import types

import ample_util
import pdb_edit

class MaxclusterData(object):

    def __init__(self):
        self.length    = None
        self.pairs     = None
        self.rmsd      = None
        self.maxsub    = None
        self.tm        = None
        self.msi       = None
        self.modelName = None
        self.pdb       = None
        return

    def __str__(self):
        """List the data attributes of this object"""
        me = {}
        for slot in dir(self):
            attr = getattr(self, slot)
            if not slot.startswith("__") and not ( isinstance(attr, types.MethodType) or
              isinstance(attr, types.FunctionType) ):
                me[slot] = attr

        return "{0} : {1}".format(self.__repr__(),str(me))

class Maxcluster(object):
    """

    # Extract the first chain from the nativePdb

    # Create a residueSequenceMap and see if the residues match

    # If not use keep_matching to create a nativePdb that has the correct residue sequence`

    # Run Maxcluster to compare the models to the native
    """

    def __init__(self,maxcluster_exe):

        self.maxclusterExe = maxcluster_exe

        return

    def compareDirectory(self, nativePdbInfo=None,
                               resSeqMap=None,
                               modelsDirectory=None,
                               workdir=None ):

        self.data = []
        self.workdir = workdir
        if not self.workdir:
            self.workdir = os.getcwd()

        #refModel = os.path.join( modelsDirectory, "S_00000001.pdb" )
        nativePdb = self.prepareNative( nativePdbInfo=nativePdbInfo, resSeqMap=resSeqMap )

        logfile = os.path.join( self.workdir, "maxclusterD.log" )
        self.runCompareDirectory( nativePdb=nativePdb, modelsDirectory=modelsDirectory, logfile=logfile )
        self.parseLogDirectory( logfile=logfile )

        return

    def compareSingle(self, nativePdb=None, modelPdb=None, sequenceIndependant=True, rmsd=False, workdir=None ):

        self.workdir = workdir
        if not self.workdir:
            self.workdir = os.getcwd()

        cmd = [ self.maxclusterExe, "-e", nativePdb, "-p", modelPdb ]

        if sequenceIndependant:
            cmd.append( "-in" )

        if rmsd:
            cmd.append( "-rmsd" )

        logfile = ample_util.filename_append( filename=modelPdb,
                                    astr="maxcluster",
                                    directory=self.workdir )

        if rmsd:
            logfile = os.path.splitext( logfile )[0] + "_rmsd.log"
        else:
            logfile = os.path.splitext( logfile )[0] + ".log"
        self.maxclusterLogfile = logfile

        #print "running cmd "," ".join( cmd )
        retcode = ample_util.run_command( cmd, logfile=self.maxclusterLogfile, dolog=False )

        if retcode != 0:
            msg = "non-zero return code for maxcluster in runMaxcluster!"
            #logging.critical( msg )

        if rmsd:
            data = self.parseLogSingleRmsd()
        else:
            data = self.parseLogSingleTm()

        return data

    def prepareNative(self, nativePdbInfo=None, resSeqMap=None ):
        """do stuff"""

        # Find out how many chains and extract the first if > 1
        if len( nativePdbInfo.models ) > 1:
            raise RuntimeError,"More than 1 model."

        # Check if > 1 chain
        chainID=None
        if len( nativePdbInfo.models[0].chains ) > 1:

            chainID=nativePdbInfo.models[0].chains[0]

            # Assume native is standardised
            # Extract the chain if > 1
            nativePdbChain = ample_util.filename_append( filename=nativePdbInfo.pdb,
                                                         astr="chain{0}".format(chainID) )
            pdb_edit.extract_chain(nativePdbInfo.pdb, nativePdbChain, chainID)
            nativePdb = nativePdbChain
        else:
            nativePdb = nativePdbInfo.pdb

        if not resSeqMap.resSeqMatch():

            # We need to create a copy of the native with numbering matching the model
            nativeRenumber = ample_util.filename_append( filename=nativePdb,
                                                         astr="ren".format(chainID) )
            pdb_edit.match_resseq( targetPdb=nativePdb, outPdb=nativeRenumber, resMap=resSeqMap )
            nativePdb = nativeRenumber

        return nativePdb

    def parseLogDirectory(self, logfile=None ):

        self.data = []
        assert logfile

        #INFO  : 1000. 2XOV_clean_ren.pdb vs. /media/data/shared/TM/2XOV/models/S_00000444.pdb  Pairs=  36, RMSD= 3.065, MaxSub=0.148, TM=0.192, MSI=0.148
        for line in open( logfile, 'r' ):

            if re.match( "INFO *: .* vs\. .* Pairs=", line ):

                # Remove spaces after =
                line = re.sub("= +", "=", line )
                # Now remove commas
                line = line.replace(",","")

                fields = line.split()

                d = MaxclusterData()
                d.pdb = fields[5]
                tmp = os.path.splitext( os.path.basename( fields[5] ) )[0]
#                 # Hack to make sure there isn't something like "1_" prepended to the name
#                 if not tmp.startswith("S"):
#                     for i,f in enumerate(tmp):
#                         if f=="S":
#                             tmp=tmp[i:]
#                             break
                d.modelName = tmp

                label, value = fields[6].split( "=" )
                assert label == "Pairs"
                d.pairs = int( value )

                label, value = fields[7].split( "=" )
                assert label == "RMSD"
                d.rmsd = float( value )

                label, value = fields[8].split( "=" )
                assert label == "MaxSub"
                d.maxsub = float( value )

                label, value = fields[9].split( "=" )
                assert label == "TM"
                d.tm = float( value )

                label, value = fields[10].split( "=" )
                assert label == "MSI"
                d.msi = float( value )

                self.data.append( d )

        return

    def parseLogSingleTm(self, logfile=None):

        if logfile is None:
            logfile = self.maxclusterLogfile

        assert logfile

        d = MaxclusterData()
        for line in open( logfile, 'r' ):

            line = line.strip()
            #"Iter 1: Pairs=  14, RMSD= 0.155, MAXSUB=0.855. Len=  15. gRMSD= 0.673, TM=0.858
            if re.match( "Iter \d?: ?Pairs=", line ):

                # Remove spaces after =
                line = re.sub("= +", "=", line )

                # Remove trailing . after numbers
                line = re.sub("\. +", " ", line )

                # Now remove commas
                line = line.replace(",","")

                fields = line.split()

                label, value = fields[2].split( "=" )
                assert label == "Pairs"
                d.pairs = int( value )

                label, value = fields[3].split( "=" )
                assert label == "RMSD"
                d.rmsd = float( value )

                label, value = fields[4].split( "=" )
                assert label == "MAXSUB"
                d.maxsub = float( value )

                label, value = fields[5].split( "=" )
                assert label == "Len"
                d.length = float( value )

                # skip gRMSD

                label, value = fields[7].split( "=" )
                assert label == "TM"
                d.tm = float( value )

        return d

    def parseLogSingleRmsd(self, logfile=None):

        if logfile is None:
            logfile = self.maxclusterLogfile

        assert logfile

        d = MaxclusterData()

        #INFO  : 1000. 2XOV_clean_ren.pdb vs. /media/data/shared/TM/2XOV/models/S_00000444.pdb  Pairs=  36, RMSD= 3.065, MaxSub=0.148, TM=0.192, MSI=0.148
        for line in open( logfile, 'r' ):

            line = line.strip()

            if line.startswith("RMSD="):
                # RMSD= 0.132 (Pairs=   8, rRMSD=0.034 ( -3.11)), URMSD= 0.049 (rURMSD=0.049)

                # Remove spaces after =
                line = re.sub("= +", "=", line )

                # Remove trailing . after numbers
                line = re.sub("\. +", " ", line )

                # Now remove commas and brackets
                line = line.replace(",","")
                line = line.replace("("," ")

                fields = line.split()

                label, value = fields[0].split( "=" )
                assert label == "RMSD"
                d.rmsd = float( value )

                label, value = fields[1].split( "=" )
                assert label == "Pairs"
                d.pairs = int( value )

        return d

    def tm(self, modelName ):
        """"""
        for d in self.data:
            if d.modelName == modelName:
                return d.tm
        s = "\n".join(traceback.format_list(traceback.extract_stack()))
        print "MaxCluster tm failed to find model name: {0}\n{1}".format(modelName,s)
        return None

    def rmsd(self, modelName ):
        """"""
        for d in self.data:
            if d.modelName == modelName:
                return d.rmsd
        s = "\n".join(traceback.format_list(traceback.extract_stack()))
        print "MaxCluster rmsd failed to find model name: {0}\n{1}".format(modelName,s)
        return None

    def maxsubSorted(self, reverse=True ):
        return sorted( self.data, key=lambda data: data.maxsub, reverse=reverse )

    def runCompareDirectory(self, nativePdb=None, modelsDirectory=None, logfile=None ):

        # Generate the list of models
        pdblist = os.path.join( self.workdir, "models.list")
        with open( pdblist, 'w' ) as f:
            l = glob.glob( os.path.join( modelsDirectory, '*S_*.pdb' ) )
            if not len(l) > 0:
                raise RuntimeError,"Could not find any pdb files in directory: {0}".format(modelsDirectory)
            f.write( os.linesep.join( l ) )

        # Run Maxcluster
        cmd = [ self.maxclusterExe, "-e", nativePdb, "-l", pdblist, ]
        retcode = ample_util.run_command( cmd, logfile=logfile, dolog=True )

        if retcode != 0:
            msg = "non-zero return code for maxcluster in runMaxcluster!"
            #logging.critical( msg )
            raise RuntimeError, msg

        return

    def tmSorted(self, reverse=True ):
        return sorted( self.data, key=lambda data: data.tm, reverse=reverse )

<commit_msg>Change to use model path and a dictionary<commit_after>#!/usr/bin/env python

import glob
import os
import re
import traceback

import ample_util
import pdb_edit

class Maxcluster(object):
    """

    # Extract the first chain from the nativePdb

    # Create a residueSequenceMap and see if the residues match

    # If not use keep_matching to create a nativePdb that has the correct residue sequence`

    # Run Maxcluster to compare the models to the native
    """

    def __init__(self,maxcluster_exe):

        self.maxclusterExe = maxcluster_exe

        return

    def compareDirectory(self, nativePdbInfo=None,
                               resSeqMap=None,
                               modelsDirectory=None,
                               workdir=None ):

        self.data = []
        self.workdir = workdir
        if not self.workdir:
            self.workdir = os.getcwd()

        #refModel = os.path.join( modelsDirectory, "S_00000001.pdb" )
        nativePdb = self.prepareNative( nativePdbInfo=nativePdbInfo, resSeqMap=resSeqMap )

        logfile = os.path.join( self.workdir, "maxclusterD.log" )
        self.runCompareDirectory( nativePdb=nativePdb, modelsDirectory=modelsDirectory, logfile=logfile )
        self.parseLogDirectory( logfile=logfile )

        return

    def compareSingle(self, nativePdb=None, modelPdb=None, sequenceIndependant=True, rmsd=False, workdir=None ):

        self.workdir = workdir
        if not self.workdir:
            self.workdir = os.getcwd()

        cmd = [ self.maxclusterExe, "-e", nativePdb, "-p", modelPdb ]

        if sequenceIndependant:
            cmd.append( "-in" )

        if rmsd:
            cmd.append( "-rmsd" )

        logfile = ample_util.filename_append( filename=modelPdb,
                                    astr="maxcluster",
                                    directory=self.workdir )

        if rmsd:
            logfile = os.path.splitext( logfile )[0] + "_rmsd.log"
        else:
            logfile = os.path.splitext( logfile )[0] + ".log"
        self.maxclusterLogfile = logfile

        #print "running cmd "," ".join( cmd )
        retcode = ample_util.run_command( cmd, logfile=self.maxclusterLogfile, dolog=False )

        if retcode != 0:
            msg = "non-zero return code for maxcluster in runMaxcluster!"
            #logging.critical( msg )
            print msg

        if rmsd:
            data = self.parseLogSingleRmsd()
        else:
            data = self.parseLogSingleTm()

        return data

    def prepareNative(self, nativePdbInfo=None, resSeqMap=None ):
        """do stuff"""

        # Find out how many chains and extract the first if > 1
        if len( nativePdbInfo.models ) > 1:
            raise RuntimeError,"More than 1 model."

        # Check if > 1 chain
        chainID=None
        if len( nativePdbInfo.models[0].chains ) > 1:

            chainID=nativePdbInfo.models[0].chains[0]

            # Assume native is standardised
            # Extract the chain if > 1
            nativePdbChain = ample_util.filename_append( filename=nativePdbInfo.pdb,
                                                         astr="chain{0}".format(chainID) )
            pdb_edit.extract_chain(nativePdbInfo.pdb, nativePdbChain, chainID)
            nativePdb = nativePdbChain
        else:
            nativePdb = nativePdbInfo.pdb

        if not resSeqMap.resSeqMatch():

            # We need to create a copy of the native with numbering matching the model
            nativeRenumber = ample_util.filename_append( filename=nativePdb,
                                                         astr="ren".format(chainID) )
            pdb_edit.match_resseq( targetPdb=nativePdb, outPdb=nativeRenumber, resMap=resSeqMap )
            nativePdb = nativeRenumber

        return nativePdb

    def parseLogDirectory(self, logfile=None ):

        self.data = []
        assert logfile

        #INFO  : 1000. 2XOV_clean_ren.pdb vs. /media/data/shared/TM/2XOV/models/S_00000444.pdb  Pairs=  36, RMSD= 3.065, MaxSub=0.148, TM=0.192, MSI=0.148
        for line in open( logfile, 'r' ):

            if re.match( "INFO *: .* vs\. .* Pairs=", line ):

                # Remove spaces after =
                line = re.sub("= +", "=", line )
                # Now remove commas
                line = line.replace(",","")

                fields = line.split()

                d = {}
                d['pdb'] = fields[5]
                tmp = os.path.splitext( os.path.basename( fields[5] ) )[0]
#                 # Hack to make sure there isn't something like "1_" prepended to the name
#                 if not tmp.startswith("S"):
#                     for i,f in enumerate(tmp):
#                         if f=="S":
#                             tmp=tmp[i:]
#                             break
                d['model_name'] = tmp

                label, value = fields[6].split( "=" )
                assert label == "Pairs"
                d['pairs'] = int( value )

                label, value = fields[7].split( "=" )
                assert label == "RMSD"
                d['rmsd'] = float( value )

                label, value = fields[8].split( "=" )
                assert label == "MaxSub"
                d['maxsub'] = float( value )

                label, value = fields[9].split( "=" )
                assert label == "TM"
                d['tm'] = float( value )

                label, value = fields[10].split( "=" )
                assert label == "MSI"
                d['msi'] = float( value )

                self.data.append( d )

        return

    def parseLogSingleTm(self, logfile=None):

        if logfile is None:
            logfile = self.maxclusterLogfile

        assert logfile

        d = {}
        for line in open( logfile, 'r' ):

            line = line.strip()
            #"Iter 1: Pairs=  14, RMSD= 0.155, MAXSUB=0.855. Len=  15. gRMSD= 0.673, TM=0.858
            if re.match( "Iter \d?: ?Pairs=", line ):

                # Remove spaces after =
                line = re.sub("= +", "=", line )

                # Remove trailing . after numbers
                line = re.sub("\. +", " ", line )

                # Now remove commas
                line = line.replace(",","")

                fields = line.split()

                label, value = fields[2].split( "=" )
                assert label == "Pairs"
                d['pairs'] = int( value )

                label, value = fields[3].split( "=" )
                assert label == "RMSD"
                d['rmsd'] = float( value )

                label, value = fields[4].split( "=" )
                assert label == "MAXSUB"
                d['maxsub'] = float( value )

                label, value = fields[5].split( "=" )
                assert label == "Len"
                d['length'] = float( value )

                # skip gRMSD

                label, value = fields[7].split( "=" )
                assert label == "TM"
                d['tm'] = float( value )

        return d

    def parseLogSingleRmsd(self, logfile=None):

        if logfile is None:
            logfile = self.maxclusterLogfile

        assert logfile

        d = {}

        #INFO  : 1000. 2XOV_clean_ren.pdb vs. /media/data/shared/TM/2XOV/models/S_00000444.pdb  Pairs=  36, RMSD= 3.065, MaxSub=0.148, TM=0.192, MSI=0.148
        for line in open( logfile, 'r' ):

            line = line.strip()

            if line.startswith("RMSD="):
                # RMSD= 0.132 (Pairs=   8, rRMSD=0.034 ( -3.11)), URMSD= 0.049 (rURMSD=0.049)

                # Remove spaces after =
                line = re.sub("= +", "=", line )

                # Remove trailing . after numbers
                line = re.sub("\. +", " ", line )

                # Now remove commas and brackets
                line = line.replace(",","")
                line = line.replace("("," ")

                fields = line.split()

                label, value = fields[0].split( "=" )
                assert label == "RMSD"
                d['rmsd'] = float( value )

                label, value = fields[1].split( "=" )
                assert label == "Pairs"
                d['pairs'] = int( value )

        return d

    def tm(self,model):
        """"""
        for d in self.data:
            if d['pdb'] == model: return d['tm']
        #s = "\n".join(traceback.format_list(traceback.extract_stack()))
        print "MaxCluster tm failed to find model name: {0}\n{1}".format(model,self.data)
        return None

    def rmsd(self,model):
        """"""
        for d in self.data:
            if d['pdb'] == model: return d['rmsd']
        #s = "\n".join(traceback.format_list(traceback.extract_stack()))
        print "MaxCluster rmsd failed to find model name: {0}\n{1}".format(model,self.data)
        return None

    def maxsubSorted(self, reverse=True ):
        return sorted( self.data, key=lambda data: data['maxsub'], reverse=reverse )

    def runCompareDirectory(self, nativePdb=None, modelsDirectory=None, logfile=None ):

        # Generate the list of models
        pdblist = os.path.join( self.workdir, "models.list")
        with open( pdblist, 'w' ) as f:
            l = glob.glob( os.path.join( modelsDirectory, '*S_*.pdb' ) )
            if not len(l) > 0:
                raise RuntimeError,"Could not find any pdb files in directory: {0}".format(modelsDirectory)
            f.write( os.linesep.join( l ) )

        # Run Maxcluster
        cmd = [ self.maxclusterExe, "-e", nativePdb, "-l", pdblist, ]
        retcode = ample_util.run_command( cmd, logfile=logfile, dolog=True )

        if retcode != 0:
            msg = "non-zero return code for maxcluster in runMaxcluster!"
            #logging.critical( msg )
            raise RuntimeError, msg

        return

    def tmSorted(self, reverse=True ):
        return sorted( self.data, key=lambda data: data['tm'], reverse=reverse )

<|endoftext|>


<commit_before>
from lock_impl import EphemeralLock

from lock_impl import ConsulLockException
from lock_impl import LockAcquisitionException
<commit_msg>making defaults available through the main module<commit_after>
from lock_impl import EphemeralLock

from lock_impl import ConsulLockException
from lock_impl import LockAcquisitionException

import defaults

<|endoftext|>