Add != operator to live calculations

pynbody · Jun 16, 2024 · d78deca · d78deca
1 parent 3d888c4
commit d78deca
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/tangos/live_calculation/builtin_functions/arithmetic.py b/tangos/live_calculation/builtin_functions/arithmetic.py
@@ -52,6 +52,10 @@ def less(halos, vals1, vals2):
 def equal(halos, vals1, vals2):
     return arithmetic_binary_op(vals1, vals2, np.equal)
 
+@BuiltinFunction.register
+def not_equal(halos, vals1, vals2):
+    return arithmetic_binary_op(vals1, vals2, np.not_equal)
+
 @BuiltinFunction.register
 def greater_equal(halos, vals1, vals2):
     return arithmetic_binary_op(vals1, vals2, np.greater_equal)

diff --git a/tangos/live_calculation/parser.py b/tangos/live_calculation/parser.py
@@ -35,6 +35,7 @@ def pack_args(for_function):
           ("|", "logical_or"),
           ("&", "logical_and"),
           ("==", "equal"),
+          ("!=", "not_equal"),
           (">=", "greater_equal"),
           ("<=", "less_equal")]
 

diff --git a/tests/test_live_calculation.py b/tests/test_live_calculation.py
@@ -216,6 +216,10 @@ def test_comparison():
     assert h.calculate("1.0>=1.0")
     assert h.calculate("1.0<=1.0")
     assert h.calculate("1.0>=0.5")
+    assert h.calculate("1==1")
+    assert h.calculate("finder_id()==1")
+    assert h.calculate("1!=2")
+    assert h.calculate("finder_id()!=2")
     assert not h.calculate("1.0<=0.5")