added weight parameter (#456)

koaning · Apr 7, 2021 · 907a771 · 907a771
1 parent 0a35875
commit 907a771
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/sklego/preprocessing/repeatingbasis.py b/sklego/preprocessing/repeatingbasis.py
@@ -36,21 +36,25 @@ class RepeatingBasisFunction(TransformerMixin, BaseEstimator):
     :type input_range: tuple or None, default=None
     :param input_range: the values at which the data repeats itself. For example, for days of
         the week this is (1,7). If input_range=None it is inferred from the training data.
+
+    :type width: float, default=1.
+    :param width: determines the width of the radial basis functions.
     """
 
-    def __init__(self, column=0, remainder="drop", n_periods=12, input_range=None):
+    def __init__(self, column=0, remainder="drop", n_periods=12, input_range=None, width=1.):
         self.column = column
         self.remainder = remainder
         self.n_periods = n_periods
         self.input_range = input_range
+        self.width = width
 
     def fit(self, X, y=None):
         self.pipeline_ = ColumnTransformer(
             [
                 (
                     "repeatingbasis",
                     _RepeatingBasisFunction(
-                        n_periods=self.n_periods, input_range=self.input_range
+                        n_periods=self.n_periods, input_range=self.input_range, width=self.width
                     ),
                     [self.column],
                 )
@@ -68,9 +72,10 @@ def transform(self, X):
 
 
 class _RepeatingBasisFunction(TransformerMixin, BaseEstimator):
-    def __init__(self, n_periods: int = 12, input_range=None):
+    def __init__(self, n_periods: int = 12, input_range=None, width: float = 1.):
         self.n_periods = n_periods
         self.input_range = input_range
+        self.width = width
 
     def fit(self, X, y=None):
         X = check_array(X, estimator=self)
@@ -83,7 +88,7 @@ def fit(self, X, y=None):
         self.bases_ = np.linspace(0, 1, self.n_periods + 1)[:-1]
 
         # curves should narrower (wider) when we have more (fewer) basis functions
-        self.width_ = 1 / self.n_periods
+        self.width_ = self.width / self.n_periods
 
         return self