def predict()

in kats/models/quadratic_model.py [0:0]

• 42 lines of code
• 4 McCabe index (conditional complexity)

    def predict(self, steps: int, include_history=False, **kwargs) -> pd.DataFrame:
        """predict with fitted quadratic model.

        Args:
            steps: the steps or length of the prediction horizon
            include_history: whether to include the historical data in the prediction

        Returns:
            The predicted dataframe with the following columns:
                `time`, `fcst`, `fcst_lower`, and `fcst_upper`
        """
        past_length = self.past_length
        if past_length is None:
            raise ValueError("Call fit() before predict().")
        model = self.model
        assert model is not None
        logging.debug(
            "Call predict() with parameters. "
            "steps:{steps}, kwargs:{kwargs}".format(steps=steps, kwargs=kwargs)
        )

        self.freq = kwargs.get("freq", "D")
        self.include_history = include_history

        if include_history:
            _X_future = list(range(0, past_length + steps))
        else:
            _X_future = list(range(past_length, past_length + steps))
        self._X_future = _X_future
        _X_fcst = np.column_stack([_X_future, np.power(_X_future, 2)])
        X_fcst = sm.add_constant(_X_fcst)
        y_fcst = model.predict(X_fcst)
        sdev, y_fcst_lower, y_fcst_upper = wls_prediction_std(
            self.model, exog=X_fcst, alpha=self.params.alpha
        )
        self.sdev = sdev
        self.y_fcst = pd.Series(y_fcst)
        self.y_fcst_lower = pd.Series(y_fcst_lower)
        self.y_fcst_upper = pd.Series(y_fcst_upper)

        # create future dates
        last_date = self.data.time.max()
        dates = pd.date_range(start=last_date, periods=steps + 1, freq=self.freq)
        self.dates = dates[dates != last_date]

        if include_history:
            self.dates = np.concatenate((pd.to_datetime(self.data.time), self.dates))

        self.fcst_df = fcst_df = pd.DataFrame(
            {
                "time": self.dates,
                "fcst": self.y_fcst,
                "fcst_lower": self.y_fcst_lower,
                "fcst_upper": self.y_fcst_upper,
            }
        )
        logging.debug("Return forecast data: {fcst_df}".format(fcst_df=self.fcst_df))
        return fcst_df