forked from inaos/iron-array-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_expr_udf.py
More file actions
79 lines (68 loc) · 2.69 KB
/
Copy pathtest_expr_udf.py
File metadata and controls
79 lines (68 loc) · 2.69 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import numpy as np
import pytest
import iarray as ia
from iarray import udf
from iarray.expr_udf import expr_udf
TOL = 1e-14
@udf.scalar(lib="lib_expr_udf")
def fsum(a: udf.float64, b: udf.float64) -> float:
return a + b
@pytest.mark.parametrize(
"sexpr, inputs",
[
("x + x", {"x": ia.arange(10, shape=(10,))}),
("x * x", {"x": ia.arange(10, shape=(10,))}),
("x + y", {"x": ia.arange(10, shape=(10,)), "y": 1}), # scalar as param!
("2 * (x + x)", {"x": ia.arange(10, shape=(10,))}),
("2 + x * x", {"x": ia.arange(10, shape=(10,))}),
("2 + sin(x) + x * x", {"x": ia.arange(10, shape=(10,))}),
("2 * (sin(x) + cos(x)) + x * x", {"x": ia.arange(10, shape=(10,))}),
("2 + x * x * (x + x)", {"x": ia.arange(10, shape=(10,))}),
("2 + x * x * ((x * x) + x)", {"x": ia.arange(10, shape=(10,))}),
("x * y * ((x * y) + y)", {"x": ia.arange(10, shape=(10,)), "y": 2}),
("lib_expr_udf.fsum(x, x)", {"x": ia.arange(10, shape=(10,))}),
("x + y", {"x": ia.arange(100, shape=(10, 10)), "y": ia.arange(100, shape=(10, 10))}),
("absolute(x) + abs(x) + negative(x) + negate(x)", {"x": ia.arange(5, step=0.5, shape=[10])}),
(
"absolute(x) + abs(x) + negative(x) + negate(x)",
{"x": ia.arange(10, shape=[10], dtype=np.float32)},
),
(
"arccos(x) + arcsin(x) + arctan(x) + arctan2(x, x) + power(x, x)",
{"x": ia.arange(10, shape=[10])},
),
],
)
def test_simple(sexpr, inputs):
out = expr_udf(sexpr, inputs).eval()
ref_out = ia.expr_from_string(sexpr, inputs).eval()
np.testing.assert_allclose(out.data, ref_out.data, rtol=TOL, atol=TOL)
@pytest.mark.parametrize(
"condition, inputs",
[
("b > 5", {"a": ia.arange(10, shape=[10]), "b": ia.arange(10, shape=[10])}),
(
"b > 5",
{
"a": ia.arange(10, shape=[10], dtype=np.float32),
"b": ia.arange(10, shape=[10], dtype=np.float32),
},
),
(
"(b > 5) and not (a > 7) or (b > 42)",
{"a": ia.arange(100, shape=[10, 10]), "b": ia.arange(100, shape=[10, 10])},
),
("not (a == 4)", {"a": ia.arange(10, shape=[10])}),
],
)
def test_masks(condition, inputs):
sexpr = f"a[{condition}]"
out = expr_udf(sexpr, inputs).eval()
# Numpy
replace = {"and": "&", "or": "|", "not": "~"}
for k, v in replace.items():
condition = condition.replace(k, v)
a = inputs["a"].data
b = inputs["b"].data if "b" in inputs else None
out_ref = np.where(eval(condition), a, np.nan)
np.testing.assert_allclose(out.data, out_ref, rtol=TOL, atol=TOL)