Overload df.rolling.kurt() (#475)

densmirn · AlexanderKalistratov · commit 58f2858714fa · 2020-01-24T17:11:56.000+03:00
* Overload df.rolling.kurt()

* Add perf.test for df.rolling.kurt()
diff --git a/examples/dataframe/rolling/dataframe_rolling_kurt.py b/examples/dataframe/rolling/dataframe_rolling_kurt.py
@@ -0,0 +1,41 @@
+# *****************************************************************************
+# Copyright (c) 2020, Intel Corporation All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#     Redistributions of source code must retain the above copyright notice,
+#     this list of conditions and the following disclaimer.
+#
+#     Redistributions in binary form must reproduce the above copyright notice,
+#     this list of conditions and the following disclaimer in the documentation
+#     and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
+# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
+# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
+# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+# EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+# *****************************************************************************
+
+import pandas as pd
+from numba import njit
+
+
+@njit
+def df_rolling_kurt():
+    df = pd.DataFrame({'A': [4, 3, 5, 2, 6], 'B': [-4, -3, -5, -2, -6]})
+    out_df = df.rolling(4).kurt()
+
+    # Expect DataFrame of
+    # {'A': [NaN, NaN, NaN, -1.2, -3.3], 'B': [NaN, NaN, NaN, -1.2, -3.3]}
+    return out_df
+
+
+print(df_rolling_kurt())
diff --git a/sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py b/sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py
@@ -149,6 +149,15 @@ def sdc_pandas_dataframe_rolling_count(self):
     return gen_df_rolling_method_impl('count', self)
 
 
+@sdc_overload_method(DataFrameRollingType, 'kurt')
+def sdc_pandas_dataframe_rolling_kurt(self):
+
+    ty_checker = TypeChecker('Method rolling.kurt().')
+    ty_checker.check(self, DataFrameRollingType)
+
+    return gen_df_rolling_method_impl('kurt', self)
+
+
 @sdc_overload_method(DataFrameRollingType, 'max')
 def sdc_pandas_dataframe_rolling_max(self):
 
@@ -194,6 +203,13 @@ def sdc_pandas_dataframe_rolling_min(self):
     'extra_params': ''
 })
 
+sdc_pandas_dataframe_rolling_kurt.__doc__ = sdc_pandas_dataframe_rolling_docstring_tmpl.format(**{
+    'method_name': 'kurt',
+    'example_caption': 'Calculate unbiased rolling kurtosis.',
+    'limitations_block': '',
+    'extra_params': ''
+})
+
 sdc_pandas_dataframe_rolling_max.__doc__ = sdc_pandas_dataframe_rolling_docstring_tmpl.format(**{
     'method_name': 'max',
     'example_caption': 'Calculate the rolling maximum.',
diff --git a/sdc/tests/test_rolling.py b/sdc/tests/test_rolling.py
@@ -583,6 +583,21 @@ def test_impl(obj, window, min_periods):
                     ref_result = test_impl(obj, window, min_periods)
                     assert_equal(jit_result, ref_result)
 
+    def _test_rolling_kurt(self, obj):
+        def test_impl(obj, window, min_periods):
+            return obj.rolling(window, min_periods).kurt()
+
+        hpat_func = self.jit(test_impl)
+        assert_equal = self._get_assert_equal(obj)
+
+        for window in range(4, len(obj) + 1):
+            for min_periods in range(window + 1):
+                with self.subTest(obj=obj, window=window,
+                                  min_periods=min_periods):
+                    ref_result = test_impl(obj, window, min_periods)
+                    jit_result = hpat_func(obj, window, min_periods)
+                    assert_equal(jit_result, ref_result)
+
     def _test_rolling_max(self, obj):
         def test_impl(obj, window, min_periods):
             return obj.rolling(window, min_periods).max()
@@ -677,6 +692,15 @@ def test_df_rolling_count(self):
 
         self._test_rolling_count(df)
 
+    @skip_sdc_jit('DataFrame.rolling.kurt() unsupported')
+    def test_df_rolling_kurt(self):
+        all_data = test_global_input_data_float64
+        length = min(len(d) for d in all_data)
+        data = {n: d[:length] for n, d in zip(string.ascii_uppercase, all_data)}
+        df = pd.DataFrame(data)
+
+        self._test_rolling_kurt(df)
+
     @skip_sdc_jit('DataFrame.rolling.max() unsupported')
     def test_df_rolling_max(self):
         all_data = test_global_input_data_float64
@@ -899,22 +923,11 @@ def test_impl(pairwise, ddof):
 
     @skip_sdc_jit('Series.rolling.kurt() unsupported Series index')
     def test_series_rolling_kurt(self):
-        def test_impl(series, window, min_periods):
-            return series.rolling(window, min_periods).kurt()
-
-        hpat_func = self.jit(test_impl)
-
         all_data = test_global_input_data_float64
         indices = [list(range(len(data)))[::-1] for data in all_data]
         for data, index in zip(all_data, indices):
             series = pd.Series(data, index, name='A')
-            for window in range(4, len(series) + 1):
-                for min_periods in range(window + 1):
-                    with self.subTest(series=series, window=window,
-                                      min_periods=min_periods):
-                        ref_result = test_impl(series, window, min_periods)
-                        jit_result = hpat_func(series, window, min_periods)
-                        pd.testing.assert_series_equal(jit_result, ref_result)
+            self._test_rolling_kurt(series)
 
     @skip_sdc_jit('Series.rolling.max() unsupported Series index')
     def test_series_rolling_max(self):
diff --git a/sdc/tests/tests_perf/test_perf_df_rolling.py b/sdc/tests/tests_perf/test_perf_df_rolling.py
@@ -86,6 +86,7 @@ def setUpClass(cls):
         cls.total_data_length = {
             'apply': [2 * 10 ** 5],
             'count': [8 * 10 ** 5],
+            'kurt': [4 * 10 ** 5],
             'max': [2 * 10 ** 5],
             'min': [2 * 10 ** 5],
         }
@@ -144,6 +145,9 @@ def test_df_rolling_apply_mean(self):
     def test_df_rolling_count(self):
         self._test_df_rolling_method('count')
 
+    def test_df_rolling_kurt(self):
+        self._test_df_rolling_method('kurt')
+
     def test_df_rolling_max(self):
         self._test_df_rolling_method('max')
 
