allow also scalar/scalar ops

basf · Feb 15, 2024 · 0966743 · 0966743
1 parent d22e06f
commit 0966743
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 12 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/rormula-rs/src/expression/ops_common.rs b/rormula-rs/src/expression/ops_common.rs
@@ -83,10 +83,15 @@ pub fn op_scalar(a: Value, b: Value, op: &impl Fn(f64, f64) -> f64) -> Value {
         }
         Value::Array(mem::take(arr))
     };
+    let sc_vs_sc = |sc1, sc2| Value::Scalar(op(sc1, sc2));
+
     match (a, b) {
         (Value::Array(mut arr), Value::Scalar(sc)) => arr_vs_sc(&mut arr, sc),
         (Value::Scalar(sc), Value::Array(mut arr)) => sc_vs_arr(sc, &mut arr),
-        _ => Value::Error("scalar op can only be applied to matrix and scalar".to_string()),
+        (Value::Scalar(sc1), Value::Scalar(sc2)) => sc_vs_sc(sc1, sc2),
+        _ => Value::Error(
+            "scalar op can only be applied to matrix and scalar or scalar and scalar".to_string(),
+        ),
     }
 }
 

diff --git a/rormula-rs/tests/test_rormula.rs b/rormula-rs/tests/test_rormula.rs
@@ -59,6 +59,13 @@ fn test_arithmetic() {
     let expr_ref = ExprArithmetic::parse(s_ref).unwrap();
     let rev_val = expr_ref.eval(&[Value::Array(Array2d::ones(5, 1))]).unwrap();
     assert_eq!(res, rev_val);
+    let s = "4*3";
+    let expr = ExprArithmetic::parse(s).unwrap();
+    let res = expr.eval_vec(vec![]).unwrap();
+    let sc_ref = Value::Scalar(12.0);
+    assert_eq!(res, sc_ref);
+    let s = "5/3 * alpha / beta * (0.2 / 200.0 / (29.22+gamma+epsilon+phi) / 7500)";
+    let _ = ExprArithmetic::parse(s).unwrap();
 }
 #[test]
 fn test_restrict() {

diff --git a/rormula/Cargo.toml b/rormula/Cargo.toml
@@ -9,7 +9,7 @@ name = "rormula"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = { version = "0.20.0", features = ["generate-import-lib"] }
+pyo3 = { version = "0.20.2", features = ["generate-import-lib"] }
 numpy = "0.20.0"
 rormula-rs = { path = "../rormula-rs" }
 

diff --git a/rormula/test/test_arithmetic.py b/rormula/test/test_arithmetic.py
@@ -61,6 +61,14 @@ def eval_asdf():
     res = rormula.eval_asdf(df)
     assert np.allclose(res.to_numpy().item(), (5.0 - 2.5) / 4.0)
 
+    data = np.ones((100, 6))
+    df = pd.DataFrame(
+        data=data, columns=["alpha", "beta", "gamma", "delta", "epsilon", "phi"]
+    )
+    s = "5/3 * alpha / beta * (0.2 / 200.0 / (29.22+gamma+epsilon+phi) / 1000)"
+    rormula = Arithmetic(s, "testslash")
+    res = rormula.eval_asdf(df)
+
 
 if __name__ == "__main__":
     test_arithmetic()