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())