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

from circuits.web import Server
from circuits.web.controllers import BaseController

class Root(BaseController):

    def index(self):
        return "Hello World!"

(Server(8000) + Root()).run()
<commit_msg>Reviewed and Fixed BaseController Demo for circuits.web [#45735111]<commit_after>#!/usr/bin/env python

from circuits.web import Server
from circuits.web.controllers import expose, BaseController

class Root(BaseController):

    @expose("index")
    def index(self):
        """Index Request Handler

        BaseController(s) do not expose methods as request handlers.
        Request Handlers _must_ be exposed explicitly using the ``@expose``
        decorator.
        """

        return "Hello World!"

(Server(9000) + Root()).run()
<|endoftext|>


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

from unittest import TestCase
from exfi.read_fasta import read_fasta

<commit_msg>Bugfix<commit_after>#!/usr/bin/env python3

from unittest import TestCase
from Bio import SeqIO
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
convnet-est-loss
"""
import numpy as np
from os import sys, path
import argparse
caffe_dir = '/home/aaskov/caffe/'

def damage_range(x):
    if float(x) < 0.0:
        raise argparse.ArgumentTypeError("%r is not positive" % x)
    return x

def run():
    parser = argparse.ArgumentParser(
        description='Neural network damage experiment. ',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('--proto', required=True, type=str,
                        help='Network prototxt.')
    parser.add_argument('--model', required=True, type=str,
                        help='Network caffemodel.')
    parser.add_argument('--meanfile', required=True, type=str,
                        help='Data mean file.')
    parser.add_argument('--data', required=True, type=str,
                        help='Data.')
    parser.add_argument('--layer', required=True, type=str,
                        help='Layer to apply damage to.')
    parser.add_argument('--prefix', required=True, type=str,
                        help='Uniqe model name.')
    parser.add_argument('--damage', default=1.0, type=damage_range,
                        help='Applied damage range.')
    parser.add_argument('--step-num', default=10, type=int,
                        help='Number of steps in damage interval.')
    parser.add_argument('--iterations', default=100, type=int,
                        help='Number of iterations to run.')
    parser.add_argument('--repeat', default=0, type=int,
                        help='Number of repeated experiments to run.')
    args = parser.parse_args()

    # Object file
    obj = 'experiment_damage_' + str(args.layer) + '_' + str(args.damage)
    obj += '_step_' + str(args.step_num) + '_iter_' + str(args.iterations)
    obj += '_prefix_' + str(args.prefix)

    loss_list = list()
    for std in np.linspace(0.0, args.damage, args.step_num):
        _loss = list()
        for t in range(args.repeat + 1):
            # Fetch a network structure and apply damage
            net = Network(args.proto, args.model, args.meanfile, args.data)
            net.__add_damage__(args.layer, std)

            # Forward to get loss
            top_1, top_5, acc, loss = net.forward(maxrun=args.iterations)
            _loss.append(loss)

        loss_list.append(_loss)

    # Store result
    save_obj(np.array(loss_list), obj)

if __name__ == '__main__' and __package__ is None:
    # Append parrent directory to sys path
    sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))
    from network import Network
    from input_output import save_obj

    # Setup Caffe
    sys.path.insert(0, caffe_dir + 'python')
    import caffe
    caffe.set_mode_cpu()

    # Run experiment
    run()

<commit_msg>Various updates<commit_after># -*- coding: utf-8 -*-
"""
convnet-est-loss
"""
import numpy as np
from os import sys, path
import argparse
caffe_dir = '/home/aaskov/caffe/'

def damage_range(x):
    if float(x) < 0.0:
        raise argparse.ArgumentTypeError("%r is not positive" % x)
    return float(x)

def run():
    parser = argparse.ArgumentParser(
        description='Neural network damage experiment. ',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument('--proto', required=True, type=str,
                        help='Network prototxt.')
    parser.add_argument('--model', required=True, type=str,
                        help='Network caffemodel.')
    parser.add_argument('--meanfile', required=True, type=str,
                        help='Data mean file.')
    parser.add_argument('--data', required=True, type=str,
                        help='Data.')
    parser.add_argument('--layer', required=True, type=str,
                        help='Layer to apply damage to.')
    parser.add_argument('--prefix', required=True, type=str,
                        help='Uniqe model name.')
    parser.add_argument('--damage', default=1.0, type=damage_range,
                        help='Applied damage range.')
    parser.add_argument('--step-num', default=10, type=int,
                        help='Number of steps in damage interval.')
    parser.add_argument('--iterations', default=100, type=int,
                        help='Number of iterations to run.')
    parser.add_argument('--repeat', default=0, type=int,
                        help='Number of repeated experiments to run.')
    args = parser.parse_args()

    # Object file
    obj = 'experiment_damage_' + str(args.layer) + '_' + str(args.damage)
    obj += '_step_' + str(args.step_num) + '_iter_' + str(args.iterations)
    obj += '_prefix_' + str(args.prefix)

    loss_list = list()
    for std in np.linspace(0.0, args.damage, args.step_num):
        _loss = list()
        for t in range(args.repeat + 1):
            # Fetch a network structure and apply damage
            net = Network(args.proto, args.model, args.meanfile, args.data)
            net.__add_damage__(args.layer, std)

            # Forward to get loss
            top_1, top_5, acc, loss = net.forward(maxrun=args.iterations)
            _loss.append(loss)

        loss_list.append(_loss)

    # Store result
    save_obj(np.array(loss_list), obj)

if __name__ == '__main__' and __package__ is None:
    # Append parrent directory to sys path
    sys.path.append(path.dirname(path.dirname(path.abspath(__file__))))
    from network import Network
    from input_output import save_obj

    # Setup Caffe
    sys.path.insert(0, caffe_dir + 'python')
    import caffe
    caffe.set_mode_cpu()

    # Run experiment
    run()

<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
##############################################################################
#
#    Copyright (C) 2014 Compassion CH (http://www.compassion.ch)
#    Releasing children from poverty in Jesus' name
#    @author: Cyril Sester, Emanuel Cino
#
#    The licence is in the file __openerp__.py
#
##############################################################################

from openerp.osv import orm, fields

class res_partner(orm.Model):
    _inherit = 'res.partner'

    def _get_related_contracts(self, cr, uid, ids, field_names, arg, context):
        """ Returns the contracts of the sponsor of given type
        ('fully_managed', 'correspondant' or 'payer')
        """
        res = dict()
        field_res = dict()
        contract_obj = self.pool.get('recurring.contract')
        for id in ids:
            correspondant_ids = contract_obj.search(
                cr, uid, [('correspondant_id', '=', id),
                          ('type', '=', 'S'),
                          ('fully_managed', '=', False)],
                order='start_date desc', context={})
            paid_ids = contract_obj.search(
                cr, uid, [('partner_id', '=', id),
                          ('type', '=', 'S'),
                          ('fully_managed', '=', False)],
                order='start_date desc', context={})
            fully_managed_ids = contract_obj.search(
                cr, uid, [('partner_id', '=', id),
                          ('type', '=', 'S'),
                          ('fully_managed', '=', True)],
                order='start_date desc', context={})

            field_res['contracts_fully_managed'] = fully_managed_ids
            field_res['contracts_paid'] = paid_ids
            field_res['contracts_correspondant'] = correspondant_ids
            res[id] = field_res.copy()

        return res

    def _write_related_contracts(self, cr, uid, id, name, value, inv_arg,
                                 context=None):
        value_obj = self.pool.get('recurring.contract')
        for line in value:
            if line[0] == 1:  # one2many update
                value_id = line[1]
                value_obj.write(cr, uid, [value_id], line[2])
        return True

    _columns = {
        'contracts_fully_managed': fields.function(
            _get_related_contracts, type="one2many",
            obj="recurring.contract",
            fnct_inv=_write_related_contracts, multi='sponsorships',
            string='Fully managed sponsorships',
            order="state asc",),
        'contracts_paid': fields.function(
            _get_related_contracts, type="one2many",
            obj="recurring.contract",
            fnct_inv=_write_related_contracts, multi='sponsorships',
            string='Sponsorships as payer only'),
        'contracts_correspondant': fields.function(
            _get_related_contracts, type="one2many",
            obj="recurring.contract",
            fnct_inv=_write_related_contracts, multi='sponsorships',
            string='Sponsorships as correspondant only'),
    }

    def show_lines(self, cr, uid, ids, context=None):
        try:
            ir_model_data = self.pool.get('ir.model.data')
            view_id = ir_model_data.get_object_reference(
                cr, uid, 'sponsorship_compassion',
                'view_invoice_line_partner_tree')[1]
        except ValueError:
            view_id = False
        context['search_default_partner_id'] = ids
        action = {
            'name': _('Related invoice lines'),
            'type': 'ir.actions.act_window',
            'view_type': 'form',
            'view_mode': 'tree, form',
            'res_model': 'account.invoice.line',
            'view_id': view_id,
            'views': [(view_id, 'tree'), (False, 'form')],
            'target': 'current',
            'context': context,
        }

        return action

    def show_move_lines(self, cr, uid, ids, context=None):
        try:
            ir_model_data = self.pool.get('ir.model.data')
            move_line_id = ir_model_data.get_object_reference(
                cr, uid, 'account',
                'view_move_line_tree')[1]
        except ValueError:
            move_line_id = False
        context['search_default_partner_id'] = ids
        action = {
            'name': 'Related invoice lines',
            'type': 'ir.actions.act_window',
            'view_type': 'form',
            'view_mode': 'tree',
            'res_model': 'account.move.line',
            'view_id': move_line_id,
            'views': [(move_line_id, 'tree')],
            'target': 'current',
            'context': context,
        }

        return action

    def create_contract(self, cr, uid, ids, context=None):
        context.update({
            'default_partner_id': ids[0],
            'default_type': 'S',
        })
        return {
            'type': 'ir.actions.act_window',
            'name': 'New Sponsorship',
            'view_type': 'form',
            'view_mode': 'form',
            'res_model': 'recurring.contract',
            'target': 'current',
            'context': context
        }

    def unreconciled_transaction_items(self, cr, uid, ids, context=None):
        context['search_default_unreconciled'] = 1
        return self.show_move_lines(cr, uid, ids, context)

    def receivable_transaction_items(self, cr, uid, ids, context=None):
        account_ids = self.pool.get('account.account').search(
            cr, uid,
            [('code', '=', '1050')],
            context=context)
        context['search_default_account_id'] = account_ids[0]
        return self.show_move_lines(cr, uid, ids, context)
<commit_msg>Fix and update translations in class<commit_after># -*- encoding: utf-8 -*-
##############################################################################
#
#    Copyright (C) 2014 Compassion CH (http://www.compassion.ch)
#    Releasing children from poverty in Jesus' name
#    @author: Cyril Sester, Emanuel Cino
#
#    The licence is in the file __openerp__.py
#
##############################################################################

from openerp.osv import orm, fields
from openerp.tools.translate import _

class res_partner(orm.Model):
    _inherit = 'res.partner'

    def _get_related_contracts(self, cr, uid, ids, field_names, arg, context):
        """ Returns the contracts of the sponsor of given type
        ('fully_managed', 'correspondant' or 'payer')
        """
        res = dict()
        field_res = dict()
        contract_obj = self.pool.get('recurring.contract')
        for id in ids:
            correspondant_ids = contract_obj.search(
                cr, uid, [('correspondant_id', '=', id),
                          ('type', '=', 'S'),
                          ('fully_managed', '=', False)],
                order='start_date desc', context={})
            paid_ids = contract_obj.search(
                cr, uid, [('partner_id', '=', id),
                          ('type', '=', 'S'),
                          ('fully_managed', '=', False)],
                order='start_date desc', context={})
            fully_managed_ids = contract_obj.search(
                cr, uid, [('partner_id', '=', id),
                          ('type', '=', 'S'),
                          ('fully_managed', '=', True)],
                order='start_date desc', context={})

            field_res['contracts_fully_managed'] = fully_managed_ids
            field_res['contracts_paid'] = paid_ids
            field_res['contracts_correspondant'] = correspondant_ids
            res[id] = field_res.copy()

        return res

    def _write_related_contracts(self, cr, uid, id, name, value, inv_arg,
                                 context=None):
        value_obj = self.pool.get('recurring.contract')
        for line in value:
            if line[0] == 1:  # one2many update
                value_id = line[1]
                value_obj.write(cr, uid, [value_id], line[2])
        return True

    _columns = {
        'contracts_fully_managed': fields.function(
            _get_related_contracts, type="one2many",
            obj="recurring.contract",
            fnct_inv=_write_related_contracts, multi='sponsorships',
            string='Fully managed sponsorships',
            order="state asc",),
        'contracts_paid': fields.function(
            _get_related_contracts, type="one2many",
            obj="recurring.contract",
            fnct_inv=_write_related_contracts, multi='sponsorships',
            string='Sponsorships as payer only'),
        'contracts_correspondant': fields.function(
            _get_related_contracts, type="one2many",
            obj="recurring.contract",
            fnct_inv=_write_related_contracts, multi='sponsorships',
            string='Sponsorships as correspondant only'),
    }

    def show_lines(self, cr, uid, ids, context=None):
        try:
            ir_model_data = self.pool.get('ir.model.data')
            view_id = ir_model_data.get_object_reference(
                cr, uid, 'sponsorship_compassion',
                'view_invoice_line_partner_tree')[1]
        except ValueError:
            view_id = False
        context['search_default_partner_id'] = ids
        action = {
            'name': _('Related invoice lines'),
            'type': 'ir.actions.act_window',
            'view_type': 'form',
            'view_mode': 'tree, form',
            'res_model': 'account.invoice.line',
            'view_id': view_id,
            'views': [(view_id, 'tree'), (False, 'form')],
            'target': 'current',
            'context': context,
        }

        return action

    def show_move_lines(self, cr, uid, ids, context=None):
        try:
            ir_model_data = self.pool.get('ir.model.data')
            move_line_id = ir_model_data.get_object_reference(
                cr, uid, 'account',
                'view_move_line_tree')[1]
        except ValueError:
            move_line_id = False
        context['search_default_partner_id'] = ids
        action = {
            'name': _('Related invoice lines'),
            'type': 'ir.actions.act_window',
            'view_type': 'form',
            'view_mode': 'tree',
            'res_model': 'account.move.line',
            'view_id': move_line_id,
            'views': [(move_line_id, 'tree')],
            'target': 'current',
            'context': context,
        }

        return action

    def create_contract(self, cr, uid, ids, context=None):
        context.update({
            'default_partner_id': ids[0],
            'default_type': 'S',
        })
        return {
            'type': 'ir.actions.act_window',
            'name': _('New Sponsorship'),
            'view_type': 'form',
            'view_mode': 'form',
            'res_model': 'recurring.contract',
            'target': 'current',
            'context': context
        }

    def unreconciled_transaction_items(self, cr, uid, ids, context=None):
        context['search_default_unreconciled'] = 1
        return self.show_move_lines(cr, uid, ids, context)

    def receivable_transaction_items(self, cr, uid, ids, context=None):
        account_ids = self.pool.get('account.account').search(
            cr, uid,
            [('code', '=', '1050')],
            context=context)
        context['search_default_account_id'] = account_ids[0]
        return self.show_move_lines(cr, uid, ids, context)
<|endoftext|>


<commit_before>from schedules_tools.handlers import ScheduleHandlerBase
from schedules_tools import models
from datetime import datetime
import os
import re
import tempfile
import logging
from lxml import etree

MSP_FLAGS_ATTRS = ('Flags', )
datetime_format = '%Y-%m-%dT%H:%M:%S'

PREFIX_FLAG = 'Flags'
PREFIX_LINK = 'Link'
PREFIX_NOTE = 'Note'

re_flags_separator = re.compile('[, ]+')
logger = logging.getLogger(__name__)

class ScheduleHandler_msp(ScheduleHandlerBase):
    provide_export = True

    @classmethod
    def is_valid_source(cls, handle=None):
        if not handle:
            handle = cls.handle
        try:
            if cls.src_storage_handler:
                handle = cls.src_storage_handler.get_local_handle()
            tree = etree.parse(handle)
        except etree.XMLSyntaxError:
            return False

        if 'http://schemas.microsoft.com/project' in tree.getroot().tag:
            return True
        return False

    def get_handle_mtime(self):
        return self._get_mtime_from_handle_file()

    # Schedule
    def import_schedule(self):
        self.schedule = models.Schedule()

        # remove project's xmlns
        tmp_file = tempfile.mkstemp()[1]
        with open(tmp_file, 'wt') as hTmp_file:
            for line in open(self.handle):
                hTmp_file.write(line.replace(' xmlns="http://schemas.microsoft.com/project"', ''))

        start_level = 1
        tree = etree.parse(tmp_file)

        eTask_list = tree.xpath('Tasks/Task[OutlineLevel >= %s]' % start_level)
        self.schedule.name = tree.xpath('Name|Title')[0].text
        self.schedule.project_name = self.schedule.name.strip()
        name_rx = re.match('(?P<name>.*?)(?P<version> [0-9]\S*)?$', self.schedule.name)
        if name_rx:
            self.schedule.project_name = name_rx.groupdict()['name']
            if name_rx.groupdict()['version']:
                # try to parse version out of string
                version_rx = re.match('(?P<major>[^._-]+)[._-]+(?P<minor>[^._-]+)?[._-]+(?P<maint>.+)?',
                                      name_rx.groupdict()['version'])
                if version_rx:
                    for number in version_rx.groupdict().iterkeys():
                        if version_rx.groupdict()[number]:
                            self.schedule._version[number] = version_rx.groupdict()[number].strip()

        self.schedule.name = self.schedule.name.strip()

        # import changelog, fill schedule.mtime
        if self.src_storage_handler:
            changelog_path = self.handle
            self.schedule.changelog = self.src_storage_handler.get_changelog(
                changelog_path)
            self.schedule.mtime = self.src_storage_handler.get_mtime(
                changelog_path)
        else:
            self.schedule.mtime = self.get_handle_mtime()

        # extended attributes
        for eExtAttr in tree.xpath('ExtendedAttributes/ExtendedAttribute'):
            fieldID = int(eExtAttr.xpath('FieldID')[0].text)
            fieldName = eExtAttr.xpath('FieldName')[0].text
            self.schedule.ext_attr[fieldName] = fieldID

        if self.schedule.ext_attr:  # choose flags field
            for ff_name in MSP_FLAGS_ATTRS:
                if ff_name in self.schedule.ext_attr:
                    self.schedule.flags_attr_id = self.schedule.ext_attr[ff_name]
                    break

        self.schedule.tasks = self._load_tasks_level(start_level, eTask_list)

        os.unlink(tmp_file)
        self.schedule.check_top_task()
        return self.schedule

    # Schedule
    def export_schedule(self, out_file=None):
        MSP_NAMESPACE = "http://schemas.microsoft.com/project"
        MSP = "{%s}" % MSP_NAMESPACE

        NSMAP = {None: MSP_NAMESPACE}  # the default namespace (no prefix)

        eProject = etree.Element(MSP + 'Project', nsmap=NSMAP)
        eName = etree.SubElement(eProject, 'Name')
        eName.text = self.schedule.name
        eTitle = etree.SubElement(eProject, 'Title')
        eTitle.text = self.schedule.name
        eMinutesPerDay = etree.SubElement(eProject, 'MinutesPerDay')
        eMinutesPerDay.text = '300'
        eMinutesPerWeek = etree.SubElement(eProject, 'MinutesPerWeek')
        eMinutesPerWeek.text = '1500'

        calendar_UID = '1'
        # add ext condition
        eCalendarUID = etree.SubElement(eProject, 'CalendarUID')
        eCalendarUID.text = calendar_UID
        eCalendars = etree.SubElement(eProject, 'Calendars')
        eCalendar = etree.SubElement(eCalendars, 'Calendar')
        eUID = etree.SubElement(eCalendar, 'UID')
        eUID.text = calendar_UID
        eName = etree.SubElement(eCalendar, 'Name')
        eName.text = '5h a day'
        eIsBaseCalendar = etree.SubElement(eCalendar, 'IsBaseCalendar')
        eIsBaseCalendar.text = '1'
        eWeekDays = etree.SubElement(eCalendar, 'WeekDays')
        for day_type in range(1, 8):
            eWeekDay = etree.SubElement(eWeekDays, 'WeekDay')
            eDayType = etree.SubElement(eWeekDay, 'DayType')
            eDayType.text = str(day_type)
            eDayWorking = etree.SubElement(eWeekDay, 'DayWorking')
            if day_type in (1, 7):
                eDayWorking.text = '0'
            else:
                eDayWorking.text = '1'
                eWorkingTimes = etree.SubElement(eWeekDay, 'WorkingTimes')
                for from_time in range(8, 17, 2):
                    eWorkingTime = etree.SubElement(eWorkingTimes, 'WorkingTime')
                    eFromTime = etree.SubElement(eWorkingTime, 'FromTime')
                    eFromTime.text = '%02d:00:00' % int(from_time)
                    eToTime = etree.SubElement(eWorkingTime, 'ToTime')
                    eToTime.text = '%02d:00:00' % (from_time + 1,)

        for r_id, resource in self.schedule.resources.items():
            eCalendar = etree.SubElement(eCalendars, 'Calendar')
            eUID = etree.SubElement(eCalendar, 'UID')
            eUID.text = str(r_id + 1)
            eName = etree.SubElement(eCalendar, 'Name')
            eName.text = resource
            eIsBaseCalendar = etree.SubElement(eCalendar, 'IsBaseCalendar')
            eIsBaseCalendar.text = '0'
            eBaseCalendarUID = etree.SubElement(eCalendar, 'BaseCalendarUID')
            eBaseCalendarUID.text = calendar_UID

        eTasks = etree.SubElement(eProject, 'Tasks')
        self.export_msp_tasks(self.schedule.tasks, eTasks, '')

        eResources = etree.SubElement(eProject, 'Resources')
        for r_id, resource in self.schedule.resources.items():
            eResource = etree.SubElement(eResources, 'Resource')

            eUID = etree.SubElement(eResource, 'UID')
            eUID.text = str(r_id)

            eID = etree.SubElement(eResource, 'ID')
            eID.text = str(r_id)

            eName = etree.SubElement(eResource, 'Name')
            eName.text = resource

            eCalendarUID = etree.SubElement(eResource, 'CalendarUID')
            eCalendarUID.text = str(r_id + 1)

        eAssignments = etree.SubElement(eProject, 'Assignments')
        for assignment in self.schedule.assignments:
            eAssignment = etree.SubElement(eAssignments, 'Assignment')

            eTaskUID = etree.SubElement(eAssignment, 'TaskUID')
            eTaskUID.text = str(assignment['t_id'])

            eResourceUID = etree.SubElement(eAssignment, 'ResourceUID')
            eResourceUID.text = str(assignment['r_id'])

            eUnits = etree.SubElement(eAssignment, 'Units')
            eUnits.text = '1'

        et = etree.ElementTree(eProject)

        if out_file:
            et.write(out_file, pretty_print=True, encoding="utf-8",
                     xml_declaration=True)

        return str(et)

    # Schedule
    def export_msp_tasks(self, tasks, eParent, outline_prefix):
        for n, task in enumerate(tasks, start=1):
            eTask = self.task_export_msp_node(task)
            eUID = etree.SubElement(eTask, 'UID')
            eUID.text = str(self.schedule._task_index)

            eID = etree.SubElement(eTask, 'ID')
            eID.text = str(self.schedule._task_index)

            if task.resource:
                self.schedule.assignments.append({
                    't_id': self.schedule._task_index,
                    'r_id': task.resource}
                )

            eOutlineNumber = etree.SubElement(eTask, 'OutlineNumber')
            eOutlineNumber.text = '%s%s' % (outline_prefix, n)

            eOutlineLevel = etree.SubElement(eTask, 'OutlineLevel')
            eOutlineLevel.text = str(len(outline_prefix.split('.')))

            eParent.append(eTask)

            self.schedule._task_index += 1
            self.export_msp_tasks(task.tasks, eParent,
                                  '%s.' % eOutlineNumber.text)

    # Schedule
    def _load_tasks_level(self, level, eTask_list):
        return_tasks = []

        while eTask_list:
            eTask = eTask_list[0]
            task_level = int(eTask.xpath('OutlineLevel')[0].text)

            if task_level > level:
                # return tasks may be empty since there could be no importable tasks yet
                if len(return_tasks):
                    return_tasks[-1].tasks = self._load_tasks_level(
                        task_level, eTask_list)
                else:
                    # remove task from list
                    eTask_list.pop(0)
                continue
            elif task_level < level:
                return return_tasks

            # process task
            task = models.Task(self.schedule)
            if self.task_load_msp_node(task, eTask):
                # update schedule start/end
                if self.schedule.dStart:
                    self.schedule.dStart = min(self.schedule.dStart, task.dStart)
                else:
                    self.schedule.dStart = task.dStart

                if self.schedule.dFinish:
                    self.schedule.dFinish = max(self.schedule.dFinish, task.dAcFinish)
                else:
                    self.schedule.dFinish = task.dAcFinish

                return_tasks.append(task)
            # remove task from list
            eTask_list.pop(0)
        return return_tasks

    # Task
    def task_load_msp_node(self, task, eTask):
        task.index = eTask.xpath('ID')[0].text

        task.name = task._workaround_it_phase_names(eTask)
        if not task.name:
            return False

        task.priority = int(eTask.xpath('Priority')[0].text)

        nlStart = eTask.xpath('Start')

        if nlStart:
            task.dStart = task.dAcStart = datetime.strptime(nlStart[0].text,
                                                   task._date_format)
            task.dFinish = task.dAcFinish = task.dStart
        else:
            return False

        nlFinish = eTask.xpath('Finish')
        if nlFinish:
            task.dFinish = task.dAcFinish = datetime.strptime(
                nlFinish[0].text, task._date_format)

        nlAcStart = eTask.xpath('ActualStart')
        if nlAcStart:
            task.dAcStart = datetime.strptime(nlAcStart[0].text,
                                              task._date_format)

        nlAcFinish = eTask.xpath('ActualFinish')
        if nlAcFinish:
            task.dAcFinish = datetime.strptime(nlAcFinish[0].text,
                                               task._date_format)

        # sanity check - if only ac start defined and beyond plan finish
        task.dAcFinish = max(task.dAcFinish, task.dAcStart)

        task.milestone = eTask.xpath('Milestone')[0].text == '1'

        ePercentComplete_list = eTask.xpath('PercentComplete')
        if ePercentComplete_list:
            task.p_complete = eTask.xpath('PercentComplete')[0].text

        notes = eTask.xpath('Notes')
        if notes:
            task.note = notes[0].text.strip()

        # load flags from ext attributes
        flag_ext_attr = eTask.xpath(
            'ExtendedAttribute[FieldID = {}]'.format(
                task._schedule.flags_attr_id)
        )
        if flag_ext_attr:
            flags_value = flag_ext_attr[0].xpath('Value')[0].text
            if flags_value:
                task.flags = [f for f in flags_value.strip(' ,\n').split(',') if ' ' not in f]
                task._schedule.used_flags |= set(task.flags)

        # workaround for SmartSheet exports - load flags, links
        ext_attr_elements = eTask.xpath('ExtendedAttribute/Value')
        for ext_attr in ext_attr_elements:
            self._parse_extended_attr(task, ext_attr)

        task.check_for_phase()
        return True

    def _parse_extended_attr(self, task, element):
        """According to content of element will guess if the value is flag
        or link definition.

        @param element: XPath element instance (/Project/Tasks/Task/ExtendedAttribute/Value)
        """
        element_val = element.text
        if not element_val:
            return

        pieces = element_val.split(':', 1)
        if len(pieces) != 2:
            # it's not a string in format 'Flag: qe, dev' - don't process
            return

        key, val = pieces
        key = key.strip().lower()
        val = val.strip()

        if key == PREFIX_FLAG.lower():
            val = val.lower()
            flags = re_flags_separator.split(val)
            for flag in flags:
                flag = flag.strip()
                if flag:
                    task.flags.append(flag)
            self.schedule.used_flags |= set(task.flags)
        elif key == PREFIX_LINK.lower():
            task.link = val
        elif key == PREFIX_NOTE.lower():
            task.note = val
        else:
            logger.warn('Extended attr "{}" wasn\'t recognized.'.format(key))

    # Task
    def task_export_msp_node(self, task):
        eTask = etree.Element('Task')

        eName = etree.SubElement(eTask, 'Name')
        eName.text = task.name

        eStart = etree.SubElement(eTask, 'Start')
        eStart.text = task.dStart.strftime(datetime_format)
        eFinish = etree.SubElement(eTask, 'Finish')
        eFinish.text = task.dFinish.strftime(datetime_format)

        ePriority = etree.SubElement(eTask, 'Priority')
        ePriority.text = str(task.priority)

        eMilestone = etree.SubElement(eTask, 'Milestone')
        eMilestone.text = str(task.milestone)

        duration = task.dFinish - task.dStart
        eDuration = etree.SubElement(eTask, 'Duration')
        h, rem = divmod(duration.seconds, 3600)
        m, s = divmod(rem, 60)
        eDuration.text = 'PT%sH%sM%sS' % (h, m, s)
        eDurationFormat = etree.SubElement(eTask, 'DurationFormat')
        eDurationFormat.text = '39'

        if task.note:
            eNotes = etree.SubElement(eTask, 'Notes')
            eNotes.text = task.note

        flags_str = ','.join(task.flags)
        if flags_str:
            ext_attr_element = etree.SubElement(eTask, 'ExtendedAttribute')
            value_element = etree.SubElement(ext_attr_element, 'Value')
            value_element.text = '{}: {}'.format(PREFIX_FLAG, flags_str)

            # this value is not used, but required by SmartSheets import
            fieldid_element = etree.SubElement(ext_attr_element, 'FieldID')
            fieldid_element.text = '188743734'
        if task.link:
            ext_attr_element = etree.SubElement(eTask, 'ExtendedAttribute')
            value_element = etree.SubElement(ext_attr_element, 'Value')
            value_element.text = '{}: {}'.format(PREFIX_LINK, task.link)

            # this value is not used, but required by SmartSheets import
            fieldid_element = etree.SubElement(ext_attr_element, 'FieldID')
            fieldid_element.text = '188743737'

        return eTask
<commit_msg>msp handler: be more robust in is_valid_source<commit_after>from schedules_tools.handlers import ScheduleHandlerBase
from schedules_tools import models
from datetime import datetime
import os
import re
import tempfile
import logging
from lxml import etree

MSP_FLAGS_ATTRS = ('Flags', )
datetime_format = '%Y-%m-%dT%H:%M:%S'

PREFIX_FLAG = 'Flags'
PREFIX_LINK = 'Link'
PREFIX_NOTE = 'Note'

re_flags_separator = re.compile('[, ]+')
logger = logging.getLogger(__name__)

class ScheduleHandler_msp(ScheduleHandlerBase):
    provide_export = True

    @classmethod
    def is_valid_source(cls, handle=None):
        if not handle:
            handle = cls.handle
        try:
            if cls.src_storage_handler:
                handle = cls.src_storage_handler.get_local_handle()
            tree = etree.parse(handle)
        except (etree.XMLSyntaxError, IOError):
            return False

        if 'http://schemas.microsoft.com/project' in tree.getroot().tag:
            return True
        return False

    def get_handle_mtime(self):
        return self._get_mtime_from_handle_file()

    # Schedule
    def import_schedule(self):
        self.schedule = models.Schedule()

        # remove project's xmlns
        tmp_file = tempfile.mkstemp()[1]
        with open(tmp_file, 'wt') as hTmp_file:
            for line in open(self.handle):
                hTmp_file.write(line.replace(' xmlns="http://schemas.microsoft.com/project"', ''))

        start_level = 1
        tree = etree.parse(tmp_file)

        eTask_list = tree.xpath('Tasks/Task[OutlineLevel >= %s]' % start_level)
        self.schedule.name = tree.xpath('Name|Title')[0].text
        self.schedule.project_name = self.schedule.name.strip()
        name_rx = re.match('(?P<name>.*?)(?P<version> [0-9]\S*)?$', self.schedule.name)
        if name_rx:
            self.schedule.project_name = name_rx.groupdict()['name']
            if name_rx.groupdict()['version']:
                # try to parse version out of string
                version_rx = re.match('(?P<major>[^._-]+)[._-]+(?P<minor>[^._-]+)?[._-]+(?P<maint>.+)?',
                                      name_rx.groupdict()['version'])
                if version_rx:
                    for number in version_rx.groupdict().iterkeys():
                        if version_rx.groupdict()[number]:
                            self.schedule._version[number] = version_rx.groupdict()[number].strip()

        self.schedule.name = self.schedule.name.strip()

        # import changelog, fill schedule.mtime
        if self.src_storage_handler:
            changelog_path = self.handle
            self.schedule.changelog = self.src_storage_handler.get_changelog(
                changelog_path)
            self.schedule.mtime = self.src_storage_handler.get_mtime(
                changelog_path)
        else:
            self.schedule.mtime = self.get_handle_mtime()

        # extended attributes
        for eExtAttr in tree.xpath('ExtendedAttributes/ExtendedAttribute'):
            fieldID = int(eExtAttr.xpath('FieldID')[0].text)
            fieldName = eExtAttr.xpath('FieldName')[0].text
            self.schedule.ext_attr[fieldName] = fieldID

        if self.schedule.ext_attr:  # choose flags field
            for ff_name in MSP_FLAGS_ATTRS:
                if ff_name in self.schedule.ext_attr:
                    self.schedule.flags_attr_id = self.schedule.ext_attr[ff_name]
                    break

        self.schedule.tasks = self._load_tasks_level(start_level, eTask_list)

        os.unlink(tmp_file)
        self.schedule.check_top_task()
        return self.schedule

    # Schedule
    def export_schedule(self, out_file=None):
        MSP_NAMESPACE = "http://schemas.microsoft.com/project"
        MSP = "{%s}" % MSP_NAMESPACE

        NSMAP = {None: MSP_NAMESPACE}  # the default namespace (no prefix)

        eProject = etree.Element(MSP + 'Project', nsmap=NSMAP)
        eName = etree.SubElement(eProject, 'Name')
        eName.text = self.schedule.name
        eTitle = etree.SubElement(eProject, 'Title')
        eTitle.text = self.schedule.name
        eMinutesPerDay = etree.SubElement(eProject, 'MinutesPerDay')
        eMinutesPerDay.text = '300'
        eMinutesPerWeek = etree.SubElement(eProject, 'MinutesPerWeek')
        eMinutesPerWeek.text = '1500'

        calendar_UID = '1'
        # add ext condition
        eCalendarUID = etree.SubElement(eProject, 'CalendarUID')
        eCalendarUID.text = calendar_UID
        eCalendars = etree.SubElement(eProject, 'Calendars')
        eCalendar = etree.SubElement(eCalendars, 'Calendar')
        eUID = etree.SubElement(eCalendar, 'UID')
        eUID.text = calendar_UID
        eName = etree.SubElement(eCalendar, 'Name')
        eName.text = '5h a day'
        eIsBaseCalendar = etree.SubElement(eCalendar, 'IsBaseCalendar')
        eIsBaseCalendar.text = '1'
        eWeekDays = etree.SubElement(eCalendar, 'WeekDays')
        for day_type in range(1, 8):
            eWeekDay = etree.SubElement(eWeekDays, 'WeekDay')
            eDayType = etree.SubElement(eWeekDay, 'DayType')
            eDayType.text = str(day_type)
            eDayWorking = etree.SubElement(eWeekDay, 'DayWorking')
            if day_type in (1, 7):
                eDayWorking.text = '0'
            else:
                eDayWorking.text = '1'
                eWorkingTimes = etree.SubElement(eWeekDay, 'WorkingTimes')
                for from_time in range(8, 17, 2):
                    eWorkingTime = etree.SubElement(eWorkingTimes, 'WorkingTime')
                    eFromTime = etree.SubElement(eWorkingTime, 'FromTime')
                    eFromTime.text = '%02d:00:00' % int(from_time)
                    eToTime = etree.SubElement(eWorkingTime, 'ToTime')
                    eToTime.text = '%02d:00:00' % (from_time + 1,)

        for r_id, resource in self.schedule.resources.items():
            eCalendar = etree.SubElement(eCalendars, 'Calendar')
            eUID = etree.SubElement(eCalendar, 'UID')
            eUID.text = str(r_id + 1)
            eName = etree.SubElement(eCalendar, 'Name')
            eName.text = resource
            eIsBaseCalendar = etree.SubElement(eCalendar, 'IsBaseCalendar')
            eIsBaseCalendar.text = '0'
            eBaseCalendarUID = etree.SubElement(eCalendar, 'BaseCalendarUID')
            eBaseCalendarUID.text = calendar_UID

        eTasks = etree.SubElement(eProject, 'Tasks')
        self.export_msp_tasks(self.schedule.tasks, eTasks, '')

        eResources = etree.SubElement(eProject, 'Resources')
        for r_id, resource in self.schedule.resources.items():
            eResource = etree.SubElement(eResources, 'Resource')

            eUID = etree.SubElement(eResource, 'UID')
            eUID.text = str(r_id)

            eID = etree.SubElement(eResource, 'ID')
            eID.text = str(r_id)

            eName = etree.SubElement(eResource, 'Name')
            eName.text = resource

            eCalendarUID = etree.SubElement(eResource, 'CalendarUID')
            eCalendarUID.text = str(r_id + 1)

        eAssignments = etree.SubElement(eProject, 'Assignments')
        for assignment in self.schedule.assignments:
            eAssignment = etree.SubElement(eAssignments, 'Assignment')

            eTaskUID = etree.SubElement(eAssignment, 'TaskUID')
            eTaskUID.text = str(assignment['t_id'])

            eResourceUID = etree.SubElement(eAssignment, 'ResourceUID')
            eResourceUID.text = str(assignment['r_id'])

            eUnits = etree.SubElement(eAssignment, 'Units')
            eUnits.text = '1'

        et = etree.ElementTree(eProject)

        if out_file:
            et.write(out_file, pretty_print=True, encoding="utf-8",
                     xml_declaration=True)

        return str(et)

    # Schedule
    def export_msp_tasks(self, tasks, eParent, outline_prefix):
        for n, task in enumerate(tasks, start=1):
            eTask = self.task_export_msp_node(task)
            eUID = etree.SubElement(eTask, 'UID')
            eUID.text = str(self.schedule._task_index)

            eID = etree.SubElement(eTask, 'ID')
            eID.text = str(self.schedule._task_index)

            if task.resource:
                self.schedule.assignments.append({
                    't_id': self.schedule._task_index,
                    'r_id': task.resource}
                )

            eOutlineNumber = etree.SubElement(eTask, 'OutlineNumber')
            eOutlineNumber.text = '%s%s' % (outline_prefix, n)

            eOutlineLevel = etree.SubElement(eTask, 'OutlineLevel')
            eOutlineLevel.text = str(len(outline_prefix.split('.')))

            eParent.append(eTask)

            self.schedule._task_index += 1
            self.export_msp_tasks(task.tasks, eParent,
                                  '%s.' % eOutlineNumber.text)

    # Schedule
    def _load_tasks_level(self, level, eTask_list):
        return_tasks = []

        while eTask_list:
            eTask = eTask_list[0]
            task_level = int(eTask.xpath('OutlineLevel')[0].text)

            if task_level > level:
                # return tasks may be empty since there could be no importable tasks yet
                if len(return_tasks):
                    return_tasks[-1].tasks = self._load_tasks_level(
                        task_level, eTask_list)
                else:
                    # remove task from list
                    eTask_list.pop(0)
                continue
            elif task_level < level:
                return return_tasks

            # process task
            task = models.Task(self.schedule)
            if self.task_load_msp_node(task, eTask):
                # update schedule start/end
                if self.schedule.dStart:
                    self.schedule.dStart = min(self.schedule.dStart, task.dStart)
                else:
                    self.schedule.dStart = task.dStart

                if self.schedule.dFinish:
                    self.schedule.dFinish = max(self.schedule.dFinish, task.dAcFinish)
                else:
                    self.schedule.dFinish = task.dAcFinish

                return_tasks.append(task)
            # remove task from list
            eTask_list.pop(0)
        return return_tasks

    # Task
    def task_load_msp_node(self, task, eTask):
        task.index = eTask.xpath('ID')[0].text

        task.name = task._workaround_it_phase_names(eTask)
        if not task.name:
            return False

        task.priority = int(eTask.xpath('Priority')[0].text)

        nlStart = eTask.xpath('Start')

        if nlStart:
            task.dStart = task.dAcStart = datetime.strptime(nlStart[0].text,
                                                   task._date_format)
            task.dFinish = task.dAcFinish = task.dStart
        else:
            return False

        nlFinish = eTask.xpath('Finish')
        if nlFinish:
            task.dFinish = task.dAcFinish = datetime.strptime(
                nlFinish[0].text, task._date_format)

        nlAcStart = eTask.xpath('ActualStart')
        if nlAcStart:
            task.dAcStart = datetime.strptime(nlAcStart[0].text,
                                              task._date_format)

        nlAcFinish = eTask.xpath('ActualFinish')
        if nlAcFinish:
            task.dAcFinish = datetime.strptime(nlAcFinish[0].text,
                                               task._date_format)

        # sanity check - if only ac start defined and beyond plan finish
        task.dAcFinish = max(task.dAcFinish, task.dAcStart)

        task.milestone = eTask.xpath('Milestone')[0].text == '1'

        ePercentComplete_list = eTask.xpath('PercentComplete')
        if ePercentComplete_list:
            task.p_complete = eTask.xpath('PercentComplete')[0].text

        notes = eTask.xpath('Notes')
        if notes:
            task.note = notes[0].text.strip()

        # load flags from ext attributes
        flag_ext_attr = eTask.xpath(
            'ExtendedAttribute[FieldID = {}]'.format(
                task._schedule.flags_attr_id)
        )
        if flag_ext_attr:
            flags_value = flag_ext_attr[0].xpath('Value')[0].text
            if flags_value:
                task.flags = [f for f in flags_value.strip(' ,\n').split(',') if ' ' not in f]
                task._schedule.used_flags |= set(task.flags)

        # workaround for SmartSheet exports - load flags, links
        ext_attr_elements = eTask.xpath('ExtendedAttribute/Value')
        for ext_attr in ext_attr_elements:
            self._parse_extended_attr(task, ext_attr)

        task.check_for_phase()
        return True

    def _parse_extended_attr(self, task, element):
        """According to content of element will guess if the value is flag
        or link definition.

        @param element: XPath element instance (/Project/Tasks/Task/ExtendedAttribute/Value)
        """
        element_val = element.text
        if not element_val:
            return

        pieces = element_val.split(':', 1)
        if len(pieces) != 2:
            # it's not a string in format 'Flag: qe, dev' - don't process
            return

        key, val = pieces
        key = key.strip().lower()
        val = val.strip()

        if key == PREFIX_FLAG.lower():
            val = val.lower()
            flags = re_flags_separator.split(val)
            for flag in flags:
                flag = flag.strip()
                if flag:
                    task.flags.append(flag)
            self.schedule.used_flags |= set(task.flags)
        elif key == PREFIX_LINK.lower():
            task.link = val
        elif key == PREFIX_NOTE.lower():
            task.note = val
        else:
            logger.warn('Extended attr "{}" wasn\'t recognized.'.format(key))

    # Task
    def task_export_msp_node(self, task):
        eTask = etree.Element('Task')

        eName = etree.SubElement(eTask, 'Name')
        eName.text = task.name

        eStart = etree.SubElement(eTask, 'Start')
        eStart.text = task.dStart.strftime(datetime_format)
        eFinish = etree.SubElement(eTask, 'Finish')
        eFinish.text = task.dFinish.strftime(datetime_format)

        ePriority = etree.SubElement(eTask, 'Priority')
        ePriority.text = str(task.priority)

        eMilestone = etree.SubElement(eTask, 'Milestone')
        eMilestone.text = str(task.milestone)

        duration = task.dFinish - task.dStart
        eDuration = etree.SubElement(eTask, 'Duration')
        h, rem = divmod(duration.seconds, 3600)
        m, s = divmod(rem, 60)
        eDuration.text = 'PT%sH%sM%sS' % (h, m, s)
        eDurationFormat = etree.SubElement(eTask, 'DurationFormat')
        eDurationFormat.text = '39'

        if task.note:
            eNotes = etree.SubElement(eTask, 'Notes')
            eNotes.text = task.note

        flags_str = ','.join(task.flags)
        if flags_str:
            ext_attr_element = etree.SubElement(eTask, 'ExtendedAttribute')
            value_element = etree.SubElement(ext_attr_element, 'Value')
            value_element.text = '{}: {}'.format(PREFIX_FLAG, flags_str)

            # this value is not used, but required by SmartSheets import
            fieldid_element = etree.SubElement(ext_attr_element, 'FieldID')
            fieldid_element.text = '188743734'
        if task.link:
            ext_attr_element = etree.SubElement(eTask, 'ExtendedAttribute')
            value_element = etree.SubElement(ext_attr_element, 'Value')
            value_element.text = '{}: {}'.format(PREFIX_LINK, task.link)

            # this value is not used, but required by SmartSheets import
            fieldid_element = etree.SubElement(ext_attr_element, 'FieldID')
            fieldid_element.text = '188743737'

        return eTask
<|endoftext|>


<commit_before>"""

acreroad_1420 Drive software

Software designed to drive the 1420 MHz telescope on the roof of the
Acre Road observatory. This class interacts with "qp", the telescope
drive software by Norman Gray (https://bitbucket.org/nxg/qp) via a
serial (USB) interface.

The serial interfacing is done through pySerial.

Parameters
----------

device : str
   The name of the unix device which the drive is connected to
port : int
   The port number of the drive
simulate : bool
   A boolean flag to set the drive in simulation mode
   (the class does not connect to the controller in simulation mode)

"""

import re, datetime

import astropy
from astropy.coordinates import SkyCoord, ICRS, EarthLocation
import astropy.units as u
import serial
from astropy.time import Time

class Drive():

    sim = 0
    acre_road = EarthLocation(lat=55.9024278*u.deg, lon=-4.307582*u.deg, height=61*u.m)

    # Position variables
    ra = 0
    dec = 0

    # String formats

    com_format = re.compile("[A-Za-z]{1,2} ?([-+]?[0-9]{0,4}\.[0-9]{0,8} ?){0,6}") # general command string format
    cal_format = re.compile("[0-9]{3} [0-9]{3}") # calibration string format

    def __init__(self, device, baud, timeout=None, simulate=0, calibration=None, location=acre_road):
        """
        Software designed to drive the 1420 MHz telescope on the roof of the
        Acre Road observatory. This class interacts with "qp", the telescope
        drive software by Norman Gray (https://bitbucket.org/nxg/qp) via a
        serial (USB) interface.

        The serial interfacing is done through pySerial.

        Parameters
        ----------

        device : str
           The name of the unix device which the drive is connected to
        baud : int
           The baud-rate of the connection.
        timeout : int
           The time, in seconds, to wait before timing-out. Default is 2 seconds.
        simulate : bool
           A boolean flag to set the drive in simulation mode
           (the class does not connect to the controller in simulation mode)
        calibration : str
           The calibration figures which have been returned by a previous run of the *calibrate()* method.
           The default is `None` which forces a calibration run to be carried-out.
        location : astropy.coordinates.EarthLocation object
           The Earth location of the telescope. The default is `None`, which sets the location as Acre Road Observatory, Glasgow.

        Examples
        --------

        >>> from acreroad_1420 import drive
        >>> connection = drive.Drive('/dev/tty.usbserial', 9600, simulate=1)

        """
        self.sim = simulate
        self.timeout = timeout
        self.location = location
        # Initialise the connection
        if not self.sim: self._openconnection(device, baud)

        self.calibrate(calibration)
        self.setTime()

        self.setLocation(location)
        pass

    def _openconnection(self, device, baud):
        self.ser = serial.Serial(device, baud, timeout=self.timeout)

    def _command(self, string):
        """
        Passes commands to the controller.
        """

        # Check that the command string is a string, and that it
        # matches the format required for a command string
        string = str(string)
        string = string+"\n"

        if not self.com_format.match(string):
            raise ValueError(string+" : This string doesn't have the format of a valid controller command.'")

        if self.sim:
            print("In simulation mode, command ignored.")
            return 1
        else:
            # Pass the command to the Arduino via pyserial
            self.ser.write(string.encode('ascii'))

            # Retrieve the return message from the controller
            ret_line =  self.ser.readline()
            if ret_line: return ret_line
            else : return 1

    def calibrate(self, values=None):
        """
        Carries-out a calibration run, returning two numbers which are offsets, or sets the calibration if known values are provided.

        Parameters
        ----------

        values : str
           The calibration values produced by a previous calibration run of the telescope, provided in the format "nnn nnn"

        """
        if values:
            # Check the format of the values string.
            if self.cal_format.match(values):
                return self._command("c "+values)
        if self.sim: return ">c 000 000"
        pass

    def setTime(self):
        """
        Sets the time on the drive controller's clock to the current system time.

        """
        time = datetime.datetime.utcnow()

        command_str = "T {} {} {} {} {} {}".format(time.year, time.month, time.day, time.hour, time.minute, time.second)

        return self._command(command_str)

    def setLocation(self, location=None, dlat=0, dlon=0, azimuth=90, altitude=-6):
        """
        Sets the location of the telescope.

        Parameters
        ----------
        location : astropy.coordinates.EarthLocation object
        """
        if not location:
            # Assume we're at Acre Road, in Glasgow
            location = self.acre_road

        latitude  = location.latitude
        longitude = location.longitude

        # Construct the command
        command_str = "O {} {} {} {} {} {}".format(latitude, longitude, dlat, dlon, azimuth, altitude)
        return self._command(command_str)

    def goto(self, skycoord, track=True):
        """
        Moves the telescope to point at a given sky location, and then commands the drive to track the point.

        Parameters
        ----------
        skycoord : astropy.SkyCoord object
           An astropy SkyCoord object which contains the sky location to slew to.
           This can also be a list of locations which the telescope will slew to sequentially.

        """

        if not type(skycoord)==astropy.coordinates.sky_coordinate.SkyCoord:
            raise ValueError("The sky coordinates provided aren't an astropy SkyCoord object!'")

        # To do : We need to make sure that this behaves nicely with a
        # list of coordinates as well as single ones.

        # We need to make sure that we send RA and DEC values to the
        # controller (it is theoretically possible to do ALTAZ too,
        # but let's start simple, eh?')

        skycoord = skycoord.transform_to(frame=ICRS)

        time = Time.now()

        ra  = skycoord.ra
        dec = skycoord.dec

        skycoord = skycoord.transform_to(AltAz(obstime=time, location=self.location))

        # construct a command string
        command_str = "gh {} {}".format(skycoord.az, skycoord.alt)
        # pass the slew-to command to the controller
        if self._command(command_str):
            # We need to do more here than just check the command has
            # passed successfully; we need to check the telescope is
            # (or thinks it is) pointing where we asked.
            if self.sim:
                self.ra  = ra
                self.dec = dec

            # If the command completes then set the controller to track the object
            if track:
                command_str = "ts"
                return self._command(command_str)
            else: return 1
        else:
            raise ControllerException("The telescope has failed to slew to the requested location")

    def home(self):
        """
        Slews the telescope to the home position.

        """
        command_str = "gH"
        return self._command(command_str)

    def status(self):
        """
        Returns a dictionary describing the status of the telescope (e.g. its location).

        Returns
        -------
        dict
           A dictionary containing the right ascension, declination, altitude, and azimuth of the telescope.

        Examples
        --------

        >>> from astropy.coordinates import SkyCoord
        >>> from astropy.coordinates import ICRS,
        >>> c = SkyCoord(frame="galactic", l="1h12m43.2s", b="+1d12m43s")
        >>> self.connection.goto(c)
        >>> ra = self.connection.status()['ra'].value

        """
        return {'ra':self.ra, 'dec': self.dec, 'alt':0, 'az':0}

class ControllerException(Exception):
    pass<commit_msg>Changed goto to work in horizontal coordinates for now.<commit_after>"""

acreroad_1420 Drive software

Software designed to drive the 1420 MHz telescope on the roof of the
Acre Road observatory. This class interacts with "qp", the telescope
drive software by Norman Gray (https://bitbucket.org/nxg/qp) via a
serial (USB) interface.

The serial interfacing is done through pySerial.

Parameters
----------

device : str
   The name of the unix device which the drive is connected to
port : int
   The port number of the drive
simulate : bool
   A boolean flag to set the drive in simulation mode
   (the class does not connect to the controller in simulation mode)

"""

import re, datetime

import astropy
from astropy.coordinates import SkyCoord, ICRS, EarthLocation, AltAz
import astropy.units as u
import serial
from astropy.time import Time

class Drive():

    sim = 0
    acre_road = EarthLocation(lat=55.9024278*u.deg, lon=-4.307582*u.deg, height=61*u.m)

    # Position variables
    ra = 0
    dec = 0

    # String formats

    com_format = re.compile("[A-Za-z]{1,2} ?([-+]?[0-9]{0,4}\.[0-9]{0,8} ?){0,6}") # general command string format
    cal_format = re.compile("[0-9]{3} [0-9]{3}") # calibration string format

    def __init__(self, device, baud, timeout=None, simulate=0, calibration=None, location=acre_road):
        """
        Software designed to drive the 1420 MHz telescope on the roof of the
        Acre Road observatory. This class interacts with "qp", the telescope
        drive software by Norman Gray (https://bitbucket.org/nxg/qp) via a
        serial (USB) interface.

        The serial interfacing is done through pySerial.

        Parameters
        ----------

        device : str
           The name of the unix device which the drive is connected to
        baud : int
           The baud-rate of the connection.
        timeout : int
           The time, in seconds, to wait before timing-out. Default is 2 seconds.
        simulate : bool
           A boolean flag to set the drive in simulation mode
           (the class does not connect to the controller in simulation mode)
        calibration : str
           The calibration figures which have been returned by a previous run of the *calibrate()* method.
           The default is `None` which forces a calibration run to be carried-out.
        location : astropy.coordinates.EarthLocation object
           The Earth location of the telescope. The default is `None`, which sets the location as Acre Road Observatory, Glasgow.

        Examples
        --------

        >>> from acreroad_1420 import drive
        >>> connection = drive.Drive('/dev/tty.usbserial', 9600, simulate=1)

        """
        self.sim = simulate
        self.timeout = timeout
        self.location = location
        # Initialise the connection
        if not self.sim: self._openconnection(device, baud)

        self.calibrate(calibration)
        self.setTime()

        self.setLocation(location)
        pass

    def _openconnection(self, device, baud):
        self.ser = serial.Serial(device, baud, timeout=self.timeout)

    def _command(self, string):
        """
        Passes commands to the controller.
        """

        # Check that the command string is a string, and that it
        # matches the format required for a command string
        string = str(string)
        string = string+"\n"

        if not self.com_format.match(string):
            raise ValueError(string+" : This string doesn't have the format of a valid controller command.'")

        if self.sim:
            print("In simulation mode, command ignored.")
            return 1
        else:
            # Pass the command to the Arduino via pyserial
            self.ser.write(string.encode('ascii'))

            # Retrieve the return message from the controller
            ret_line =  self.ser.readline()
            if ret_line: return ret_line
            else : return 1

    def calibrate(self, values=None):
        """
        Carries-out a calibration run, returning two numbers which are offsets, or sets the calibration if known values are provided.

        Parameters
        ----------

        values : str
           The calibration values produced by a previous calibration run of the telescope, provided in the format "nnn nnn"

        """
        if values:
            # Check the format of the values string.
            if self.cal_format.match(values):
                return self._command("c "+values)
        if self.sim: return ">c 000 000"
        pass

    def setTime(self):
        """
        Sets the time on the drive controller's clock to the current system time.

        """
        time = datetime.datetime.utcnow()

        command_str = "T {} {} {} {} {} {}".format(time.year, time.month, time.day, time.hour, time.minute, time.second)

        return self._command(command_str)

    def setLocation(self, location=None, dlat=0, dlon=0, azimuth=90, altitude=-6):
        """
        Sets the location of the telescope.

        Parameters
        ----------
        location : astropy.coordinates.EarthLocation object
        """
        if not location:
            # Assume we're at Acre Road, in Glasgow
            location = self.acre_road

        latitude  = location.latitude
        longitude = location.longitude

        # Construct the command
        command_str = "O {} {} {} {} {} {}".format(latitude, longitude, dlat, dlon, azimuth, altitude)
        return self._command(command_str)

    def goto(self, skycoord, track=True):
        """
        Moves the telescope to point at a given sky location, and then commands the drive to track the point.

        Parameters
        ----------
        skycoord : astropy.SkyCoord object
           An astropy SkyCoord object which contains the sky location to slew to.
           This can also be a list of locations which the telescope will slew to sequentially.

        """

        if not type(skycoord)==astropy.coordinates.sky_coordinate.SkyCoord:
            raise ValueError("The sky coordinates provided aren't an astropy SkyCoord object!'")

        # To do : We need to make sure that this behaves nicely with a
        # list of coordinates as well as single ones.

        # We need to make sure that we send RA and DEC values to the
        # controller (it is theoretically possible to do ALTAZ too,
        # but let's start simple, eh?')

        skycoord = skycoord.transform_to(frame=ICRS)

        time = Time.now()

        ra  = skycoord.ra
        dec = skycoord.dec

        skycoord = skycoord.transform_to(AltAz(obstime=time, location=self.location))

        # construct a command string
        command_str = "gh {} {}".format(skycoord.az, skycoord.alt)
        # pass the slew-to command to the controller
        if self._command(command_str):
            # We need to do more here than just check the command has
            # passed successfully; we need to check the telescope is
            # (or thinks it is) pointing where we asked.
            if self.sim:
                self.ra  = ra
                self.dec = dec

            # If the command completes then set the controller to track the object
            if track:
                command_str = "ts"
                return self._command(command_str)
            else: return 1
        else:
            raise ControllerException("The telescope has failed to slew to the requested location")

    def home(self):
        """
        Slews the telescope to the home position.

        """
        command_str = "gH"
        return self._command(command_str)

    def status(self):
        """
        Returns a dictionary describing the status of the telescope (e.g. its location).

        Returns
        -------
        dict
           A dictionary containing the right ascension, declination, altitude, and azimuth of the telescope.

        Examples
        --------

        >>> from astropy.coordinates import SkyCoord
        >>> from astropy.coordinates import ICRS,
        >>> c = SkyCoord(frame="galactic", l="1h12m43.2s", b="+1d12m43s")
        >>> self.connection.goto(c)
        >>> ra = self.connection.status()['ra'].value

        """
        return {'ra':self.ra, 'dec': self.dec, 'alt':0, 'az':0}

class ControllerException(Exception):
    pass<|endoftext|>


<commit_before>#! /usr/bin/env python
# -*- encoding: UTF-8 -*-

# DOCUMENTATION
#http://doc.aldebaran.com/2-5/naoqi/peopleperception/alpeopleperception-api.html#PeoplePerception
#http://doc.aldebaran.com/2-5/naoqi/peopleperception/alfacedetection.html

import qi
import argparse
import sys
import os
import time
import threading
import json

from naoqi import ALProxy

from utils import point2world
import conditions
from conditions import set_condition

##PARAMS
confidenceThreshold=0.4 #default =0.4
min_size={'w':0.10 ,'h':0.10}

def writeJson(path,data):

    if not os.path.exists(path):
        os.makedirs(path)

    jsonData=json.dumps(data,indent=2 )

    fh = open(path+'userlearned.json',"w")
    fh.write(jsonData)
    fh.close()

def onFaceDetection(facevalues):

    for faceInfo in facevalues:
        # First Field = Shape info.
        faceShapeInfo = faceInfo[0]
        # Second Field = Extra info .
        faceExtraInfo = faceInfo[1]

        #print 'eyeLeft',round(faceExtraInfo[3][0],2),'eyeRight',round(faceExtraInfo[4][0],2)
        pose={'alpha': round(faceShapeInfo[1],2), 'beta': round(faceShapeInfo[2],2), 'width':round(faceShapeInfo[3],2), 'height': round(faceShapeInfo[4],2)}
        face={'faceid': faceExtraInfo[0], 'user': faceExtraInfo[2], 'confidence':round(faceExtraInfo[1],3), 'pose': pose}

    return face

def rhMonitorThread (memory_service):
    t = threading.currentThread()
    print "learnperson thread started"
    print "waiting for face"
    while getattr(t, "do_run", True):

        facevalues =  memory_service.getData("FaceDetected")
        peoplevisible =  memory_service.getData("PeoplePerception/VisiblePeopleList")

        if facevalues:

            facesarray=facevalues[1][0:len(facevalues[1]) -1]

            #### TODO: Now only valid when detects 1 face
            if len(facesarray) is 1:

                face=onFaceDetection(facesarray)

                print 'pose',face['pose']
                if float(face['pose']['width'])>= float(min_size['w']) and float(face['pose']['height'])>=min_size['h']:

                    #####
                    ##### ID ????
                    print 'Learning face ',face['faceid']
                    learned=faces_service.learnFace('user-'+str(face['faceid']))
                    print 'Face learned =',learned
                    face['user']='user-'+str(face['faceid'])

                    if learned is True:

                        person=None
                        poseinworld=None

                        if len(peoplevisible) is 1:

                            personid=peoplevisible[0]
                            print 'personid= ',personid
                            #try:

        #                    lookingat =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/LookingAtRobotScore")
        #                    gazedirection =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/GazeDirection")
        #                    age =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/AgeProperties")
        #                    gender =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/GenderProperties")
        #                    expression =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/ExpressionProperties")
        #                    smile =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/SmileProperties")
        #                    print 'lookingat=',lookingat
        #                    print 'gazedirection=',gazedirection
        #                    print 'age=',age
        #                    print 'gender=',gender
        #                    print 'expression=',expression
        #                    print 'smile=',smile
        #

                            angles =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/AnglesYawPitch")
                            distance =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/Distance")
                            height =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/RealHeight")
                            shirtcolor =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/ShirtColor")
                            shirtcolorHSV =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/ShirtColorHSV")
                            PositionInRobotFrame =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/PositionInRobotFrame")
        #                    issitting=memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/IsSitting")
        #
        #
        #                    print 'issitting=',issitting

                            # Write data in json format

                            personinfo={'height': round(height,2), 'shirtcolor': shirtcolor, 'shirtcolorHSV':shirtcolorHSV}

                            w_px, w_py = point2world(memory_service,[PositionInRobotFrame[0],PositionInRobotFrame[1]])
                            poseinworld={'x':w_px,'y': w_py}

                            person={'personid': personid, 'info': personinfo}

                        posetopological={'current_node':'TODO','closest_node':'TODO'}
                        lastlocation= {'world':poseinworld , 'topological':posetopological}

                        user={'face': face,'person': person,'lastlocation':lastlocation}

                        writeJson('./data/',user)

                        print 'EXIT TRUE'
                        quit()
                        break

                        set_condition(memory_service,'learnperson','true')

                    elif learned is False:

                        print 'EXIT FALSE'
                        quit()
                        set_condition(memory_service,'learnperson','false')

        time.sleep(0.2)
    print "learnperson thread quit"

def init(session):
    global memory_service
    global monitorThread

    global faces_service
    print "Learn Person init"

    #Starting services
    memory_service  = session.service("ALMemory")

    faces_service = session.service("ALFaceDetection")
    faces_service.setRecognitionConfidenceThreshold(confidenceThreshold)
    faces_service.setRecognitionEnabled(True)
    learnedlist = faces_service.getLearnedFacesList()
    print 'learnedlist=',learnedlist

    people_service = session.service("ALPeoplePerception")

    print "Creating the thread"

    #create a thead that monitors directly the signal
    monitorThread = threading.Thread(target = rhMonitorThread, args = (memory_service,))
    monitorThread.start()

def quit():
    global monitorThread
    print "Learn person quit"
    monitorThread.do_run = False

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("--pip", type=str, default=os.environ['PEPPER_IP'],
                        help="Robot IP address.  On robot or Local Naoqi: use '127.0.0.1'.")
    parser.add_argument("--pport", type=int, default=9559,
                        help="Naoqi port number")
    args = parser.parse_args()
    pip = args.pip
    pport = args.pport

    #Starting application
    try:
        connection_url = "tcp://" + pip + ":" + str(pport)
        print "Connecting to ",    connection_url
        app = qi.Application(["PersonHere", "--qi-url=" + connection_url ])
    except RuntimeError:
        print ("Can't connect to Naoqi at ip \"" + pip + "\" on port " + str(pport) +".\n"
               "Please check your script arguments. Run with -h option for help.")
        sys.exit(1)

    app.start()
    session = app.session
    init(session)

    app.run()

if __name__ == "__main__":
    main()
<commit_msg>Update learnperson.py<commit_after>#! /usr/bin/env python
# -*- encoding: UTF-8 -*-

# DOCUMENTATION
#http://doc.aldebaran.com/2-5/naoqi/peopleperception/alpeopleperception-api.html#PeoplePerception
#http://doc.aldebaran.com/2-5/naoqi/peopleperception/alfacedetection.html

import qi
import argparse
import sys
import os
import time
import threading
import json

from naoqi import ALProxy

from utils import point2world
import conditions
from conditions import set_condition

##PARAMS
confidenceThreshold=0.4 #default =0.4
min_size={'w':0.10 ,'h':0.10}

def writeJson(path,data):

    if not os.path.exists(path):
        os.makedirs(path)

    jsonData=json.dumps(data,indent=2 )

    fh = open(path+'userlearned.json',"w")
    fh.write(jsonData)
    fh.close()

def onFaceDetection(facevalues):

    for faceInfo in facevalues:
        # First Field = Shape info.
        faceShapeInfo = faceInfo[0]
        # Second Field = Extra info .
        faceExtraInfo = faceInfo[1]

        #print 'eyeLeft',round(faceExtraInfo[3][0],2),'eyeRight',round(faceExtraInfo[4][0],2)
        pose={'alpha': round(faceShapeInfo[1],2), 'beta': round(faceShapeInfo[2],2), 'width':round(faceShapeInfo[3],2), 'height': round(faceShapeInfo[4],2)}
        face={'faceid': faceExtraInfo[0], 'user': faceExtraInfo[2], 'confidence':round(faceExtraInfo[1],3), 'pose': pose}

    return face

def rhMonitorThread (memory_service):
    t = threading.currentThread()
    print "learnperson thread started"
    print "waiting for face"
    while getattr(t, "do_run", True):

        facevalues =  memory_service.getData("FaceDetected")
        peoplevisible =  memory_service.getData("PeoplePerception/VisiblePeopleList")

        if facevalues:

            facesarray=facevalues[1][0:len(facevalues[1]) -1]

            #### TODO: Now only valid when detects 1 face
            if len(facesarray) is 1:

                face=onFaceDetection(facesarray)

                print 'pose',face['pose']
                if float(face['pose']['width'])>= float(min_size['w']) and float(face['pose']['height'])>=min_size['h']:

                    #####
                    ##### ID ????
                    print 'Learning face ',face['faceid']
                    learned=faces_service.learnFace('user-'+str(face['faceid']))
                    print 'Face learned =',learned
                    face['user']='user-'+str(face['faceid'])

                    if learned is True:

                        person=None
                        poseinworld=None

                        if len(peoplevisible) is 1:

                            personid=peoplevisible[0]
                            print 'personid= ',personid
                            #try:

        #                    lookingat =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/LookingAtRobotScore")
        #                    gazedirection =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/GazeDirection")
        #                    age =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/AgeProperties")
        #                    gender =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/GenderProperties")
        #                    expression =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/ExpressionProperties")
        #                    smile =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/SmileProperties")
        #                    print 'lookingat=',lookingat
        #                    print 'gazedirection=',gazedirection
        #                    print 'age=',age
        #                    print 'gender=',gender
        #                    print 'expression=',expression
        #                    print 'smile=',smile
        #

                            angles =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/AnglesYawPitch")
                            distance =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/Distance")
                            height =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/RealHeight")
                            shirtcolor =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/ShirtColor")
                            shirtcolorHSV =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/ShirtColorHSV")
                            PositionInRobotFrame =memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/PositionInRobotFrame")
        #                    issitting=memory_service.getData( "PeoplePerception/Person/" +str(personid)+"/IsSitting")
        #
        #
        #                    print 'issitting=',issitting

                            # Write data in json format

                            personinfo={'height': round(height,2), 'shirtcolor': shirtcolor, 'shirtcolorHSV':shirtcolorHSV}

                            w_px, w_py = point2world(memory_service,[PositionInRobotFrame[0],PositionInRobotFrame[1]])
                            poseinworld={'x':w_px,'y': w_py}

                            person={'personid': personid, 'info': personinfo}

                        posetopological={'current_node':'TODO','closest_node':'TODO'}
                        lastlocation= {'world':poseinworld , 'topological':posetopological}

                        user={'face': face,'person': person,'lastlocation':lastlocation}

                        writeJson('./data/',user)

                        print 'EXIT TRUE'

                        set_condition(memory_service,'learnperson','true')

                        quit()

                    elif learned is False:

                        print 'EXIT FALSE'

                        set_condition(memory_service,'learnperson','false')
                        quit()

        time.sleep(0.2)
    print "learnperson thread quit"

def init(session):
    global memory_service
    global monitorThread

    global faces_service
    print "Learn Person init"

    #Starting services
    memory_service  = session.service("ALMemory")

    faces_service = session.service("ALFaceDetection")
    faces_service.setRecognitionConfidenceThreshold(confidenceThreshold)
    faces_service.setRecognitionEnabled(True)
    learnedlist = faces_service.getLearnedFacesList()
    print 'learnedlist=',learnedlist

    people_service = session.service("ALPeoplePerception")

    print "Creating the thread"

    #create a thead that monitors directly the signal
    monitorThread = threading.Thread(target = rhMonitorThread, args = (memory_service,))
    monitorThread.start()

def quit():
    global monitorThread
    print "Learn person quit"
    monitorThread.do_run = False

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("--pip", type=str, default=os.environ['PEPPER_IP'],
                        help="Robot IP address.  On robot or Local Naoqi: use '127.0.0.1'.")
    parser.add_argument("--pport", type=int, default=9559,
                        help="Naoqi port number")
    args = parser.parse_args()
    pip = args.pip
    pport = args.pport

    #Starting application
    try:
        connection_url = "tcp://" + pip + ":" + str(pport)
        print "Connecting to ",    connection_url
        app = qi.Application(["PersonHere", "--qi-url=" + connection_url ])
    except RuntimeError:
        print ("Can't connect to Naoqi at ip \"" + pip + "\" on port " + str(pport) +".\n"
               "Please check your script arguments. Run with -h option for help.")
        sys.exit(1)

    app.start()
    session = app.session
    init(session)

    app.run()

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>import numpy as np
from scipy.spatial.distance import cdist

def distance(a,b):
    # catch nested (ligands + atomic number groups) coordinates of shape (m-ligands, k-groups, n-atoms, m-dim)
    if len(a.shape) == 3 and len(b.shape) == 4:
        return np.sqrt((b[:,np.newaxis,:,np.newaxis,:,0] - a[:, np.newaxis, :, 0, np.newaxis])**2 + (b[:,np.newaxis,:, np.newaxis,:,1] - a[:, np.newaxis, :, 1, np.newaxis])**2 + (b[:,np.newaxis,:, np.newaxis, :,2] - a[:, np.newaxis, :, 2, np.newaxis])**2)
    # catch nested coordinates of shape (k-groups, n-atoms, m-dim)
    if len(a.shape) == 3 and len(b.shape) == 3:
        return np.sqrt((b[:,np.newaxis,:,0] - a[:, np.newaxis, :, 0, np.newaxis])**2 + (b[:, np.newaxis,:,1] - a[:, np.newaxis, :, 1, np.newaxis])**2 + (b[:, np.newaxis, :,2] - a[:, np.newaxis, :, 2, np.newaxis])**2)
    # otherwise simple group of atom coordinates of shape (n-atoms, m-dim)
    elif len(a.shape) == 2 and len(b.shape) == 2:
        # numpy fastest solution (slower than scipy.spatial.distance.cdist)
        #return np.sqrt(((b[:,0] - a[:,0, np.newaxis])**2 + (b[:,1] - a[:,1, np.newaxis])**2 + (b[:,2] - a[:,2, np.newaxis])**2))
        # scipy
        return cdist(a,b)
    else:
        raise ValueError('Unsuported shape of arrays (%i, %i)' % (len(a.shape), len(b.shape)))

class Molecule:
    """ Class which holds dictionaries aiding descriptors calculation """
    def __init__(self, mol):
        """ mol: pybel molecule """
        self.m = mol

    def coordinate_dict(self, atomic_nums):
        """
        Returns vector of atomic coordinates of atoms with atomic number = atomic_num in given molecule

        atomic nums: array of atomic numbers to compute the dictionary
        """
        mol_atoms = {}
        for a in atomic_nums:
            mol_atoms[a] = []
        for atom in self.m:
            if atom.atomicnum in atomic_nums:
                   mol_atoms[atom.atomicnum].append(atom.coords)
        for a in atomic_nums:
            mol_atoms[a] = np.array(mol_atoms[a])

#        # define dictionaries
#        mol_atoms = np.empty((len(atomic_nums), len(self.m.atoms),3))
#        mol_atoms.fill(np.nan)
#        # iterate through atoms
#        for atom in self.m:
#            if atom.atomicnum in atomic_nums:
#                    mol_atoms[atomic_nums.index(atom.atomicnum)][atom.idx-1] = atom.coords
        return mol_atoms

# DESCRIPTORS

def close_contact(mol1_atoms, mol2_atoms, cutoff):
    """
    Builds descriptor from two dictionaries (protein and ligand) generated by mol_dict()
    """

    # !!! FIX: change to pure numpy implementation, but distance must support nested atom coordinates
    #return np.sum(distance(mol2_atoms, mol1_atoms) < cutoff, axis=-1)

    # !!! FIX 2: check if scipy.spatial.KDTree performs well, otherwise leave numpy alone

    desc = []
    for mol2_a in sorted(mol2_atoms.keys()):
        for mol1_a in sorted(mol1_atoms.keys()):
            if len(mol1_atoms[mol1_a]) > 0 and len(mol2_atoms[mol2_a]):
                desc.append(np.sum(cdist(mol1_atoms[mol1_a], mol2_atoms[mol2_a]) < cutoff))
            else:
                desc.append(0)

    return np.array(desc)

    #return np.sum(distance(mol1_atoms, mol2_atoms) < cutoff, axis=(-1,-2))

<commit_msg>Code cleanup<commit_after>import numpy as np
from scipy.spatial.distance import cdist

def distance(a,b):
    # catch nested (ligands + atomic number groups) coordinates of shape (m-ligands, k-groups, n-atoms, m-dim)
    if len(a.shape) == 3 and len(b.shape) == 4:
        return np.sqrt((b[:,np.newaxis,:,np.newaxis,:,0] - a[:, np.newaxis, :, 0, np.newaxis])**2 + (b[:,np.newaxis,:, np.newaxis,:,1] - a[:, np.newaxis, :, 1, np.newaxis])**2 + (b[:,np.newaxis,:, np.newaxis, :,2] - a[:, np.newaxis, :, 2, np.newaxis])**2)
    # catch nested coordinates of shape (k-groups, n-atoms, m-dim)
    if len(a.shape) == 3 and len(b.shape) == 3:
        return np.sqrt((b[:,np.newaxis,:,0] - a[:, np.newaxis, :, 0, np.newaxis])**2 + (b[:, np.newaxis,:,1] - a[:, np.newaxis, :, 1, np.newaxis])**2 + (b[:, np.newaxis, :,2] - a[:, np.newaxis, :, 2, np.newaxis])**2)
    # otherwise simple group of atom coordinates of shape (n-atoms, m-dim)
    elif len(a.shape) == 2 and len(b.shape) == 2:
        # numpy fastest solution (slower than scipy.spatial.distance.cdist)
        #return np.sqrt(((b[:,0] - a[:,0, np.newaxis])**2 + (b[:,1] - a[:,1, np.newaxis])**2 + (b[:,2] - a[:,2, np.newaxis])**2))
        # scipy
        return cdist(a,b)
    else:
        raise ValueError('Unsuported shape of arrays (%i, %i)' % (len(a.shape), len(b.shape)))

class Molecule:
    """ Class which holds dictionaries aiding descriptors calculation """
    def __init__(self, mol):
        """ mol: pybel molecule """
        self.m = mol

    def coordinate_dict(self, atomic_nums):
        """
        Returns vector of atomic coordinates of atoms with atomic number = atomic_num in given molecule

        atomic nums: array of atomic numbers to compute the dictionary
        """
        mol_atoms = {}
        for a in atomic_nums:
            mol_atoms[a] = []
        for atom in self.m:
            if atom.atomicnum in atomic_nums:
                   mol_atoms[atom.atomicnum].append(atom.coords)
        for a in atomic_nums:
            mol_atoms[a] = np.array(mol_atoms[a])
        return mol_atoms

# DESCRIPTORS

def close_contact(mol1_atoms, mol2_atoms, cutoff):
    """
    Builds descriptor from two dictionaries (protein and ligand) generated by mol_dict()
    """
    desc = []
    for mol2_a in sorted(mol2_atoms.keys()):
        for mol1_a in sorted(mol1_atoms.keys()):
            if len(mol1_atoms[mol1_a]) > 0 and len(mol2_atoms[mol2_a]):
                desc.append(np.sum(cdist(mol1_atoms[mol1_a], mol2_atoms[mol2_a]) < cutoff))
            else:
                desc.append(0)

    return np.array(desc)
<|endoftext|>