def test_rolling_mean(self):
mean_1d = _lowlevel.make_builtin_mean1d_arrfunc('float64', -1)
rolling_mean = _lowlevel.make_rolling_arrfunc(mean_1d, 4)
in0 = nd.array([3.0, 2, 1, 3, 8, nd.nan, nd.nan])
out = rolling_mean(in0)
result = nd.as_py(out)
self.assertTrue(np.all(np.isnan(result[:3])))
self.assertTrue(np.isnan(result[-1]))
self.assertEqual(result[3:-1], [9.0/4, 14.0/4, 12.0/3])
for typ in types:
x.add_overload(
"(%s) -> %s" % (typ.__name__, typ.__name__),
_lowlevel.arrfunc_from_ufunc(np_op, (typ,) * 3, False),
associative=True,
commutative=True,
identity=ident,
)
locals()[name] = x
# ------------------------------------------------------------------------
# Other Funcs
# ------------------------------------------------------------------------
rolling_mean = RollingWindowBlazeFunc("blaze", "rolling_mean")
mean1d = _lowlevel.make_builtin_mean1d_arrfunc("float64", 0)
rolling_mean.add_overload("(M * float64) -> M * float64", mean1d)
diff = BlazeFunc("blaze", "diff")
subtract_doubles_ck = _lowlevel.arrfunc_from_ufunc(np.subtract, (np.float64, np.float64, np.float64), False)
diff_pair_ck = _lowlevel.lift_reduction_arrfunc(
subtract_doubles_ck, "strided * float64", axis=0, commutative=False, associative=False
)
diff_ck = _lowlevel.make_rolling_arrfunc(diff_pair_ck, 2)
diff.add_overload("(M * float64) -> M * float64", diff_ck)
take = BlazeFunc("blaze", "take")
# Masked take
take.add_overload("(M * T, M * bool) -> var * T", _lowlevel.make_take_arrfunc())
# Indexed take
take.add_overload("(M * T, N * intptr) -> N * T", _lowlevel.make_take_arrfunc())
