diff --git a/openmmtools/tests/test_utils.py b/openmmtools/tests/test_utils.py index e39b2c3b1..c8f4131bb 100644 --- a/openmmtools/tests/test_utils.py +++ b/openmmtools/tests/test_utils.py @@ -57,7 +57,10 @@ def test_math_eval(): ('(x + lambda) / z * 4', {'x': 1, 'lambda': 2, 'z': 3}, 4.0), ('-((x + y) / z * 4)**2', {'x': 1, 'y': 2, 'z': 3}, -16.0), ('ceil(0.8) + acos(x) + step(0.5 - x) + step(0.5)', {'x': 1}, 2), - ('step_hm(x)', {'x': 0}, 0.5)] + ('step_hm(x)', {'x': 0}, 0.5), + ('myset & myset2', {'myset': {1,2,3}, 'myset2': {2,3,4}}, {2, 3}), + ('myset or myset2', {'myset': {1,2,3}, 'myset2': {2,3,4}}, {1, 2, 3, 4}), + ('(myset or my2set) & myset3', {'myset': {1, 2}, 'my2set': {3, 4}, 'myset3': {2, 3}}, {2, 3})] for expression, variables, result in test_cases: evaluated_expression = math_eval(expression, variables) assert evaluated_expression == result, '{}, {}, {}'.format( diff --git a/openmmtools/utils.py b/openmmtools/utils.py index f57709ab8..f47ad9ce7 100644 --- a/openmmtools/utils.py +++ b/openmmtools/utils.py @@ -233,6 +233,12 @@ def math_eval(expression, variables=None, functions=None): - step_hm(x) : Heaviside step function with half-maximum convention. - sign(x) : sign function (0.0 for x=0.0) + Available operators are ``+``, ``-``, ``*``, ``/``, ``**``, ``-x`` (negative), + ``&``, ``and``, ``|``, and ``or`` + + **The operators ``and`` and ``or`` operate BITWISE and behave the same as ``&`` and ``|`` respectively as this + function is not designed to handle logical operations.** If you provide sets, they must be as variables. + Parameters ---------- expression : str @@ -245,7 +251,7 @@ def math_eval(expression, variables=None, functions=None): Returns ------- - float + result The result of the evaluated expression. Examples @@ -259,7 +265,10 @@ def math_eval(expression, variables=None, functions=None): # Supported operators. operators = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv, - ast.Pow: operator.pow, ast.USub: operator.neg} + ast.Pow: operator.pow, ast.USub: operator.neg, + ast.BitAnd: operator.and_, ast.And: operator.and_, + ast.BitOr: operator.or_, ast.Or: operator.or_ + } # Supported functions, not defined in math. stock_functions = {'step': lambda x: 1 * (x >= 0), @@ -278,6 +287,15 @@ def _math_eval(node): elif isinstance(node, ast.BinOp): return operators[type(node.op)](_math_eval(node.left), _math_eval(node.right)) + elif isinstance(node, ast.BoolOp): + # Parse ternary operator + if len(node.values) > 2: + # Left-to-right precedence. + left_value = copy.deepcopy(node) + left_value.values.pop(-1) + else: + left_value = node.values[0] + return operators[type(node.op)](_math_eval(left_value), _math_eval(node.values[-1])) elif isinstance(node, ast.Name): try: return variables[node.id]