def predict()

in kats/models/ensemble/kats_ensemble.py [0:0]

• 39 lines of code
• 13 McCabe index (conditional complexity)

    def predict(self, steps: int) -> KatsEnsemble:
        """Predit future for each individual model

        Args:
            steps : number of steps ahead to forecast

        Returns:
            None
        """

        fitted = self.fitted
        if fitted is None:
            raise _logged_error("fit must be called before predict.")

        self.steps = steps
        if self.seasonality:
            sea_data = self.sea_data
            assert sea_data is not None
            # we should pred two types of model
            desea_fitted = {k: v for k, v in fitted.items() if "_smodel" not in k}
            desea_predict = {
                k: v.predict(self.steps).set_index("time")
                for k, v in desea_fitted.items()
            }

            # re-seasonalize
            predicted = KatsEnsemble.reseasonalize(
                sea_data=sea_data,
                desea_predict=desea_predict,
                decomposition_method=self.decomposition_method,
                seasonality_length=self.params["seasonality_length"],
                steps=self.steps,
            )

            # add extra model prediction results from smodels
            fitted_smodel = {k: v for k, v in fitted.items() if "_smodel" in k}
            extra_predict = {
                k: v.predict(self.steps).set_index("time")
                for k, v in fitted_smodel.items()
            }

            predicted.update(extra_predict)
            self.predicted = predicted
        else:
            predicted = {
                k: v.predict(self.steps).set_index("time") for k, v in fitted.items()
            }

            # add dummy C.I if the model doesn't have native C.I
            # this is a hack for median ensemble; everyone model needs to have
            # its native C.I if user choose weighted average ensemble.
            for k, v in predicted.items():
                # if predicted df doesn't have fcst_lower and fcst_upper
                if not {"fcst_lower", "fcst_upper"}.issubset(v.columns):
                    # add dummy C.I
                    tmp_v = copy(v)
                    tmp_v["fcst_lower"] = np.nan
                    tmp_v["fcst_upper"] = np.nan
                    predicted[k] = tmp_v
            self.predicted = predicted
        return self