Merge pull request #301 from choderalab/math_operators

Add logical and bitwise operators for math_eval

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(

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]