add quantiles (#8)

Nixtla · Nov 29, 2023 · 230a510 · 230a510
1 parent a0fe601
commit 230a510
Show file tree

Hide file tree

Showing 7 changed files with 492 additions and 17 deletions.
diff --git a/coreforecast/grouped_array.py b/coreforecast/grouped_array.py
@@ -1,6 +1,7 @@
 import ctypes
 import platform
 import sys
+from typing import Union
 
 import numpy as np
 
@@ -59,6 +60,13 @@ def __len__(self):
     def __getitem__(self, i):
         return self.data[self.indptr[i] : self.indptr[i + 1]]
 
+    def _pyfloat_to_c(self, x: float) -> Union[ctypes.c_float, ctypes.c_double]:
+        if self.prefix == "GroupedArrayFloat32":
+            out = ctypes.c_float(x)
+        else:
+            out = ctypes.c_double(x)
+        return out
+
     def scaler_fit(self, scaler_type: str) -> np.ndarray:
         stats = np.full_like(self.data, np.nan, shape=(len(self), 2))
         _LIB[f"{self.prefix}_{scaler_type}ScalerStats"](
@@ -116,6 +124,20 @@ def rolling_transform(
         )
         return out
 
+    def rolling_quantile_transform(
+        self, lag: int, p: float, window_size: int, min_samples: int
+    ) -> np.ndarray:
+        out = np.full_like(self.data, np.nan)
+        _LIB[f"{self.prefix}_RollingQuantileTransform"](
+            self._handle,
+            ctypes.c_int(lag),
+            self._pyfloat_to_c(p),
+            ctypes.c_int(window_size),
+            ctypes.c_int(min_samples),
+            _data_as_void_ptr(out),
+        )
+        return out
+
     def rolling_update(
         self, stat_name: str, lag: int, window_size: int, min_samples: int
     ) -> np.ndarray:
@@ -129,6 +151,20 @@ def rolling_update(
         )
         return out
 
+    def rolling_quantile_update(
+        self, lag: int, p: float, window_size: int, min_samples: int
+    ) -> np.ndarray:
+        out = np.empty_like(self.data, shape=len(self))
+        _LIB[f"{self.prefix}_RollingQuantileUpdate"](
+            self._handle,
+            ctypes.c_int(lag),
+            self._pyfloat_to_c(p),
+            ctypes.c_int(window_size),
+            ctypes.c_int(min_samples),
+            _data_as_void_ptr(out),
+        )
+        return out
+
     def seasonal_rolling_transform(
         self,
         stat_name: str,
@@ -167,6 +203,46 @@ def seasonal_rolling_update(
         )
         return out
 
+    def seasonal_rolling_quantile_transform(
+        self,
+        lag: int,
+        p: float,
+        season_length: int,
+        window_size: int,
+        min_samples: int,
+    ) -> np.ndarray:
+        out = np.full_like(self.data, np.nan)
+        _LIB[f"{self.prefix}_SeasonalRollingQuantileTransform"](
+            self._handle,
+            ctypes.c_int(lag),
+            ctypes.c_int(season_length),
+            self._pyfloat_to_c(p),
+            ctypes.c_int(window_size),
+            ctypes.c_int(min_samples),
+            _data_as_void_ptr(out),
+        )
+        return out
+
+    def seasonal_rolling_quantile_update(
+        self,
+        lag: int,
+        p: float,
+        season_length: int,
+        window_size: int,
+        min_samples: int,
+    ) -> np.ndarray:
+        out = np.empty_like(self.data, shape=len(self))
+        _LIB[f"{self.prefix}_SeasonalRollingQuantileUpdate"](
+            self._handle,
+            ctypes.c_int(lag),
+            ctypes.c_int(season_length),
+            self._pyfloat_to_c(p),
+            ctypes.c_int(window_size),
+            ctypes.c_int(min_samples),
+            _data_as_void_ptr(out),
+        )
+        return out
+
     def expanding_transform_with_aggs(
         self,
         stat_name: str,
@@ -195,21 +271,37 @@ def expanding_transform(
         )
         return out
 
+    def expanding_quantile_transform(self, lag: int, p: float) -> np.ndarray:
+        out = np.full_like(self.data, np.nan)
+        _LIB[f"{self.prefix}_ExpandingQuantileTransform"](
+            self._handle,
+            ctypes.c_int(lag),
+            self._pyfloat_to_c(p),
+            _data_as_void_ptr(out),
+        )
+        return out
+
+    def expanding_quantile_update(self, lag: int, p: float) -> np.ndarray:
+        out = np.empty_like(self.data, shape=len(self))
+        _LIB[f"{self.prefix}_ExpandingQuantileUpdate"](
+            self._handle,
+            ctypes.c_int(lag),
+            self._pyfloat_to_c(p),
+            _data_as_void_ptr(out),
+        )
+        return out
+
     def exponentially_weighted_transform(
         self,
         stat_name: str,
         lag: int,
         alpha: float,
     ) -> np.ndarray:
         out = np.full_like(self.data, np.nan)
-        if self.prefix == "GroupedArrayFloat32":
-            alpha = ctypes.c_float(alpha)
-        else:
-            alpha = ctypes.c_double(alpha)
         _LIB[f"{self.prefix}_ExponentiallyWeighted{stat_name}Transform"](
             self._handle,
             ctypes.c_int(lag),
-            alpha,
+            self._pyfloat_to_c(alpha),
             _data_as_void_ptr(out),
         )
         return out
diff --git a/coreforecast/lag_transforms.py b/coreforecast/lag_transforms.py
@@ -4,14 +4,17 @@
     "RollingStd",
     "RollingMin",
     "RollingMax",
+    "RollingQuantile",
     "SeasonalRollingMean",
     "SeasonalRollingStd",
     "SeasonalRollingMin",
     "SeasonalRollingMax",
+    "SeasonalRollingQuantile",
     "ExpandingMean",
     "ExpandingStd",
     "ExpandingMin",
     "ExpandingMax",
+    "ExpandingQuantile",
     "ExponentiallyWeightedMean",
 ]
 
@@ -86,6 +89,24 @@ class RollingMax(RollingBase):
     stat_name = "Max"
 
 
+class RollingQuantile(RollingBase):
+    def __init__(
+        self, lag: int, p: float, window_size: int, min_samples: Optional[int] = None
+    ):
+        super().__init__(lag=lag, window_size=window_size, min_samples=min_samples)
+        self.p = p
+
+    def transform(self, ga: GroupedArray) -> np.ndarray:
+        return ga.rolling_quantile_transform(
+            self.lag, self.p, self.window_size, self.min_samples
+        )
+
+    def update(self, ga: GroupedArray) -> np.ndarray:
+        return ga.rolling_quantile_update(
+            self.lag - 1, self.p, self.window_size, self.min_samples
+        )
+
+
 class SeasonalRollingBase(RollingBase):
     season_length: int
 
@@ -134,6 +155,42 @@ class SeasonalRollingMax(SeasonalRollingBase):
     stat_name = "Max"
 
 
+class SeasonalRollingQuantile(SeasonalRollingBase):
+    def __init__(
+        self,
+        lag: int,
+        p: float,
+        season_length: int,
+        window_size: int,
+        min_samples: Optional[int] = None,
+    ):
+        super().__init__(
+            lag=lag,
+            season_length=season_length,
+            window_size=window_size,
+            min_samples=min_samples,
+        )
+        self.p = p
+
+    def transform(self, ga: GroupedArray) -> np.ndarray:
+        return ga.seasonal_rolling_quantile_transform(
+            lag=self.lag,
+            p=self.p,
+            season_length=self.season_length,
+            window_size=self.window_size,
+            min_samples=self.min_samples,
+        )
+
+    def update(self, ga: GroupedArray) -> np.ndarray:
+        return ga.seasonal_rolling_quantile_update(
+            lag=self.lag - 1,
+            p=self.p,
+            season_length=self.season_length,
+            window_size=self.window_size,
+            min_samples=self.min_samples,
+        )
+
+
 class ExpandingBase(BaseLagTransform):
     def __init__(self, lag: int):
         self.lag = lag
@@ -193,6 +250,18 @@ class ExpandingMax(ExpandingComp):
     comp_fn = np.maximum
 
 
+class ExpandingQuantile(BaseLagTransform):
+    def __init__(self, lag: int, p: float):
+        self.lag = lag
+        self.p = p
+
+    def transform(self, ga: GroupedArray) -> np.ndarray:
+        return ga.expanding_quantile_transform(self.lag, self.p)
+
+    def update(self, ga: GroupedArray) -> np.ndarray:
+        return ga.expanding_quantile_update(self.lag - 1, self.p)
+
+
 class ExponentiallyWeightedMean(BaseLagTransform):
     def __init__(self, lag: int, alpha: float):
         self.lag = lag

diff --git a/include/coreforecast.h b/include/coreforecast.h
@@ -106,6 +106,15 @@ GroupedArrayFloat64_RollingMaxTransform(GroupedArrayHandle handle, int lag,
                                         int window_size, int min_samples,
                                         double *out);
 
+DLL_EXPORT int
+GroupedArrayFloat32_RollingQuantileTransform(GroupedArrayHandle handle, int lag,
+                                             float p, int window_size,
+                                             int min_samples, float *out);
+DLL_EXPORT int
+GroupedArrayFloat64_RollingQuantileTransform(GroupedArrayHandle handle, int lag,
+                                             double p, int window_size,
+                                             int min_samples, double *out);
+
 DLL_EXPORT int GroupedArrayFloat32_RollingMeanUpdate(GroupedArrayHandle handle,
                                                      int lag, int window_size,
                                                      int min_samples,
@@ -142,6 +151,15 @@ DLL_EXPORT int GroupedArrayFloat64_RollingMaxUpdate(GroupedArrayHandle handle,
                                                     int min_samples,
                                                     double *out);
 
+DLL_EXPORT int
+GroupedArrayFloat32_RollingQuantileUpdate(GroupedArrayHandle handle, int lag,
+                                          float p, int window_size,
+                                          int min_samples, float *out);
+DLL_EXPORT int
+GroupedArrayFloat64_RollingQuantileUpdate(GroupedArrayHandle handle, int lag,
+                                          double p, int window_size,
+                                          int min_samples, double *out);
+
 DLL_EXPORT int GroupedArrayFloat32_SeasonalRollingMeanTransform(
     GroupedArrayHandle handle, int lag, int season_length, int window_size,
     int min_samples, float *out);
@@ -170,6 +188,13 @@ DLL_EXPORT int GroupedArrayFloat64_SeasonalRollingMaxTransform(
     GroupedArrayHandle handle, int lag, int season_length, int window_size,
     int min_samples, double *out);
 
+DLL_EXPORT int GroupedArrayFloat32_SeasonalRollingQuantileTransform(
+    GroupedArrayHandle handle, int lag, int season_length, float p,
+    int window_size, int min_samples, float *out);
+DLL_EXPORT int GroupedArrayFloat64_SeasonalRollingQuantileTransform(
+    GroupedArrayHandle handle, int lag, int season_length, double p,
+    int window_size, int min_samples, double *out);
+
 DLL_EXPORT int GroupedArrayFloat32_SeasonalRollingMeanUpdate(
     GroupedArrayHandle handle, int lag, int season_length, int window_size,
     int min_samples, float *out);
@@ -204,6 +229,13 @@ GroupedArrayFloat64_SeasonalRollingMaxUpdate(GroupedArrayHandle handle, int lag,
                                              int season_length, int window_size,
                                              int min_samples, double *out);
 
+DLL_EXPORT int GroupedArrayFloat32_SeasonalRollingQuantileUpdate(
+    GroupedArrayHandle handle, int lag, int season_length, float p,
+    int window_size, int min_samples, float *out);
+DLL_EXPORT int GroupedArrayFloat64_SeasonalRollingQuantileUpdate(
+    GroupedArrayHandle handle, int lag, int season_length, double p,
+    int window_size, int min_samples, double *out);
+
 DLL_EXPORT int
 GroupedArrayFloat32_ExpandingMeanTransform(GroupedArrayHandle handle, int lag,
                                            float *out, float *agg);
@@ -232,6 +264,20 @@ DLL_EXPORT int
 GroupedArrayFloat64_ExpandingMaxTransform(GroupedArrayHandle handle, int lag,
                                           double *out);
 
+DLL_EXPORT int
+GroupedArrayFloat32_ExpandingQuantileTransform(GroupedArrayHandle handle,
+                                               int lag, float p, float *out);
+DLL_EXPORT int
+GroupedArrayFloat64_ExpandingQuantileTransform(GroupedArrayHandle handle,
+                                               int lag, double p, double *out);
+
+DLL_EXPORT int
+GroupedArrayFloat32_ExpandingQuantileUpdate(GroupedArrayHandle handle, int lag,
+                                            float p, float *out);
+DLL_EXPORT int
+GroupedArrayFloat64_ExpandingQuantileUpdate(GroupedArrayHandle handle, int lag,
+                                            double p, double *out);
+
 DLL_EXPORT int GroupedArrayFloat32_ExponentiallyWeightedMeanTransform(
     GroupedArrayHandle handle, int lag, float alpha, float *out);
 DLL_EXPORT int GroupedArrayFloat64_ExponentiallyWeightedMeanTransform(

diff --git a/pyproject.toml b/pyproject.toml
@@ -43,5 +43,5 @@ wheel.packages = ["coreforecast"]
 wheel.py-api = "py3"
 
 [tool.cibuildwheel]
-test-requires = "pytest window-ops"
+test-requires = "pandas pytest window-ops"
 test-command = "pytest {project}/tests -k correct"
diff --git a/requirements-test.txt b/requirements-test.txt
@@ -1,6 +1,7 @@
 importlib_resources ; python_version < "3.10"
 numba
 numpy
+pandas
 pytest
 pytest-benchmark
 utilsforecast>=0.0.7