<commit_before>#  Copyright 2008-2011 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from runonfailure import RunOnFailure

class NoFlexApplicationSelected(Exception):
    pass

class Flex(RunOnFailure):

    def select_flex_application(self, locator, alias=None):
        """Select flex application to work with.

        All further Flex-keywords will operate on the selected application.

        `locator` is the value `id` or `name` attribute of the movie in HTML.

        Return index if this application that can be used with `Switch Flex
        Application`
        """
        self.page_should_contain_element(locator)
        # It seems that selenium timeout is used instead of given timeout.
        self._selenium.do_command("waitForFlexReady", [locator, 'timeout'])
        return self._flex_apps.register(locator, alias)

    def wait_for_flex_element(self, locator, timeout=None):
        error = "Element '%s' did not appear in %%(timeout)s" % locator
        self._wait_for_flex_element(self._create_flex_pilot_locator(locator),
                                    timeout, error)

    def _wait_for_flex_element(self, locator, timeout, error):
        self._wait_until(lambda: self._element_exists(locator),
                         error, timeout)

    def _element_exists(self, locator):
        try:
            self._flex_command('flexAssertDisplayObject', locator)
        except NoFlexApplicationSelected:
            raise
        except Exception:
            return False
        else:
            return True

    def switch_flex_application(self, index_or_alias):
        """Switches between active flex applications  using index or alias.

        `index_or_alias` is got from `Select Flex Application` and alias can
        be given to it.
        """
        self._flex_apps.switch(index_or_alias)

    def flex_element_should_exist(self, locator):
        """Verifies that Flex component can be found by `locator`.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexAssertDisplayObject',
                           self._flex_locator(locator))

    def flex_element_should_not_exist(self, locator):
        """Verifies that Flex component cannot be found by `locator`.

        See `introduction` about rules for locating Flex elements.
        """
        try:
            self.flex_element_should_exist(locator)
        except Exception:
            pass
        else:
            raise AssertionError("Element '%s' exists" % locator)

    def click_flex_element(self, locator):
        """Click the Flex element found by `locator`.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexClick', self._flex_locator(locator))

    def flex_element_text_should_be(self, locator, expected):
        """Verifies the value of the text field found by `locator` is `expected`.

        See `introduction` about rules for locating Flex elements.
        """
        locator = self._flex_locator(locator)
        self._flex_command('flexAssertText',
                           '%s,validator=%s' % (locator, expected))

    def flex_element_property_should_be(self, locator, name, expected):
        """Verifies property value of an element found by `locator`.

        `name` is the name of the property and `expected` is the expected
        value.

        See `introduction` about rules for locating Flex elements.
        """
        validator = '%s|%s' % (name, expected)
        locator = self._flex_locator(locator)
        self._flex_command('flexAssertProperty', '%s,validator=%s'
                           % (locator, validator))

    def input_text_into_flex_element(self, locator, text):
        """Input `value` in the text field found by `locator`.

        See `introduction` about rules for locating Flex elements.
        """
        locator = self._flex_locator(locator)
        self._flex_command('flexType', '%s, text=%s' % (locator, text))

    def select_from_flex_element(self, locator, value):
        """Select `value` from Flex element found by `locator`.

        `value` may be either an index, visible text, or associated data of
        the item to be selected.

        Examples:
        | Select From Flex Element | Text | # Select by visible text |
        | Select From Flex Element | index=1 | # Select by index |
        | Select From Flex Element | data=someData | # Select by associated data |

        *NOTE* This keyword generates mx.events.ListEvent.CHANGE event, which
        means that event handlers associated with opening or closing a drop down
        menu will not be executed.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexSelect',
                           '%s, %s' %  (self._flex_locator(locator),
                                        self._choice_locator(value)))

    def _choice_locator(self, given_locator):
        for prefix in ['index=', 'label=', 'text=', 'data=', 'value=']:
            if given_locator.startswith(prefix):
                return given_locator
        return 'label=' + given_locator

    def _flex_locator(self, locator):
        error = "Element '%s' does not exist" % locator
        locator = self._create_flex_pilot_locator(locator)
        self._wait_for_flex_element(locator, '0.5s', error)
        return locator

    def _create_flex_pilot_locator(self, locator):
        locator = locator.strip()
        # prefixes gotten from org/flex_pilot/FPLocator.as
        for prefix in ['id=', 'automationName=', 'name=', 'chain=', 'label=',
                       'htmlText=']:
            if locator.startswith(prefix):
                return locator
        return 'id=%s' % locator

    def _flex_command(self, command, options):
        # TODO: Howto handle commas in option values??
        app = self._flex_apps.current
        if not app:
            raise NoFlexApplicationSelected
        self._debug("Executing command '%s' for application '%s' with options '%s'"
                    % (command, app, options))
        self._selenium.do_command(command, [app, options])

<commit_msg>Changed the way flex commands are created<commit_after>#  Copyright 2008-2011 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from runonfailure import RunOnFailure

class NoFlexApplicationSelected(Exception):
    pass

class Flex(RunOnFailure):

    def select_flex_application(self, locator, alias=None):
        """Select flex application to work with.

        All further Flex-keywords will operate on the selected application.

        `locator` is the value `id` or `name` attribute of the movie in HTML.

        Return index if this application that can be used with `Switch Flex
        Application`
        """
        self.page_should_contain_element(locator)
        # It seems that selenium timeout is used instead of given timeout.
        self._selenium.do_command("waitForFlexReady", [locator, 'timeout'])
        return self._flex_apps.register(locator, alias)

    def wait_for_flex_element(self, locator, timeout=None):
        error = "Element '%s' did not appear in %%(timeout)s" % locator
        self._wait_for_flex_element(self._flex_locator(locator),
                                    timeout, error)

    def _wait_for_flex_element(self, locator, timeout, error):
        self._wait_until(lambda: self._element_exists(locator),
                         error, timeout)

    def _element_exists(self, locator):
        try:
            self._do_flex_command('flexAssertDisplayObject', locator)
        except NoFlexApplicationSelected:
            raise
        except Exception:
            return False
        else:
            return True

    def switch_flex_application(self, index_or_alias):
        """Switches between active flex applications  using index or alias.

        `index_or_alias` is got from `Select Flex Application` and alias can
        be given to it.
        """
        self._flex_apps.switch(index_or_alias)

    def flex_element_should_exist(self, locator):
        """Verifies that Flex component can be found by `locator`.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexAssertDisplayObject', locator)

    def flex_element_should_not_exist(self, locator):
        """Verifies that Flex component cannot be found by `locator`.

        See `introduction` about rules for locating Flex elements.
        """
        try:
            self.flex_element_should_exist(locator)
        except Exception:
            pass
        else:
            raise AssertionError("Element '%s' exists" % locator)

    def click_flex_element(self, locator):
        """Click the Flex element found by `locator`.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexClick', locator)

    def flex_element_text_should_be(self, locator, expected):
        """Verifies the value of the text field found by `locator` is `expected`.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexAssertText', locator, 'validator='+expected)

    def flex_element_property_should_be(self, locator, name, expected):
        """Verifies property value of an element found by `locator`.

        `name` is the name of the property and `expected` is the expected
        value.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexAssertProperty', locator,
                           'validator=%s|%s' % (name, expected))

    def input_text_into_flex_element(self, locator, text):
        """Input `value` in the text field found by `locator`.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexType', locator, 'text='+text)

    def select_from_flex_element(self, locator, value):
        """Select `value` from Flex element found by `locator`.

        `value` may be either an index, visible text, or associated data of
        the item to be selected.

        Examples:
        | Select From Flex Element | Text | # Select by visible text |
        | Select From Flex Element | index=1 | # Select by index |
        | Select From Flex Element | data=someData | # Select by associated data |

        *NOTE* This keyword generates mx.events.ListEvent.CHANGE event, which
        means that event handlers associated with opening or closing a drop down
        menu will not be executed.

        See `introduction` about rules for locating Flex elements.
        """
        self._flex_command('flexSelect', locator, self._choice_locator(value))

    def _choice_locator(self, locator):
        for prefix in ['index=', 'label=', 'text=', 'data=', 'value=']:
            if locator.startswith(prefix):
                return locator
        return 'label='+locator

    def _flex_command(self, command, locator, option=None):
        locator = self._flex_locator(locator)
        self._verify_flex_element_exists(locator)
        self._do_flex_command(command, locator, option)

    def _flex_locator(self, locator):
        locator = locator.strip()
        # prefixes gotten from org/flex_pilot/FPLocator.as
        for prefix in ['id=', 'automationName=', 'name=', 'chain=', 'label=',
                       'htmlText=']:
            if locator.startswith(prefix):
                return locator
        return 'id=%s' % locator

    def _verify_flex_element_exists(self, locator):
        if not self._flex_apps.current:
            raise NoFlexApplicationSelected
        error = "Element '%s' does not exist" % locator
        self._wait_for_flex_element(locator, '0.5s', error)

    def _do_flex_command(self, command, *options):
        app = self._flex_apps.current
        opts = ', '.join([o for o in options if o is not None])
        self._debug("Executing command '%s' for application '%s' with options '%s'"
                    % (command, app, opts))
        self._selenium.do_command(command, [app, opts])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright (c) 2019, Frappe Technologies and Contributors
# See license.txt
from __future__ import unicode_literals

import frappe
import unittest
import requests
from frappe.utils import get_site_url

scripts = [
	dict(
		name='test_todo',
		script_type = 'DocType Event',
		doctype_event = 'Before Insert',
		reference_doctype = 'ToDo',
		script = '''
if "test" in doc.description:
	doc.status = 'Closed'
'''
	),
	dict(
		name='test_todo_validate',
		script_type = 'DocType Event',
		doctype_event = 'Before Insert',
		reference_doctype = 'ToDo',
		script = '''
if "validate" in doc.description:
	raise frappe.ValidationError
'''
	),
	dict(
		name='test_api',
		script_type = 'API',
		api_method = 'test_server_script',
		allow_guest = 1,
		script = '''
frappe.response['message'] = 'hello'
'''
	),
	dict(
		name='test_return_value',
		script_type = 'API',
		api_method = 'test_return_value',
		allow_guest = 1,
		script = '''
frappe.flags = 'hello'
'''
	)
]
class TestServerScript(unittest.TestCase):
	@classmethod
	def setUpClass(cls):
		frappe.db.commit()
		frappe.db.sql('truncate `tabServer Script`')
		frappe.get_doc('User', 'Administrator').add_roles('Script Manager')
		for script in scripts:
			script_doc = frappe.get_doc(doctype ='Server Script')
			script_doc.update(script)
			script_doc.insert()

		frappe.db.commit()

	@classmethod
	def tearDownClass(cls):
		frappe.db.commit()
		frappe.db.sql('truncate `tabServer Script`')

	def setUp(self):
		frappe.cache().delete_value('server_script_map')

	def test_doctype_event(self):
		todo = frappe.get_doc(dict(doctype='ToDo', description='hello')).insert()
		self.assertEqual(todo.status, 'Open')

		todo = frappe.get_doc(dict(doctype='ToDo', description='test todo')).insert()
		self.assertEqual(todo.status, 'Closed')

		self.assertRaises(frappe.ValidationError, frappe.get_doc(dict(doctype='ToDo', description='validate me')).insert)

	def test_api(self):
		response = requests.post(get_site_url(frappe.local.site) + "/api/method/test_server_script")
		self.assertEqual(response.status_code, 200)
		self.assertEqual("hello", response.json()["message"])

	def test_api_return(self):
		self.assertEqual(frappe.get_doc('Server Script', 'test_return_value').execute_method(), 'hello')
<commit_msg>test: Add tests for throwing AttributeError if method not found<commit_after># -*- coding: utf-8 -*-
# Copyright (c) 2019, Frappe Technologies and Contributors
# See license.txt
from __future__ import unicode_literals

import frappe
import unittest
import requests
from frappe.utils import get_site_url

scripts = [
	dict(
		name='test_todo',
		script_type = 'DocType Event',
		doctype_event = 'Before Insert',
		reference_doctype = 'ToDo',
		script = '''
if "test" in doc.description:
	doc.status = 'Closed'
'''
	),
	dict(
		name='test_todo_validate',
		script_type = 'DocType Event',
		doctype_event = 'Before Insert',
		reference_doctype = 'ToDo',
		script = '''
if "validate" in doc.description:
	raise frappe.ValidationError
'''
	),
	dict(
		name='test_api',
		script_type = 'API',
		api_method = 'test_server_script',
		allow_guest = 1,
		script = '''
frappe.response['message'] = 'hello'
'''
	),
	dict(
		name='test_return_value',
		script_type = 'API',
		api_method = 'test_return_value',
		allow_guest = 1,
		script = '''
frappe.flags = 'hello'
'''
	),
	dict(
		name='test_invalid_namespace_method',
		script_type = 'DocType Event',
		doctype_event = 'Before Insert',
		reference_doctype = 'Note',
		script = '''
frappe.method_that_doesnt_exist("do some magic")
'''
	)
]
class TestServerScript(unittest.TestCase):
	@classmethod
	def setUpClass(cls):
		frappe.db.commit()
		frappe.db.sql('truncate `tabServer Script`')
		frappe.get_doc('User', 'Administrator').add_roles('Script Manager')
		for script in scripts:
			script_doc = frappe.get_doc(doctype ='Server Script')
			script_doc.update(script)
			script_doc.insert()

		frappe.db.commit()

	@classmethod
	def tearDownClass(cls):
		frappe.db.commit()
		frappe.db.sql('truncate `tabServer Script`')

	def setUp(self):
		frappe.cache().delete_value('server_script_map')

	def test_doctype_event(self):
		todo = frappe.get_doc(dict(doctype='ToDo', description='hello')).insert()
		self.assertEqual(todo.status, 'Open')

		todo = frappe.get_doc(dict(doctype='ToDo', description='test todo')).insert()
		self.assertEqual(todo.status, 'Closed')

		self.assertRaises(frappe.ValidationError, frappe.get_doc(dict(doctype='ToDo', description='validate me')).insert)

	def test_api(self):
		response = requests.post(get_site_url(frappe.local.site) + "/api/method/test_server_script")
		self.assertEqual(response.status_code, 200)
		self.assertEqual("hello", response.json()["message"])

	def test_api_return(self):
		self.assertEqual(frappe.get_doc('Server Script', 'test_return_value').execute_method(), 'hello')

	def test_attribute_error(self):
		"""Raise AttributeError if method not found in Namespace"""
		note = frappe.get_doc({"doctype": "Note", "title": "Test Note: Server Script"})
		self.assertRaises(AttributeError, note.insert)
<|endoftext|>


<commit_before>from aimacode.logic import PropKB
from aimacode.planning import Action
from aimacode.search import (
    Node, Problem,
)
from aimacode.utils import expr
from lp_utils import (
    FluentState, encode_state, decode_state,
)
from my_planning_graph import PlanningGraph

class AirCargoProblem(Problem):
    def __init__(self, cargos, planes, airports, initial: FluentState, goal: list):
        """

        :param cargos: list of str
            cargos in the problem
        :param planes: list of str
            planes in the problem
        :param airports: list of str
            airports in the problem
        :param initial: FluentState object
            positive and negative literal fluents (as expr) describing initial state
        :param goal: list of expr
            literal fluents required for goal test
        """
        self.state_map = initial.pos + initial.neg
        self.initial_state_TF = encode_state(initial, self.state_map)
        Problem.__init__(self, self.initial_state_TF, goal=goal)
        self.cargos = cargos
        self.planes = planes
        self.airports = airports
        self.actions_list = self.get_actions()

    def get_actions(self):
        '''
        This method creates concrete actions (no variables) for all actions in the problem
        domain action schema and turns them into complete Action objects as defined in the
        aimacode.planning module. It is computationally expensive to call this method directly;
        however, it is called in the constructor and the results cached in the `actions_list` property.

        Returns:
        ----------
        list<Action>
            list of Action objects
        '''

        # TODO create concrete Action objects based on the domain action schema for: Load, Unload, and Fly
        # concrete actions definition: specific literal action that does not include variables as with the schema
        # for example, the action schema 'Load(c, p, a)' can represent the concrete actions 'Load(C1, P1, SFO)'
        # or 'Load(C2, P2, JFK)'.  The actions for the planning problem must be concrete because the problems in
        # forward search and Planning Graphs must use Propositional Logic

        def load_actions():
            '''Create all concrete Load actions and return a list

            :return: list of Action objects
            '''
            loads = []
            for p in self.planes:
                for c in self.cargos:
                    for a in self.airports:
                        precond_pos = [
                            expr('At({}, {})'.format(c, a)),
                            expr('At({}, {})'.format(p, a))
                        ]
                        precond_neg = []
                        effect_add = [expr('In({}, {})'.format(c, p))]
                        effect_rem = [expr('At({}, {})'.format(c, a))]
                        load = Action(expr('Load({}, {}, {})'.format(c, p, a)),
                                      [precond_pos, precond_neg],
                                      [effect_add, effect_rem])
                        loads.append(load)
            return loads

        def unload_actions():
            '''Create all concrete Unload actions and return a list

            :return: list of Action objects
            '''
            unloads = []
            for p in self.planes:
                for c in self.cargos:
                    for a in self.airports:
                        precond_pos = [
                            expr('In({}, {})'.format(c, p)),
                            expr('At({}, {})'.format(p, a))
                        ]
                        precond_neg = []
                        effect_add = [expr('At({}, {})'.format(c, a))]
                        effect_rem = [expr('In({}, {})'.format(c, p))]
                        unload = Action(expr('Unload({}, {}, {})'.format(c, p, a)),
                                        [precond_pos, precond_neg],
                                        [effect_add, effect_rem])
                        unloads.append(unload)
            return unloads

        def fly_actions():
            '''Create all concrete Fly actions and return a list

            :return: list of Action objects
            '''
            flys = []
            for fr in self.airports:
                for to in self.airports:
                    if fr != to:
                        for p in self.planes:
                            precond_pos = [expr('At({}, {})'.format(p, fr)),
                                           ]
                            precond_neg = []
                            effect_add = [expr('At({}, {})'.format(p, to))]
                            effect_rem = [expr('At({}, {})'.format(p, fr))]
                            fly = Action(expr('Fly({}, {}, {})'.format(p, fr, to)),
                                         [precond_pos, precond_neg],
                                         [effect_add, effect_rem])
                            flys.append(fly)
            return flys

        return load_actions() + unload_actions() + fly_actions()

    def actions(self, state: str) -> list:
        """ Return the actions that can be executed in the given state.

        :param state: str
            state represented as T/F string of mapped fluents (state variables)
            e.g. 'FTTTFF'
        :return: list of Action objects
        """
        possible_actions = []
        kb = PropKB()
        kb.tell(decode_state(state, self.state_map).pos_sentence())
        for action in self.actions_list:
            is_possible = True
            for clause in action.precond_pos:
                if clause not in kb.clauses:
                    is_possible = False
            for clause in action.precond_neg:
                if clause in kb.clauses:
                    is_possible = False
            if is_possible:
                possible_actions.append(action)
        return possible_actions

    def result(self, state: str, action: Action):
        """ Return the state that results from executing the given
        action in the given state. The action must be one of
        self.actions(state).

        :param state: state entering node
        :param action: Action applied
        :return: resulting state after action
        """
        # TODO implement
        new_state = FluentState([], [])
        old_state = decode_state(state, self.state_map)
        for fluent in old_state.pos:
            if fluent not in action.effect_rem:
                new_state.pos.append(fluent)
        for fluent in action.effect_add:
            if fluent not in new_state.pos:
                new_state.pos.append(fluent)
        for fluent in old_state.neg:
            if fluent not in action.effect_add:
                new_state.neg.append(fluent)
        for fluent in action.effect_rem:
            if fluent not in new_state.neg:
                new_state.neg.append(fluent)
        return encode_state(new_state, self.state_map)

    def goal_test(self, state: str) -> bool:
        """ Test the state to see if goal is reached

        :param state: str representing state
        :return: bool
        """
        kb = PropKB()
        kb.tell(decode_state(state, self.state_map).pos_sentence())
        for clause in self.goal:
            if clause not in kb.clauses:
                return False
        return True

    def h_1(self, node: Node):
        # note that this is not a true heuristic
        h_const = 1
        return h_const

    def h_pg_levelsum(self, node: Node):
        '''
        This heuristic uses a planning graph representation of the problem
        state space to estimate the sum of all actions that must be carried
        out from the current state in order to satisfy each individual goal
        condition.
        '''
        # requires implemented PlanningGraph class
        pg = PlanningGraph(self, node.state)
        pg_levelsum = pg.h_levelsum()
        return pg_levelsum

    def h_ignore_preconditions(self, node: Node):
        '''
        This heuristic estimates the minimum number of actions that must be
        carried out from the current state in order to satisfy all of the goal
        conditions by ignoring the preconditions required for an action to be
        executed.
        '''
        # TODO implement (see Russell-Norvig Ed-3 10.2.3  or Russell-Norvig Ed-2 11.2)
        count = 0
        return count

def air_cargo_p1() -> AirCargoProblem:
    cargos = ['C1', 'C2']
    planes = ['P1', 'P2']
    airports = ['JFK', 'SFO']
    pos = [expr('At(C1, SFO)'),
           expr('At(C2, JFK)'),
           expr('At(P1, SFO)'),
           expr('At(P2, JFK)'),
           ]
    neg = [expr('At(C2, SFO)'),
           expr('In(C2, P1)'),
           expr('In(C2, P2)'),
           expr('At(C1, JFK)'),
           expr('In(C1, P1)'),
           expr('In(C1, P2)'),
           expr('At(P1, JFK)'),
           expr('At(P2, SFO)'),
           ]
    init = FluentState(pos, neg)
    goal = [expr('At(C1, JFK)'),
            expr('At(C2, SFO)'),
            ]
    return AirCargoProblem(cargos, planes, airports, init, goal)

def air_cargo_p2() -> AirCargoProblem:
    cargos = ['C1', 'C2', 'C3']
    planes = ['P1', 'P2', 'P3']
    airports = ['JFK', 'SFO', 'ATL']
    pos = [expr('At(C1, SFO)'),
           expr('At(C2, JFK)'),
           expr('At(C3, ATL)'),
           expr('At(P1, SFO)'),
           expr('At(P2, JFK)'),
           expr('At(P3, ATL)'),
           ]
    neg = [expr('At(C1, JFK)'),
           expr('At(C1, ATL)'),
           expr('In(C1, P1)'),
           expr('In(C1, P2)'),
           expr('In(C1, P3)'),
           expr('At(C2, SFO)'),
           expr('At(C2, ATL)'),
           expr('In(C2, P1)'),
           expr('In(C2, P2)'),
           expr('In(C2, P3)'),
           expr('At(C3, SFO)'),
           expr('At(C3, JFK)'),
           expr('In(C3, P1)'),
           expr('In(C3, P2)'),
           expr('In(C3, P3)'),
           expr('At(P1, JFK)'),
           expr('At(P1, ATL)'),
           expr('At(P2, SFO)'),
           expr('At(P2, ATL)'),
           expr('At(P3, SFO)'),
           expr('At(P3, JFK)'),
           ]
    init = FluentState(pos, neg)
    goal = [expr('At(C1, JFK)'),
            expr('At(C2, SFO)'),
            expr('At(C3, SFO)'),
            ]
    return AirCargoProblem(cargos, planes, airports, init, goal)

def air_cargo_p3() -> AirCargoProblem:
    cargos = ['C1', 'C2', 'C3', 'C4']
    planes = ['P1', 'P2', 'P3', 'P4']
    airports = ['JFK', 'SFO', 'ATL', 'ORD']
    pos = [expr('At(C1, SFO)'),
           expr('At(C2, JFK)'),
           expr('At(C3, ATL)'),
           expr('At(C4, ORD)'),
           expr('At(P1, SFO)'),
           expr('At(P2, JFK)'),
           ]
    neg = [expr('At(C1, JFK)'),
           expr('At(C1, ATL)'),
           expr('At(C1, ORD)'),
           expr('In(C1, P1)'),
           expr('In(C1, P2)'),
           expr('At(C2, SFO)'),
           expr('At(C2, ATL)'),
           expr('At(C2, ORD)'),
           expr('In(C2, P1)'),
           expr('In(C2, P2)'),
           expr('At(C3, SFO)'),
           expr('At(C3, JFK)'),
           expr('At(C3, ORD)'),
           expr('In(C3, P1)'),
           expr('In(C3, P2)'),
           expr('At(C4, SFO)'),
           expr('At(C4, JFK)'),
           expr('At(C4, ATL)'),
           expr('In(C4, P1)'),
           expr('In(C4, P2)'),
           expr('At(P1, JFK)'),
           expr('At(P1, ATL)'),
           expr('At(P1, ORD)'),
           expr('At(P2, SFO)'),
           expr('At(P2, ATL)'),
           expr('At(P2, ORD)'),
           ]
    init = FluentState(pos, neg)
    goal = [expr('At(C1, JFK)'),
            expr('At(C2, SFO)'),
            expr('At(C3, JFK)'),
            expr('At(C4, SFO)'),
            ]
    return AirCargoProblem(cargos, planes, airports, init, goal)
<commit_msg>added heuristic<commit_after>from aimacode.logic import PropKB
from aimacode.planning import Action
from aimacode.search import (
    Node, Problem,
)
from aimacode.utils import expr
from lp_utils import (
    FluentState, encode_state, decode_state,
)
from my_planning_graph import PlanningGraph

class AirCargoProblem(Problem):
    def __init__(self, cargos, planes, airports, initial: FluentState, goal: list):
        """

        :param cargos: list of str
            cargos in the problem
        :param planes: list of str
            planes in the problem
        :param airports: list of str
            airports in the problem
        :param initial: FluentState object
            positive and negative literal fluents (as expr) describing initial state
        :param goal: list of expr
            literal fluents required for goal test
        """
        self.state_map = initial.pos + initial.neg
        self.initial_state_TF = encode_state(initial, self.state_map)
        Problem.__init__(self, self.initial_state_TF, goal=goal)
        self.cargos = cargos
        self.planes = planes
        self.airports = airports
        self.actions_list = self.get_actions()

    def get_actions(self):
        '''
        This method creates concrete actions (no variables) for all actions in the problem
        domain action schema and turns them into complete Action objects as defined in the
        aimacode.planning module. It is computationally expensive to call this method directly;
        however, it is called in the constructor and the results cached in the `actions_list` property.

        Returns:
        ----------
        list<Action>
            list of Action objects
        '''

        # TODO create concrete Action objects based on the domain action schema for: Load, Unload, and Fly
        # concrete actions definition: specific literal action that does not include variables as with the schema
        # for example, the action schema 'Load(c, p, a)' can represent the concrete actions 'Load(C1, P1, SFO)'
        # or 'Load(C2, P2, JFK)'.  The actions for the planning problem must be concrete because the problems in
        # forward search and Planning Graphs must use Propositional Logic

        def load_actions():
            '''Create all concrete Load actions and return a list

            :return: list of Action objects
            '''
            loads = []
            for p in self.planes:
                for c in self.cargos:
                    for a in self.airports:
                        precond_pos = [
                            expr('At({}, {})'.format(c, a)),
                            expr('At({}, {})'.format(p, a))
                        ]
                        precond_neg = []
                        effect_add = [expr('In({}, {})'.format(c, p))]
                        effect_rem = [expr('At({}, {})'.format(c, a))]
                        load = Action(expr('Load({}, {}, {})'.format(c, p, a)),
                                      [precond_pos, precond_neg],
                                      [effect_add, effect_rem])
                        loads.append(load)
            return loads

        def unload_actions():
            '''Create all concrete Unload actions and return a list

            :return: list of Action objects
            '''
            unloads = []
            for p in self.planes:
                for c in self.cargos:
                    for a in self.airports:
                        precond_pos = [
                            expr('In({}, {})'.format(c, p)),
                            expr('At({}, {})'.format(p, a))
                        ]
                        precond_neg = []
                        effect_add = [expr('At({}, {})'.format(c, a))]
                        effect_rem = [expr('In({}, {})'.format(c, p))]
                        unload = Action(expr('Unload({}, {}, {})'.format(c, p, a)),
                                        [precond_pos, precond_neg],
                                        [effect_add, effect_rem])
                        unloads.append(unload)
            return unloads

        def fly_actions():
            '''Create all concrete Fly actions and return a list

            :return: list of Action objects
            '''
            flys = []
            for fr in self.airports:
                for to in self.airports:
                    if fr != to:
                        for p in self.planes:
                            precond_pos = [expr('At({}, {})'.format(p, fr)),
                                           ]
                            precond_neg = []
                            effect_add = [expr('At({}, {})'.format(p, to))]
                            effect_rem = [expr('At({}, {})'.format(p, fr))]
                            fly = Action(expr('Fly({}, {}, {})'.format(p, fr, to)),
                                         [precond_pos, precond_neg],
                                         [effect_add, effect_rem])
                            flys.append(fly)
            return flys

        return load_actions() + unload_actions() + fly_actions()

    def actions(self, state: str) -> list:
        """ Return the actions that can be executed in the given state.

        :param state: str
            state represented as T/F string of mapped fluents (state variables)
            e.g. 'FTTTFF'
        :return: list of Action objects
        """
        possible_actions = []
        kb = PropKB()
        kb.tell(decode_state(state, self.state_map).pos_sentence())
        for action in self.actions_list:
            is_possible = True
            for clause in action.precond_pos:
                if clause not in kb.clauses:
                    is_possible = False
            for clause in action.precond_neg:
                if clause in kb.clauses:
                    is_possible = False
            if is_possible:
                possible_actions.append(action)
        return possible_actions

    def result(self, state: str, action: Action):
        """ Return the state that results from executing the given
        action in the given state. The action must be one of
        self.actions(state).

        :param state: state entering node
        :param action: Action applied
        :return: resulting state after action
        """
        # TODO implement
        new_state = FluentState([], [])
        old_state = decode_state(state, self.state_map)
        for fluent in old_state.pos:
            if fluent not in action.effect_rem:
                new_state.pos.append(fluent)
        for fluent in action.effect_add:
            if fluent not in new_state.pos:
                new_state.pos.append(fluent)
        for fluent in old_state.neg:
            if fluent not in action.effect_add:
                new_state.neg.append(fluent)
        for fluent in action.effect_rem:
            if fluent not in new_state.neg:
                new_state.neg.append(fluent)
        return encode_state(new_state, self.state_map)

    def goal_test(self, state: str) -> bool:
        """ Test the state to see if goal is reached

        :param state: str representing state
        :return: bool
        """
        kb = PropKB()
        kb.tell(decode_state(state, self.state_map).pos_sentence())
        for clause in self.goal:
            if clause not in kb.clauses:
                return False
        return True

    def h_1(self, node: Node):
        # note that this is not a true heuristic
        h_const = 1
        return h_const

    def h_pg_levelsum(self, node: Node):
        '''
        This heuristic uses a planning graph representation of the problem
        state space to estimate the sum of all actions that must be carried
        out from the current state in order to satisfy each individual goal
        condition.
        '''
        # requires implemented PlanningGraph class
        pg = PlanningGraph(self, node.state)
        pg_levelsum = pg.h_levelsum()
        return pg_levelsum

    def h_ignore_preconditions(self, node: Node):
        '''
        This heuristic estimates the minimum number of actions that must be
        carried out from the current state in order to satisfy all of the goal
        conditions by ignoring the preconditions required for an action to be
        executed.
        '''
        # TODO implement (see Russell-Norvig Ed-3 10.2.3  or Russell-Norvig Ed-2 11.2)
        count = 0
        kb = PropKB()
        kb.tell(decode_state(node.state, self.state_map).pos_sentence())
        for clause in self.goal:
            if clause not in kb.clauses:
                count += 1
        return count

def air_cargo_p1() -> AirCargoProblem:
    cargos = ['C1', 'C2']
    planes = ['P1', 'P2']
    airports = ['JFK', 'SFO']
    pos = [expr('At(C1, SFO)'),
           expr('At(C2, JFK)'),
           expr('At(P1, SFO)'),
           expr('At(P2, JFK)'),
           ]
    neg = [expr('At(C2, SFO)'),
           expr('In(C2, P1)'),
           expr('In(C2, P2)'),
           expr('At(C1, JFK)'),
           expr('In(C1, P1)'),
           expr('In(C1, P2)'),
           expr('At(P1, JFK)'),
           expr('At(P2, SFO)'),
           ]
    init = FluentState(pos, neg)
    goal = [expr('At(C1, JFK)'),
            expr('At(C2, SFO)'),
            ]
    return AirCargoProblem(cargos, planes, airports, init, goal)

def air_cargo_p2() -> AirCargoProblem:
    cargos = ['C1', 'C2', 'C3']
    planes = ['P1', 'P2', 'P3']
    airports = ['JFK', 'SFO', 'ATL']
    pos = [expr('At(C1, SFO)'),
           expr('At(C2, JFK)'),
           expr('At(C3, ATL)'),
           expr('At(P1, SFO)'),
           expr('At(P2, JFK)'),
           expr('At(P3, ATL)'),
           ]
    neg = [expr('At(C1, JFK)'),
           expr('At(C1, ATL)'),
           expr('In(C1, P1)'),
           expr('In(C1, P2)'),
           expr('In(C1, P3)'),
           expr('At(C2, SFO)'),
           expr('At(C2, ATL)'),
           expr('In(C2, P1)'),
           expr('In(C2, P2)'),
           expr('In(C2, P3)'),
           expr('At(C3, SFO)'),
           expr('At(C3, JFK)'),
           expr('In(C3, P1)'),
           expr('In(C3, P2)'),
           expr('In(C3, P3)'),
           expr('At(P1, JFK)'),
           expr('At(P1, ATL)'),
           expr('At(P2, SFO)'),
           expr('At(P2, ATL)'),
           expr('At(P3, SFO)'),
           expr('At(P3, JFK)'),
           ]
    init = FluentState(pos, neg)
    goal = [expr('At(C1, JFK)'),
            expr('At(C2, SFO)'),
            expr('At(C3, SFO)'),
            ]
    return AirCargoProblem(cargos, planes, airports, init, goal)

def air_cargo_p3() -> AirCargoProblem:
    cargos = ['C1', 'C2', 'C3', 'C4']
    planes = ['P1', 'P2', 'P3', 'P4']
    airports = ['JFK', 'SFO', 'ATL', 'ORD']
    pos = [expr('At(C1, SFO)'),
           expr('At(C2, JFK)'),
           expr('At(C3, ATL)'),
           expr('At(C4, ORD)'),
           expr('At(P1, SFO)'),
           expr('At(P2, JFK)'),
           ]
    neg = [expr('At(C1, JFK)'),
           expr('At(C1, ATL)'),
           expr('At(C1, ORD)'),
           expr('In(C1, P1)'),
           expr('In(C1, P2)'),
           expr('At(C2, SFO)'),
           expr('At(C2, ATL)'),
           expr('At(C2, ORD)'),
           expr('In(C2, P1)'),
           expr('In(C2, P2)'),
           expr('At(C3, SFO)'),
           expr('At(C3, JFK)'),
           expr('At(C3, ORD)'),
           expr('In(C3, P1)'),
           expr('In(C3, P2)'),
           expr('At(C4, SFO)'),
           expr('At(C4, JFK)'),
           expr('At(C4, ATL)'),
           expr('In(C4, P1)'),
           expr('In(C4, P2)'),
           expr('At(P1, JFK)'),
           expr('At(P1, ATL)'),
           expr('At(P1, ORD)'),
           expr('At(P2, SFO)'),
           expr('At(P2, ATL)'),
           expr('At(P2, ORD)'),
           ]
    init = FluentState(pos, neg)
    goal = [expr('At(C1, JFK)'),
            expr('At(C2, SFO)'),
            expr('At(C3, JFK)'),
            expr('At(C4, SFO)'),
            ]
    return AirCargoProblem(cargos, planes, airports, init, goal)
<|endoftext|>


<commit_before>"""
Unit tests for NCMC switching engine.

"""

__author__ = 'John D. Chodera'

################################################################################
# IMPORTS
################################################################################

from simtk import openmm, unit
from simtk.openmm import app
import os, os.path
import sys, math
import numpy as np
from functools import partial
from pkg_resources import resource_filename
from openeye import oechem
if sys.version_info >= (3, 0):
    from io import StringIO
    from subprocess import getstatusoutput
else:
    from cStringIO import StringIO
    from commands import getstatusoutput

################################################################################
# CONSTANTS
################################################################################

kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
temperature = 300.0 * unit.kelvin
kT = kB * temperature
beta = 1.0/kT

################################################################################
# TESTS
################################################################################

def align_molecules(mol1, mol2):
    """
    MCSS two OEmols. Return the mapping of new : old atoms
    """
    mcs = oechem.OEMCSSearch(oechem.OEMCSType_Exhaustive)
    atomexpr = oechem.OEExprOpts_AtomicNumber
    bondexpr = 0
    mcs.Init(mol1, atomexpr, bondexpr)
    mcs.SetMCSFunc(oechem.OEMCSMaxAtomsCompleteCycles())
    unique = True
    match = [m for m in mcs.Match(mol2, unique)][0]
    new_to_old_atom_mapping = {}
    for matchpair in match.GetAtoms():
        old_index = matchpair.pattern.GetIdx()
        new_index = matchpair.target.GetIdx()
        new_to_old_atom_mapping[new_index] = old_index
    return new_to_old_atom_mapping

def simulate(system, positions, nsteps=500, timestep=1.0*unit.femtoseconds, temperature=temperature, collision_rate=5.0/unit.picoseconds, platform=None):
    integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    if platform == None:
        context = openmm.Context(system, integrator)
    else:
        context = openmm.Context(system, integrator, platform)
    context.setPositions(positions)
    context.setVelocitiesToTemperature(temperature)
    integrator.step(nsteps)
    positions = context.getState(getPositions=True).getPositions(asNumpy=True)
    velocities = context.getState(getVelocities=True).getVelocities(asNumpy=True)
    return [positions, velocities]

def check_alchemical_null_elimination(topology_proposal, positions, ncmc_nsteps=50, NSIGMA_MAX=6.0):
    """
    Test alchemical elimination engine on null transformations, where some atoms are deleted and then reinserted in a cycle.

    Parameters
    ----------
    topology_proposal : TopologyProposal
        The topology proposal to test.
        This must be a null transformation, where topology_proposal.old_system == topology_proposal.new_system
    ncmc_steps : int, optional, default=50
        Number of NCMC switching steps, or 0 for instantaneous switching.
    NSIGMA_MAX : float, optional, default=6.0
        Number of standard errors away from analytical solution tolerated before Exception is thrown
    """
    # Initialize engine
    from perses.annihilation.ncmc_switching import NCMCEngine
    ncmc_engine = NCMCEngine(temperature=temperature, nsteps=ncmc_nsteps)

    # Make sure that old system and new system are identical.
    if not (topology_proposal.old_system == topology_proposal.new_system):
        raise Exception("topology_proposal must be a null transformation for this test (old_system == new_system)")
    for (k,v) in topology_proposal.new_to_old_atom_map.items():
        if k != v:
            raise Exception("topology_proposal must be a null transformation for this test (retailed atoms must map onto themselves)")

    nequil = 5 # number of equilibration iterations
    niterations = 20 # number of round-trip switching trials
    logP_insert_n = np.zeros([niterations], np.float64)
    logP_delete_n = np.zeros([niterations], np.float64)
    logP_switch_n = np.zeros([niterations], np.float64)
    print("")
    for iteration in range(nequil):
        [positions, velocities] = simulate(topology_proposal.old_system, positions)
    for iteration in range(niterations):
        # Equilibrate
        [positions, velocities] = simulate(topology_proposal.old_system, positions)

        # Check that positions are not NaN
        if(np.any(np.isnan(positions / unit.angstroms))):
            raise Exception("Positions became NaN during equilibration")

        # Delete atoms
        [positions, logP_delete, potential_delete] = ncmc_engine.integrate(topology_proposal, positions, direction='delete')

        # Check that positions are not NaN
        if(np.any(np.isnan(positions / unit.angstroms))):
            raise Exception("Positions became NaN on NCMC deletion")

        # Insert atoms
        [positions, logP_insert, potential_insert] = ncmc_engine.integrate(topology_proposal, positions, direction='insert')

        # Check that positions are not NaN
        if(np.any(np.isnan(positions / unit.angstroms))):
            raise Exception("Positions became NaN on NCMC insertion")

        # Compute probability of switching geometries.
        logP_switch = - (potential_insert - potential_delete)

        # Compute total probability
        logP_delete_n[iteration] = logP_delete
        logP_insert_n[iteration] = logP_insert
        logP_switch_n[iteration] = logP_switch
        #print("Iteration %5d : delete %16.8f kT | insert %16.8f kT | geometry switch %16.8f" % (iteration, logP_delete, logP_insert, logP_switch))

    # Check free energy difference is withing NSIGMA_MAX standard errors of zero.
    logP_n = logP_delete_n + logP_insert_n + logP_switch_n
    work_n = - logP_n
    from pymbar import EXP
    [df, ddf] = EXP(work_n)
    #print("df = %12.6f +- %12.5f kT" % (df, ddf))
    if (abs(df) > NSIGMA_MAX * ddf):
        msg = 'Delta F (%d steps switching) = %f +- %f kT; should be within %f sigma of 0\n' % (ncmc_nsteps, df, ddf, NSIGMA_MAX)
        msg += 'delete logP:\n'
        msg += str(logP_delete_n) + '\n'
        msg += 'insert logP:\n'
        msg += str(logP_insert_n) + '\n'
        msg += 'logP:\n'
        msg += str(logP_n) + '\n'
        raise Exception(msg)

def test_alchemical_elimination_mutation():
    """
    Test alchemical elimination for mutations.
    """

    ff_filename = "amber99sbildn.xml"
    max_point_mutants = 1
    proposal_metadata = {'ffxmls':[ff_filename]}

    # Create peptide.
    from openmmtools import testsystems
    testsystem = testsystems.AlanineDipeptideVacuum()
    [topology, system, positions] = [testsystem.topology, testsystem.system, testsystem.positions]

    # Create forcefield.
    ff = app.ForceField(ff_filename)
    chain_id = ' '
    allowed_mutations = [[('2','GLY')]]

    from perses.rjmc.topology_proposal import SystemGenerator
    system_generator = SystemGenerator([ff_filename])

    # Create a topology proposal fro mutating ALA -> GLY
    from perses.rjmc.topology_proposal import PointMutationEngine
    proposal_engine = PointMutationEngine(system_generator, max_point_mutants, chain_id, proposal_metadata=proposal_metadata, allowed_mutations=allowed_mutations)
    topology_proposal = proposal_engine.propose(system, topology)

    # Modify atom mapping to get a null transformation.
    from perses.rjmc.topology_proposal import TopologyProposal
    new_to_old_atom_map = { atom1 : atom1 for atom1 in topology_proposal.new_to_old_atom_map }
    topology_proposal = TopologyProposal(
                new_topology=topology_proposal.old_topology, new_system=topology_proposal.old_system, old_topology=topology_proposal.old_topology, old_system=topology_proposal.old_system,
                old_chemical_state_key='AA', new_chemical_state_key='AG', logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata=topology_proposal.metadata)

    for ncmc_nsteps in [0, 1, 2, 50]:
        f = partial(check_alchemical_null_elimination, topology_proposal, positions, ncmc_nsteps=ncmc_nsteps)
        f.description = "Testing alchemical null transformation of ALA sidechain in alanine dipeptide with %d NCMC steps" % ncmc_nsteps
        yield f

def test_ncmc_alchemical_integrator_stability_molecules():
    """
    Test NCMCAlchemicalIntegrator

    """
    molecule_names = ['pentane', 'biphenyl', 'imatinib']
    for molecule_name in molecule_names:
        from perses.tests.utils import createSystemFromIUPAC
        [molecule, system, positions, topology] = createSystemFromIUPAC(molecule_name)

        # Eliminate half of the molecule
        # TODO: Use a more rigorous scheme to make sure we are really cutting the molecule in half and not just eliminating hydrogens or something.
        alchemical_atoms = [ index for index in range(int(system.getNumParticles()/2)) ]

        # Create an alchemically-modified system.
        from alchemy import AbsoluteAlchemicalFactory
        alchemical_factory = AbsoluteAlchemicalFactory(system, ligand_atoms=alchemical_atoms, annihilate_electrostatics=True, annihilate_sterics=True)

        # Return the alchemically-modified system in fully-interacting form.
        alchemical_system = alchemical_factory.createPerturbedSystem()

        # Create an NCMC switching integrator.
        from perses.annihilation.ncmc_switching import NCMCAlchemicalIntegrator
        temperature = 300.0 * unit.kelvin
        functions = { 'lambda_sterics' : 'lambda', 'lambda_electrostatics' : 'lambda^0.5', 'lambda_torsions' : 'lambda', 'lambda_angles' : 'lambda^2' }
        ncmc_integrator = NCMCAlchemicalIntegrator(temperature, alchemical_system, functions, direction='delete', nsteps=10, timestep=1.0*unit.femtoseconds)

        # Create a Context
        context = openmm.Context(alchemical_system, ncmc_integrator)
        context.setPositions(positions)

        # Run the integrator
        ncmc_integrator.step(1)

        # Check positions are finite
        positions = context.getState(getPositions=True).getPositions(asNumpy=True)
        if np.isnan(np.any(positions / positions.unit)):
            raise Exception('NCMCAlchemicalIntegrator gave NaN positions')
        if np.isnan(ncmc_integrator.getLogAcceptanceProbability()):
            raise Exception('NCMCAlchemicalIntegrator gave NaN logAcceptanceProbability')

        del context, ncmc_integrator

def test_ncmc_engine_molecule():
    """
    Check alchemical elimination for alanine dipeptide in vacuum with 0, 1, 2, and 50 switching steps.
    """
    molecule_names = ['ethane', 'pentane', 'benzene', 'phenol', 'biphenyl', 'imatinib']
    for molecule_name in molecule_names:
        from perses.tests.utils import createSystemFromIUPAC
        [molecule, system, positions, topology] = createSystemFromIUPAC(molecule_name)
        natoms = system.getNumParticles()
        # Eliminate half of the molecule
        # TODO: Use a more rigorous scheme to make sure we are really cutting the molecule in half and not just eliminating hydrogens or something.
        new_to_old_atom_map = { index : index for index in range(int(natoms/2)) }

        from perses.rjmc.topology_proposal import TopologyProposal
        topology_proposal = TopologyProposal(
            new_topology=topology, new_system=system, old_topology=topology, old_system=system,
            old_chemical_state_key='', new_chemical_state_key='', logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata={'test':0.0})
        for ncmc_nsteps in [0, 1, 50]:
            f = partial(check_alchemical_null_elimination, topology_proposal, positions, ncmc_nsteps=ncmc_nsteps)
            f.description = "Testing alchemical null elimination for '%s' with %d NCMC steps" % (molecule_name, ncmc_nsteps)
            yield f

def test_alchemical_elimination_peptide():
    """
    Test alchemical elimination for the alanine dipeptide.
    """
    # Create an alanine dipeptide null transformation, where N-methyl group is deleted and then inserted.
    from openmmtools import testsystems
    testsystem = testsystems.AlanineDipeptideVacuum()
    from perses.rjmc.topology_proposal import TopologyProposal
    new_to_old_atom_map = { index : index for index in range(testsystem.system.getNumParticles()) if (index > 3) } # all atoms but N-methyl
    topology_proposal = TopologyProposal(
        old_system=testsystem.system, old_topology=testsystem.topology,
        old_chemical_state_key='AA', new_chemical_state_key='AA',
        new_system=testsystem.system, new_topology=testsystem.topology,
        logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata=dict())

    for ncmc_nsteps in [0, 1, 50]:
        f = partial(check_alchemical_null_elimination, topology_proposal, testsystem.positions, ncmc_nsteps=ncmc_nsteps)
        f.description = "Testing alchemical elimination using alanine dipeptide with %d NCMC steps" % ncmc_nsteps
        yield f

if __name__ == "__main__":
    print(__name__)
    for x in test_alchemical_elimination_molecule():
        x()
<commit_msg>fixed input arguments in test<commit_after>"""
Unit tests for NCMC switching engine.

"""

__author__ = 'John D. Chodera'

################################################################################
# IMPORTS
################################################################################

from simtk import openmm, unit
from simtk.openmm import app
import os, os.path
import sys, math
import numpy as np
from functools import partial
from pkg_resources import resource_filename
from openeye import oechem
if sys.version_info >= (3, 0):
    from io import StringIO
    from subprocess import getstatusoutput
else:
    from cStringIO import StringIO
    from commands import getstatusoutput

################################################################################
# CONSTANTS
################################################################################

kB = unit.BOLTZMANN_CONSTANT_kB * unit.AVOGADRO_CONSTANT_NA
temperature = 300.0 * unit.kelvin
kT = kB * temperature
beta = 1.0/kT

################################################################################
# TESTS
################################################################################

def align_molecules(mol1, mol2):
    """
    MCSS two OEmols. Return the mapping of new : old atoms
    """
    mcs = oechem.OEMCSSearch(oechem.OEMCSType_Exhaustive)
    atomexpr = oechem.OEExprOpts_AtomicNumber
    bondexpr = 0
    mcs.Init(mol1, atomexpr, bondexpr)
    mcs.SetMCSFunc(oechem.OEMCSMaxAtomsCompleteCycles())
    unique = True
    match = [m for m in mcs.Match(mol2, unique)][0]
    new_to_old_atom_mapping = {}
    for matchpair in match.GetAtoms():
        old_index = matchpair.pattern.GetIdx()
        new_index = matchpair.target.GetIdx()
        new_to_old_atom_mapping[new_index] = old_index
    return new_to_old_atom_mapping

def simulate(system, positions, nsteps=500, timestep=1.0*unit.femtoseconds, temperature=temperature, collision_rate=5.0/unit.picoseconds, platform=None):
    integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    if platform == None:
        context = openmm.Context(system, integrator)
    else:
        context = openmm.Context(system, integrator, platform)
    context.setPositions(positions)
    context.setVelocitiesToTemperature(temperature)
    integrator.step(nsteps)
    positions = context.getState(getPositions=True).getPositions(asNumpy=True)
    velocities = context.getState(getVelocities=True).getVelocities(asNumpy=True)
    return [positions, velocities]

def check_alchemical_null_elimination(topology_proposal, positions, ncmc_nsteps=50, NSIGMA_MAX=6.0):
    """
    Test alchemical elimination engine on null transformations, where some atoms are deleted and then reinserted in a cycle.

    Parameters
    ----------
    topology_proposal : TopologyProposal
        The topology proposal to test.
        This must be a null transformation, where topology_proposal.old_system == topology_proposal.new_system
    ncmc_steps : int, optional, default=50
        Number of NCMC switching steps, or 0 for instantaneous switching.
    NSIGMA_MAX : float, optional, default=6.0
        Number of standard errors away from analytical solution tolerated before Exception is thrown
    """
    # Initialize engine
    from perses.annihilation.ncmc_switching import NCMCEngine
    ncmc_engine = NCMCEngine(temperature=temperature, nsteps=ncmc_nsteps)

    # Make sure that old system and new system are identical.
    if not (topology_proposal.old_system == topology_proposal.new_system):
        raise Exception("topology_proposal must be a null transformation for this test (old_system == new_system)")
    for (k,v) in topology_proposal.new_to_old_atom_map.items():
        if k != v:
            raise Exception("topology_proposal must be a null transformation for this test (retailed atoms must map onto themselves)")

    nequil = 5 # number of equilibration iterations
    niterations = 20 # number of round-trip switching trials
    logP_insert_n = np.zeros([niterations], np.float64)
    logP_delete_n = np.zeros([niterations], np.float64)
    logP_switch_n = np.zeros([niterations], np.float64)
    print("")
    for iteration in range(nequil):
        [positions, velocities] = simulate(topology_proposal.old_system, positions)
    for iteration in range(niterations):
        # Equilibrate
        [positions, velocities] = simulate(topology_proposal.old_system, positions)

        # Check that positions are not NaN
        if(np.any(np.isnan(positions / unit.angstroms))):
            raise Exception("Positions became NaN during equilibration")

        # Delete atoms
        [positions, logP_delete, potential_delete] = ncmc_engine.integrate(topology_proposal, positions, direction='delete')

        # Check that positions are not NaN
        if(np.any(np.isnan(positions / unit.angstroms))):
            raise Exception("Positions became NaN on NCMC deletion")

        # Insert atoms
        [positions, logP_insert, potential_insert] = ncmc_engine.integrate(topology_proposal, positions, direction='insert')

        # Check that positions are not NaN
        if(np.any(np.isnan(positions / unit.angstroms))):
            raise Exception("Positions became NaN on NCMC insertion")

        # Compute probability of switching geometries.
        logP_switch = - (potential_insert - potential_delete)

        # Compute total probability
        logP_delete_n[iteration] = logP_delete
        logP_insert_n[iteration] = logP_insert
        logP_switch_n[iteration] = logP_switch
        #print("Iteration %5d : delete %16.8f kT | insert %16.8f kT | geometry switch %16.8f" % (iteration, logP_delete, logP_insert, logP_switch))

    # Check free energy difference is withing NSIGMA_MAX standard errors of zero.
    logP_n = logP_delete_n + logP_insert_n + logP_switch_n
    work_n = - logP_n
    from pymbar import EXP
    [df, ddf] = EXP(work_n)
    #print("df = %12.6f +- %12.5f kT" % (df, ddf))
    if (abs(df) > NSIGMA_MAX * ddf):
        msg = 'Delta F (%d steps switching) = %f +- %f kT; should be within %f sigma of 0\n' % (ncmc_nsteps, df, ddf, NSIGMA_MAX)
        msg += 'delete logP:\n'
        msg += str(logP_delete_n) + '\n'
        msg += 'insert logP:\n'
        msg += str(logP_insert_n) + '\n'
        msg += 'logP:\n'
        msg += str(logP_n) + '\n'
        raise Exception(msg)

def test_alchemical_elimination_mutation():
    """
    Test alchemical elimination for mutations.
    """

    ff_filename = "amber99sbildn.xml"
    proposal_metadata = {'ffxmls':[ff_filename]}

    # Create peptide.
    from openmmtools import testsystems
    testsystem = testsystems.AlanineDipeptideVacuum()
    [topology, system, positions] = [testsystem.topology, testsystem.system, testsystem.positions]

    # Create forcefield.
    ff = app.ForceField(ff_filename)
    chain_id = ' '
    allowed_mutations = [[('2','GLY')]]

    from perses.rjmc.topology_proposal import SystemGenerator
    system_generator = SystemGenerator([ff_filename])

    # Create a topology proposal fro mutating ALA -> GLY
    from perses.rjmc.topology_proposal import PointMutationEngine
    proposal_engine = PointMutationEngine(system_generator, chain_id, proposal_metadata=proposal_metadata, allowed_mutations=allowed_mutations)
    topology_proposal = proposal_engine.propose(system, topology)

    # Modify atom mapping to get a null transformation.
    from perses.rjmc.topology_proposal import TopologyProposal
    new_to_old_atom_map = { atom1 : atom1 for atom1 in topology_proposal.new_to_old_atom_map }
    topology_proposal = TopologyProposal(
                new_topology=topology_proposal.old_topology, new_system=topology_proposal.old_system, old_topology=topology_proposal.old_topology, old_system=topology_proposal.old_system,
                old_chemical_state_key='AA', new_chemical_state_key='AG', logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata=topology_proposal.metadata)

    for ncmc_nsteps in [0, 1, 2, 50]:
        f = partial(check_alchemical_null_elimination, topology_proposal, positions, ncmc_nsteps=ncmc_nsteps)
        f.description = "Testing alchemical null transformation of ALA sidechain in alanine dipeptide with %d NCMC steps" % ncmc_nsteps
        yield f

def test_ncmc_alchemical_integrator_stability_molecules():
    """
    Test NCMCAlchemicalIntegrator

    """
    molecule_names = ['pentane', 'biphenyl', 'imatinib']
    for molecule_name in molecule_names:
        from perses.tests.utils import createSystemFromIUPAC
        [molecule, system, positions, topology] = createSystemFromIUPAC(molecule_name)

        # Eliminate half of the molecule
        # TODO: Use a more rigorous scheme to make sure we are really cutting the molecule in half and not just eliminating hydrogens or something.
        alchemical_atoms = [ index for index in range(int(system.getNumParticles()/2)) ]

        # Create an alchemically-modified system.
        from alchemy import AbsoluteAlchemicalFactory
        alchemical_factory = AbsoluteAlchemicalFactory(system, ligand_atoms=alchemical_atoms, annihilate_electrostatics=True, annihilate_sterics=True)

        # Return the alchemically-modified system in fully-interacting form.
        alchemical_system = alchemical_factory.createPerturbedSystem()

        # Create an NCMC switching integrator.
        from perses.annihilation.ncmc_switching import NCMCAlchemicalIntegrator
        temperature = 300.0 * unit.kelvin
        functions = { 'lambda_sterics' : 'lambda', 'lambda_electrostatics' : 'lambda^0.5', 'lambda_torsions' : 'lambda', 'lambda_angles' : 'lambda^2' }
        ncmc_integrator = NCMCAlchemicalIntegrator(temperature, alchemical_system, functions, direction='delete', nsteps=10, timestep=1.0*unit.femtoseconds)

        # Create a Context
        context = openmm.Context(alchemical_system, ncmc_integrator)
        context.setPositions(positions)

        # Run the integrator
        ncmc_integrator.step(1)

        # Check positions are finite
        positions = context.getState(getPositions=True).getPositions(asNumpy=True)
        if np.isnan(np.any(positions / positions.unit)):
            raise Exception('NCMCAlchemicalIntegrator gave NaN positions')
        if np.isnan(ncmc_integrator.getLogAcceptanceProbability()):
            raise Exception('NCMCAlchemicalIntegrator gave NaN logAcceptanceProbability')

        del context, ncmc_integrator

def test_ncmc_engine_molecule():
    """
    Check alchemical elimination for alanine dipeptide in vacuum with 0, 1, 2, and 50 switching steps.
    """
    molecule_names = ['ethane', 'pentane', 'benzene', 'phenol', 'biphenyl', 'imatinib']
    for molecule_name in molecule_names:
        from perses.tests.utils import createSystemFromIUPAC
        [molecule, system, positions, topology] = createSystemFromIUPAC(molecule_name)
        natoms = system.getNumParticles()
        # Eliminate half of the molecule
        # TODO: Use a more rigorous scheme to make sure we are really cutting the molecule in half and not just eliminating hydrogens or something.
        new_to_old_atom_map = { index : index for index in range(int(natoms/2)) }

        from perses.rjmc.topology_proposal import TopologyProposal
        topology_proposal = TopologyProposal(
            new_topology=topology, new_system=system, old_topology=topology, old_system=system,
            old_chemical_state_key='', new_chemical_state_key='', logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata={'test':0.0})
        for ncmc_nsteps in [0, 1, 50]:
            f = partial(check_alchemical_null_elimination, topology_proposal, positions, ncmc_nsteps=ncmc_nsteps)
            f.description = "Testing alchemical null elimination for '%s' with %d NCMC steps" % (molecule_name, ncmc_nsteps)
            yield f

def test_alchemical_elimination_peptide():
    """
    Test alchemical elimination for the alanine dipeptide.
    """
    # Create an alanine dipeptide null transformation, where N-methyl group is deleted and then inserted.
    from openmmtools import testsystems
    testsystem = testsystems.AlanineDipeptideVacuum()
    from perses.rjmc.topology_proposal import TopologyProposal
    new_to_old_atom_map = { index : index for index in range(testsystem.system.getNumParticles()) if (index > 3) } # all atoms but N-methyl
    topology_proposal = TopologyProposal(
        old_system=testsystem.system, old_topology=testsystem.topology,
        old_chemical_state_key='AA', new_chemical_state_key='AA',
        new_system=testsystem.system, new_topology=testsystem.topology,
        logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata=dict())

    for ncmc_nsteps in [0, 1, 50]:
        f = partial(check_alchemical_null_elimination, topology_proposal, testsystem.positions, ncmc_nsteps=ncmc_nsteps)
        f.description = "Testing alchemical elimination using alanine dipeptide with %d NCMC steps" % ncmc_nsteps
        yield f

if __name__ == "__main__":
    print(__name__)
    for x in test_alchemical_elimination_molecule():
        x()
<|endoftext|>


<commit_before>from __future__ import print_function, division, absolute_import
from fontTools.misc.py23 import *
from fontTools.misc.textTools import safeEval
from . import DefaultTable
import operator
import struct

class table_V_O_R_G_(DefaultTable.DefaultTable):

	""" This table is structured so that you can treat it like a dictionary keyed by glyph name.
	ttFont['VORG'][<glyphName>] will return the vertical origin for any glyph
	ttFont['VORG'][<glyphName>] = <value> will set the vertical origin for any glyph.
	"""

	def decompile(self, data, ttFont):
		self.getGlyphName = ttFont.getGlyphName # for use in get/set item functions, for access by GID
		self.majorVersion, self.minorVersion, self.defaultVertOriginY, self.numVertOriginYMetrics = struct.unpack(">HHhH", data[:8])
		assert (self.majorVersion <= 1), "Major version of VORG table is higher than I know how to handle"
		data = data[8:]
		vids = []
		gids = []
		pos = 0
		for i in range(self.numVertOriginYMetrics):
			gid, vOrigin = struct.unpack(">Hh", data[pos:pos+4])
			pos += 4
			gids.append(gid)
			vids.append(vOrigin)

		self.VOriginRecords = vOrig = {}
		glyphOrder = ttFont.getGlyphOrder()
		try:
			names = [glyphOrder[gid] for gid in gids]
		except IndexError:
			getGlyphName = self.getGlyphName
			names = map(getGlyphName, gids)

		for name, vid in zip(names, vids):
			vOrig[name] = vid

	def compile(self, ttFont):
		vorgs = list(self.VOriginRecords.values())
		names = list(self.VOriginRecords.keys())
		nameMap = ttFont.getReverseGlyphMap()
		try:
			gids = [nameMap[name] for name in names]
		except KeyError:
			nameMap = ttFont.getReverseGlyphMap(rebuild=True)
			gids = [nameMap[name] for name in names]
		vOriginTable = list(zip(gids, vorgs))
		self.numVertOriginYMetrics = len(vorgs)
		vOriginTable.sort() # must be in ascending GID order
		dataList = [ struct.pack(">Hh", rec[0], rec[1]) for rec in vOriginTable]
		header = struct.pack(">HHhH", self.majorVersion, self.minorVersion, self.defaultVertOriginY, self.numVertOriginYMetrics)
		dataList.insert(0, header)
		data = bytesjoin(dataList)
		return data

	def toXML(self, writer, ttFont):
		writer.simpletag("majorVersion", value=self.majorVersion)
		writer.newline()
		writer.simpletag("minorVersion", value=self.minorVersion)
		writer.newline()
		writer.simpletag("defaultVertOriginY", value=self.defaultVertOriginY)
		writer.newline()
		writer.simpletag("numVertOriginYMetrics", value=self.numVertOriginYMetrics)
		writer.newline()
		vOriginTable = []
		glyphNames = self.VOriginRecords.keys()
		for glyphName in glyphNames:
			try:
				gid = ttFont.getGlyphID(glyphName)
			except:
				assert 0, "VORG table contains a glyph name not in ttFont.getGlyphNames(): " + str(glyphName)
			vOriginTable.append([gid, glyphName, self.VOriginRecords[glyphName]])
		vOriginTable.sort()
		for entry in vOriginTable:
			vOriginRec = VOriginRecord(entry[1], entry[2])
			vOriginRec.toXML(writer, ttFont)

	def fromXML(self, name, attrs, content, ttFont):
		if not hasattr(self, "VOriginRecords"):
			self.VOriginRecords = {}
		self.getGlyphName = ttFont.getGlyphName # for use in get/set item functions, for access by GID
		if name == "VOriginRecord":
			vOriginRec = VOriginRecord()
			for element in content:
				if isinstance(element, basestring):
					continue
				name, attrs, content = element
				vOriginRec.fromXML(name, attrs, content, ttFont)
			self.VOriginRecords[vOriginRec.glyphName] = vOriginRec.vOrigin
		elif "value" in attrs:
			setattr(self, name, safeEval(attrs["value"]))

	def __getitem__(self, glyphSelector):
		if isinstance(glyphSelector, int):
			# its a gid, convert to glyph name
			glyphSelector = self.getGlyphName(glyphSelector)

		if glyphSelector not in self.VOriginRecords:
			return self.defaultVertOriginY

		return self.VOriginRecords[glyphSelector]

	def __setitem__(self, glyphSelector, value):
		if isinstance(glyphSelector, int):
			# its a gid, convert to glyph name
			glyphSelector = self.getGlyphName(glyphSelector)

		if  value != self.defaultVertOriginY:
			self.VOriginRecords[glyphSelector] = value
		elif glyphSelector in self.VOriginRecords:
			del self.VOriginRecords[glyphSelector]

	def __delitem__(self, glyphSelector):
		del self.VOriginRecords[glyphSelector]

class VOriginRecord(object):

	def __init__(self, name=None, vOrigin=None):
		self.glyphName = name
		self.vOrigin = vOrigin

	def toXML(self, writer, ttFont):
		writer.begintag("VOriginRecord")
		writer.newline()
		writer.simpletag("glyphName", value=self.glyphName)
		writer.newline()
		writer.simpletag("vOrigin", value=self.vOrigin)
		writer.newline()
		writer.endtag("VOriginRecord")
		writer.newline()

	def fromXML(self, name, attrs, content, ttFont):
		value = attrs["value"]
		if name == "glyphName":
			setattr(self, name, value)
		else:
			setattr(self, name, safeEval(value))
<commit_msg>whitespace<commit_after>from __future__ import print_function, division, absolute_import
from fontTools.misc.py23 import *
from fontTools.misc.textTools import safeEval
from . import DefaultTable
import operator
import struct

class table_V_O_R_G_(DefaultTable.DefaultTable):

	""" This table is structured so that you can treat it like a dictionary keyed by glyph name.
	ttFont['VORG'][<glyphName>] will return the vertical origin for any glyph
	ttFont['VORG'][<glyphName>] = <value> will set the vertical origin for any glyph.
	"""

	def decompile(self, data, ttFont):
		self.getGlyphName = ttFont.getGlyphName # for use in get/set item functions, for access by GID
		self.majorVersion, self.minorVersion, self.defaultVertOriginY, self.numVertOriginYMetrics = struct.unpack(">HHhH", data[:8])
		assert (self.majorVersion <= 1), "Major version of VORG table is higher than I know how to handle"
		data = data[8:]
		vids = []
		gids = []
		pos = 0
		for i in range(self.numVertOriginYMetrics):
			gid, vOrigin = struct.unpack(">Hh", data[pos:pos+4])
			pos += 4
			gids.append(gid)
			vids.append(vOrigin)

		self.VOriginRecords = vOrig = {}
		glyphOrder = ttFont.getGlyphOrder()
		try:
			names = [glyphOrder[gid] for gid in gids]
		except IndexError:
			getGlyphName = self.getGlyphName
			names = map(getGlyphName, gids)

		for name, vid in zip(names, vids):
			vOrig[name] = vid

	def compile(self, ttFont):
		vorgs = list(self.VOriginRecords.values())
		names = list(self.VOriginRecords.keys())
		nameMap = ttFont.getReverseGlyphMap()
		try:
			gids = [nameMap[name] for name in names]
		except KeyError:
			nameMap = ttFont.getReverseGlyphMap(rebuild=True)
			gids = [nameMap[name] for name in names]
		vOriginTable = list(zip(gids, vorgs))
		self.numVertOriginYMetrics = len(vorgs)
		vOriginTable.sort() # must be in ascending GID order
		dataList = [struct.pack(">Hh", rec[0], rec[1]) for rec in vOriginTable]
		header = struct.pack(">HHhH", self.majorVersion, self.minorVersion, self.defaultVertOriginY, self.numVertOriginYMetrics)
		dataList.insert(0, header)
		data = bytesjoin(dataList)
		return data

	def toXML(self, writer, ttFont):
		writer.simpletag("majorVersion", value=self.majorVersion)
		writer.newline()
		writer.simpletag("minorVersion", value=self.minorVersion)
		writer.newline()
		writer.simpletag("defaultVertOriginY", value=self.defaultVertOriginY)
		writer.newline()
		writer.simpletag("numVertOriginYMetrics", value=self.numVertOriginYMetrics)
		writer.newline()
		vOriginTable = []
		glyphNames = self.VOriginRecords.keys()
		for glyphName in glyphNames:
			try:
				gid = ttFont.getGlyphID(glyphName)
			except:
				assert 0, "VORG table contains a glyph name not in ttFont.getGlyphNames(): " + str(glyphName)
			vOriginTable.append([gid, glyphName, self.VOriginRecords[glyphName]])
		vOriginTable.sort()
		for entry in vOriginTable:
			vOriginRec = VOriginRecord(entry[1], entry[2])
			vOriginRec.toXML(writer, ttFont)

	def fromXML(self, name, attrs, content, ttFont):
		if not hasattr(self, "VOriginRecords"):
			self.VOriginRecords = {}
		self.getGlyphName = ttFont.getGlyphName # for use in get/set item functions, for access by GID
		if name == "VOriginRecord":
			vOriginRec = VOriginRecord()
			for element in content:
				if isinstance(element, basestring):
					continue
				name, attrs, content = element
				vOriginRec.fromXML(name, attrs, content, ttFont)
			self.VOriginRecords[vOriginRec.glyphName] = vOriginRec.vOrigin
		elif "value" in attrs:
			setattr(self, name, safeEval(attrs["value"]))

	def __getitem__(self, glyphSelector):
		if isinstance(glyphSelector, int):
			# its a gid, convert to glyph name
			glyphSelector = self.getGlyphName(glyphSelector)

		if glyphSelector not in self.VOriginRecords:
			return self.defaultVertOriginY

		return self.VOriginRecords[glyphSelector]

	def __setitem__(self, glyphSelector, value):
		if isinstance(glyphSelector, int):
			# its a gid, convert to glyph name
			glyphSelector = self.getGlyphName(glyphSelector)

		if  value != self.defaultVertOriginY:
			self.VOriginRecords[glyphSelector] = value
		elif glyphSelector in self.VOriginRecords:
			del self.VOriginRecords[glyphSelector]

	def __delitem__(self, glyphSelector):
		del self.VOriginRecords[glyphSelector]

class VOriginRecord(object):

	def __init__(self, name=None, vOrigin=None):
		self.glyphName = name
		self.vOrigin = vOrigin

	def toXML(self, writer, ttFont):
		writer.begintag("VOriginRecord")
		writer.newline()
		writer.simpletag("glyphName", value=self.glyphName)
		writer.newline()
		writer.simpletag("vOrigin", value=self.vOrigin)
		writer.newline()
		writer.endtag("VOriginRecord")
		writer.newline()

	def fromXML(self, name, attrs, content, ttFont):
		value = attrs["value"]
		if name == "glyphName":
			setattr(self, name, value)
		else:
			setattr(self, name, safeEval(value))
<|endoftext|>