Adapters Reference¶

def fit(self, data: pd.DataFrame, max_train_date: pd.Timestamp | None = None) -> None:
    """Fit the Meridian model to data.

    Args:
        data: Training data
        max_train_date: Optional maximum training date for holdout validation

    """
    # Reset state to ensure clean start when new data is provided
    self._reset_state()

    # build Meridian data object
    self.training_data = construct_meridian_data_object(data, self.config)
    self.max_train_date = max_train_date

    model_spec_kwargs = dict(self.config.model_spec_config)

    # if max train date is provided, construct a mask that is True for all dates before max_train_date
    if self.max_train_date:
        self.holdout_mask = construct_holdout_mask(self.max_train_date, self.training_data.kpi.time)
        # model expects a 2D array of shape (n_geos, n_times) so have to duplicate the values across each geo
        model_spec_kwargs["holdout_id"] = np.repeat(
            self.holdout_mask[None, :], repeats=len(self.training_data.kpi.geo), axis=0
        )

    # Create and fit the Meridian model
    model_spec = ModelSpec(**model_spec_kwargs)
    self.model = Meridian(
        input_data=self.training_data,  # type: ignore
        model_spec=model_spec,
    )
    self.trace = self.model.sample_posterior(**dict(self.config.sample_posterior_config))

    # used to compute channel contributions for ROI calculations
    self.analyzer = Analyzer(self.model)
    self.is_fitted = True

def fit_and_predict(self, train: pd.DataFrame, test: pd.DataFrame) -> np.ndarray:
    """Fit the Meridian model and make predictions given new input data.

    The full dataset must be passed to `fit()`, since making out-of-sample predictions
    is only possible by way of specifying a holdout mask when sampling from the
    posterior.

    Args:
        train: Training data
        test: Test data

    Returns:
        Predicted values for the test period

    """
    train_and_test = pd.concat([train, test])
    max_train_date = train[self.date_column].squeeze().max()
    self.fit(train_and_test, max_train_date=max_train_date)
    return self.predict()

def get_channel_names(self) -> list[str]:
    """Get the channel names that would be used as the index in get_channel_roi results.

    For Meridian, this returns the media_channels which are the human-readable
    channel names used in the ROI results.

    Returns
        List of channel names

    """
    return self.media_channels

def get_channel_roi(
    self,
    start_date: pd.Timestamp | None = None,
    end_date: pd.Timestamp | None = None,
) -> pd.Series:
    """Get channel ROI estimates.

    Args:
        start_date: Optional start date for ROI calculation
        end_date: Optional end date for ROI calculation

    Returns:
        Series containing ROI estimates for each channel

    """
    if not self.is_fitted or self.training_data is None or self.analyzer is None:
        raise RuntimeError("Model must be fit before computing ROI")

    # restrict ROI calculation to a particular period if start/end date args are
    # passed
    training_date_index = pd.to_datetime(self.training_data.kpi.time)
    roi_date_index = training_date_index.copy()
    if start_date:
        roi_date_index = roi_date_index[roi_date_index >= start_date]
    if end_date:
        roi_date_index = roi_date_index[roi_date_index < end_date]

    selected_times = [bool(date) for date in training_date_index.isin(roi_date_index)]

    # analyzer.roi() returns a tensor of shape (n_chains, n_draws, n_channels)
    rois_per_channel = np.mean(self.analyzer.roi(selected_times=selected_times), axis=(0, 1))

    rois = {}
    for channel, roi in zip(self.media_channels, rois_per_channel, strict=False):
        rois[channel] = float(roi)
    return pd.Series(rois)

def predict(self, data: pd.DataFrame | None = None) -> np.ndarray:
    """Make predictions using the fitted model.

    This returns predictions for the entirety of the dataset passed to fit() unless
    `max_train_date` is specified when calling fit(); in that case it only returns
    predictions for the time periods indicated by the holdout mask.

    Note: Meridian doesn't require input data for prediction - it uses the fitted
    model state, so the `data` argument will be ignored if passed.

    Args:
        data: Ignored - Meridian uses the fitted model state for predictions.

    Returns:
        Predicted values

    Raises:
        RuntimeError: If model is not fitted

    """
    if not self.is_fitted or self.analyzer is None:
        raise RuntimeError("Model must be fit before prediction")

    # shape (n_chains, n_draws, n_times)
    preds_tensor = self.analyzer.expected_outcome(aggregate_geos=True, aggregate_times=False)
    posterior_mean = np.mean(preds_tensor, axis=(0, 1))

    # if holdout mask is provided, use it to mask the predictions to restrict only to the
    # holdout period
    if self.holdout_mask is not None:
        posterior_mean = posterior_mean[self.holdout_mask]

    return posterior_mean

def fit(self, data: pd.DataFrame) -> None:
    """Fit the model and compute ROIs.

    Args:
        data: DataFrame containing the training data adhering to the PyMCInputDataSchema.

    """
    # Reset to original values at the start of each fit call
    self.channel_spend_columns = self._original_channel_spend_columns.copy()
    self.model_kwargs = self._original_model_kwargs.copy()

    # Identify channel spend columns that sum to zero and remove them from modelling.
    # We cannot reliably make any prediction based on these channels when making
    # predictions on new data.
    channel_spend_data = data[self.channel_spend_columns]
    zero_spend_channels = channel_spend_data.columns[channel_spend_data.sum() == 0].tolist()

    if zero_spend_channels:
        logger.info(f"Dropping channels with zero spend: {zero_spend_channels}")
        # Remove zero-spend channels from the list passed to the MMM constructor
        self.channel_spend_columns = [col for col in self.channel_spend_columns if col not in zero_spend_channels]
        # also update the model config field to reflect the new channel spend columns
        self.model_kwargs["channel_columns"] = self.channel_spend_columns

        # Check for vector priors that might cause shape mismatches
        _check_vector_priors_when_dropping_channels(self.model_kwargs["model_config"], zero_spend_channels)

        data = data.drop(columns=zero_spend_channels)

    X = data.drop(columns=[InputDataframeConstants.RESPONSE_COL, InputDataframeConstants.MEDIA_CHANNEL_REVENUE_COL])
    y = data[InputDataframeConstants.RESPONSE_COL]

    self.model = MMM(**self.model_kwargs)
    self.trace = self.model.fit(X=X, y=y, **self.fit_kwargs)

    self._channel_roi_df = self._compute_channel_contributions(data)
    self.is_fitted = True

def fit_and_predict(self, train: pd.DataFrame, test: pd.DataFrame) -> np.ndarray:
    """Fit on training data and make predictions on test data.

    Arguments:
        train: training dataset
        test: test dataset

    Returns:
        model predictions.

    """
    self.fit(train)
    return self.predict(test)

def get_channel_names(self) -> list[str]:
    """Get the channel names that would be used as the index in get_channel_roi results.

    For PyMC, this returns the `channel_spend_columns` which are used as the index
    in the ROI results.

    Returns
        List of channel names

    """
    return self.channel_spend_columns

def get_channel_roi(
    self,
    start_date: pd.Timestamp | None = None,
    end_date: pd.Timestamp | None = None,
) -> pd.Series:
    """Return the ROIs for each channel, optionally within a given date range.

    Args:
        start_date: Optional start date for ROI calculation
        end_date: Optional end date for ROI calculation

    Returns:
        Series containing ROI estimates for each channel

    """
    if not self.is_fitted or self._channel_roi_df is None:
        raise RuntimeError("Model must be fit before computing ROI.")

    _validate_start_end_dates(start_date, end_date, pd.DatetimeIndex(self._channel_roi_df.index))

    # Filter the contribution DataFrame by date range
    date_range_df = self._channel_roi_df.loc[start_date:end_date]

    if date_range_df.empty:
        raise ValueError(f"No data found for date range {start_date} to {end_date}")

    return pd.Series(self._calculate_rois(date_range_df))

def predict(self, data: pd.DataFrame | None = None) -> np.ndarray:
    """Predict the response variable for new data.

    Args:
        data: Input data for prediction. This parameter is required for PyMC
            predictions and cannot be None.

    Returns:
        Predicted values

    Raises:
        RuntimeError: If model is not fitted
        ValueError: If data is None (PyMC requires data for prediction)

    """
    if not self.is_fitted or self.model is None:
        raise RuntimeError("Model must be fit before prediction.")

    if data is None:
        raise ValueError("PyMC adapter requires data for prediction")

    if InputDataframeConstants.RESPONSE_COL in data.columns:
        data = data.drop(columns=[InputDataframeConstants.RESPONSE_COL])
    predictions = predictions = self.model.predict(
        data, extend_idata=False, include_last_observations=True, **self.predict_kwargs
    )
    return predictions