def aggregate()

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

• 44 lines of code
• 9 McCabe index (conditional complexity)

    def aggregate(self) -> pd.DataFrame:
        """Aggregate the results from predict method

        Args:
            None

        Returns:
            final results in pd.DataFrame
        """
        predicted = self.predicted
        if predicted is None:
            raise _logged_error("predict must be called before aggregate.")

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

        # collect the fcst, fcst_lower, and fcst_upper into dataframes
        fcsts = {}
        for col in ["fcst", "fcst_lower", "fcst_upper"]:
            fcsts[col] = pd.concat(
                [x[col].reset_index(drop=True) for x in predicted.values()], axis=1
            )
            fcsts[col].columns = predicted.keys()

        if self.params["aggregation"].lower() == "median":
            # clean up dataframes with C.I as np.nan or zero
            fcsts = self.clean_dummy_CI(fcsts, use_zero=False)
            self.fcst_df = fcst_df = pd.DataFrame(
                {
                    "time": dates,
                    "fcst": fcsts["fcst"].median(axis=1),
                    "fcst_lower": fcsts["fcst_lower"].median(axis=1),
                    "fcst_upper": fcsts["fcst_upper"].median(axis=1),
                }
            )
        else:
            if (
                fcsts["fcst_lower"].isnull().values.any()
                or fcsts["fcst_upper"].isnull().values.any()
            ):
                msg = "Conf. interval contains NaN, please check individual model."
                raise _logged_error(msg)
            weights = self.weights
            assert weights is not None
            weights = np.array(list(weights.values()))
            self.fcst_df = fcst_df = pd.DataFrame(
                {
                    "time": dates,
                    "fcst": fcsts["fcst"].dot(weights),
                    "fcst_lower": fcsts["fcst_lower"].dot(weights),
                    "fcst_upper": fcsts["fcst_upper"].dot(weights),
                }
            )

        logging.debug("Return forecast data: {fcst_df}".format(fcst_df=fcst_df))
        return fcst_df