From dd432e87598ff703615409fb5f24641f93589982 Mon Sep 17 00:00:00 2001 From: Kevin Sheppard Date: Tue, 25 Aug 2020 12:59:22 +0100 Subject: [PATCH] ENH: Return self from KDEUnivariate fit Return the fitted instance to simplify chaining --- statsmodels/nonparametric/kde.py | 6 ++++++ statsmodels/nonparametric/tests/test_kde.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/statsmodels/nonparametric/kde.py b/statsmodels/nonparametric/kde.py index 12da35f7455..9cca87cd866 100644 --- a/statsmodels/nonparametric/kde.py +++ b/statsmodels/nonparametric/kde.py @@ -120,6 +120,11 @@ def fit(self, kernel="gau", bw="normal_reference", fft=True, weights=None, -/+ cut*bw*{min(X) or max(X)} adjust : float An adjustment factor for the bw. Bandwidth becomes bw * adjust. + + Returns + ------- + KDEUnivariate + The instance fit, """ try: bw = float(bw) @@ -152,6 +157,7 @@ def fit(self, kernel="gau", bw="normal_reference", fft=True, weights=None, if weights is not None: self.kernel.weights /= weights.sum() self._cache = {} + return self @cache_readonly def cdf(self): diff --git a/statsmodels/nonparametric/tests/test_kde.py b/statsmodels/nonparametric/tests/test_kde.py index f14b26c43fa..2b83af65767 100644 --- a/statsmodels/nonparametric/tests/test_kde.py +++ b/statsmodels/nonparametric/tests/test_kde.py @@ -341,3 +341,10 @@ def test_kde_bw_positive(): kde = KDE(x) kde.fit() assert kde.bw > 0 + + +def test_fit_self(reset_randomstate): + x = np.random.standard_normal(100) + kde = KDE(x) + assert isinstance(kde, KDE) + assert isinstance(kde.fit(), KDE)