Prompt:
# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Plotting functions pre and post model fitting."""

import functools
import logging

# Using these types from typing instead of their generic types in the type hints
# in order to be compatible with Python 3.7 and 3.8.
from typing import Any, List, Optional, Sequence, Tuple

import arviz
import jax
import jax.numpy as jnp
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpyro
import pandas as pd
import seaborn as sns
from sklearn import metrics

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import models
from lightweight_mmm import preprocessing
from lightweight_mmm import utils

plt.style.use("default")

_PALETTE = sns.color_palette(n_colors=100)

@functools.partial(jax.jit, static_argnames=("media_mix_model"))
def _make_single_prediction(media_mix_model: lightweight_mmm.LightweightMMM,
                            mock_media: jnp.ndarray,
                            extra_features: Optional[jnp.ndarray],
                            seed: Optional[int]
                            ) -> jnp.ndarray:
  """Makes a prediction of a single row.

  Serves as a helper function for making predictions individually for each media
  channel and one row at a time. It is meant to be used vmaped otherwise it can
  be slow as it's meant to be used for plotting curve responses only. Use
  lightweight_mmm.LightweightMMM for regular predict functionality.

  Args:
    media_mix_model: Media mix model to use for getting the predictions.
    mock_media: Mock media for this iteration of predictions.
    extra_features: Extra features to use for predictions.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    A point estimate for the given data.
  """
  return media_mix_model.predict(
      media=jnp.expand_dims(mock_media, axis=0),
      extra_features=extra_features,
      seed=seed).mean(axis=0)

@functools.partial(
    jax.jit,
    static_argnames=("media_mix_model", "target_scaler"))
def _generate_diagonal_predictions(
    media_mix_model: lightweight_mmm.LightweightMMM,
    media_values: jnp.ndarray,
    extra_features: Optional[jnp.ndarray],
    target_scaler: Optional[preprocessing.CustomScaler],
    prediction_offset: jnp.ndarray,
    seed: Optional[int]):
  """Generates predictions for one value per channel leaving the rest to zero.

  This function does the following steps:
    - Vmaps the single prediction function on axis=0 of the media arg.
    - Diagonalizes the media input values so that each value is represented
      along side zeros on for the rest of the channels.
    - Generate predictions.
    - Unscale prediction if target_scaler is given.

  Args:
    media_mix_model: Media mix model to use for plotting the response curves.
    media_values: Media values.
    extra_features: Extra features values.
    target_scaler: Scaler used for scaling the target, to unscaled values and
      plot in the original scale.
    prediction_offset: The value of a prediction of an all zero media input.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    The predictions for the given data.
  """
  make_predictions = jax.vmap(fun=_make_single_prediction,
                              in_axes=(None, 0, None, None))
  diagonal = jnp.eye(media_values.shape[0])
  if media_values.ndim == 2:  # Only two since we only provide one row
    diagonal = jnp.expand_dims(diagonal, axis=-1)
    media_values = jnp.expand_dims(media_values, axis=0)
  diag_media_values = diagonal * media_values
  predictions = make_predictions(
      media_mix_model,
      diag_media_values,
      extra_features,
      seed) - prediction_offset
  predictions = jnp.squeeze(predictions)
  if target_scaler:
    predictions = target_scaler.inverse_transform(predictions)
  if predictions.ndim == 2:
    predictions = jnp.sum(predictions, axis=-1)
  return predictions

def _calculate_number_rows_plot(n_media_channels: int, n_columns: int):
  """Calculates the number of rows of plots needed to fit n + 1 plots in n_cols.

  Args:
    n_media_channels: Number of media channels. The total of plots needed is
      n_media_channels + 1.
    n_columns: Number of columns in the plot grid.

  Returns:
    The number of rows of plots needed to fit n + 1 plots in n cols
  """

  if n_media_channels % n_columns == 0:
    return n_media_channels // n_columns + 1
  return n_media_channels // n_columns + 2


Prompt:
# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Plotting functions pre and post model fitting."""

import functools
import logging

# Using these types from typing instead of their generic types in the type hints
# in order to be compatible with Python 3.7 and 3.8.
from typing import Any, List, Optional, Sequence, Tuple

import arviz
import jax
import jax.numpy as jnp
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpyro
import pandas as pd
import seaborn as sns
from sklearn import metrics

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import models
from lightweight_mmm import preprocessing
from lightweight_mmm import utils

plt.style.use("default")

_PALETTE = sns.color_palette(n_colors=100)

@functools.partial(jax.jit, static_argnames=("media_mix_model"))
def _make_single_prediction(media_mix_model: lightweight_mmm.LightweightMMM,
                            mock_media: jnp.ndarray,
                            extra_features: Optional[jnp.ndarray],
                            seed: Optional[int]
                            ) -> jnp.ndarray:
  """Makes a prediction of a single row.

  Serves as a helper function for making predictions individually for each media
  channel and one row at a time. It is meant to be used vmaped otherwise it can
  be slow as it's meant to be used for plotting curve responses only. Use
  lightweight_mmm.LightweightMMM for regular predict functionality.

  Args:
    media_mix_model: Media mix model to use for getting the predictions.
    mock_media: Mock media for this iteration of predictions.
    extra_features: Extra features to use for predictions.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    A point estimate for the given data.
  """
  return media_mix_model.predict(
      media=jnp.expand_dims(mock_media, axis=0),
      extra_features=extra_features,
      seed=seed).mean(axis=0)

@functools.partial(
    jax.jit,
    static_argnames=("media_mix_model", "target_scaler"))
def _generate_diagonal_predictions(
    media_mix_model: lightweight_mmm.LightweightMMM,
    media_values: jnp.ndarray,
    extra_features: Optional[jnp.ndarray],
    target_scaler: Optional[preprocessing.CustomScaler],
    prediction_offset: jnp.ndarray,
    seed: Optional[int]):
  """Generates predictions for one value per channel leaving the rest to zero.

  This function does the following steps:
    - Vmaps the single prediction function on axis=0 of the media arg.
    - Diagonalizes the media input values so that each value is represented
      along side zeros on for the rest of the channels.
    - Generate predictions.
    - Unscale prediction if target_scaler is given.

  Args:
    media_mix_model: Media mix model to use for plotting the response curves.
    media_values: Media values.
    extra_features: Extra features values.
    target_scaler: Scaler used for scaling the target, to unscaled values and
      plot in the original scale.
    prediction_offset: The value of a prediction of an all zero media input.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    The predictions for the given data.
  """
  make_predictions = jax.vmap(fun=_make_single_prediction,
                              in_axes=(None, 0, None, None))
  diagonal = jnp.eye(media_values.shape[0])
  if media_values.ndim == 2:  # Only two since we only provide one row
    diagonal = jnp.expand_dims(diagonal, axis=-1)
    media_values = jnp.expand_dims(media_values, axis=0)
  diag_media_values = diagonal * media_values
  predictions = make_predictions(
      media_mix_model,
      diag_media_values,
      extra_features,
      seed) - prediction_offset
  predictions = jnp.squeeze(predictions)
  if target_scaler:
    predictions = target_scaler.inverse_transform(predictions)
  if predictions.ndim == 2:
    predictions = jnp.sum(predictions, axis=-1)
  return predictions

def _calculate_number_rows_plot(n_media_channels: int, n_columns: int):
  """Calculates the number of rows of plots needed to fit n + 1 plots in n_cols.

  Args:
    n_media_channels: Number of media channels. The total of plots needed is
      n_media_channels + 1.
    n_columns: Number of columns in the plot grid.

  Returns:
    The number of rows of plots needed to fit n + 1 plots in n cols
  """
  if n_media_channels % n_columns == 0:
    return n_media_channels // n_columns + 1
  return n_media_channels // n_columns + 2

def _calculate_media_contribution(
    media_mix_model: lightweight_mmm.LightweightMMM) -> jnp.ndarray:
  """Computes contribution for each sample, time, channel.

  Serves as a helper function for making predictions for each channel, time
  and estimate sample. It is meant to be used in creating media baseline
  contribution dataframe and visualize media attribution over spend proportion
  plot.

  Args:
    media_mix_model: Media mix model.

  Returns:
    Estimation of contribution for each sample, time, channel.

  Raises:
    NotFittedModelError: if the model is not fitted before computation
  """

  if not hasattr(media_mix_model, "trace"):
    raise lightweight_mmm.NotFittedModelError(
        "Model needs to be fit first before attempting to plot its fit.")

  if media_mix_model.trace["media_transformed"].ndim > 3:
    # s for samples, t for time, c for media channels, g for geo
    einsum_str = "stcg, scg->stcg"
  elif media_mix_model.trace["media_transformed"].ndim == 3:
    # s for samples, t for time, c for media channels
    einsum_str = "stc, sc->stc"

  media_contribution = jnp.einsum(einsum_str,
                                  media_mix_model.trace["media_transformed"],
                                  media_mix_model.trace["coef_media"])
  if media_mix_model.trace["media_transformed"].ndim > 3:
    # Aggregate media channel contribution across geos.
    media_contribution = media_contribution.sum(axis=-1)
  return media_contribution


Prompt:
  # Create media contribution matrix.
  scaled_media_contribution = _calculate_media_contribution(media_mix_model)

  # Aggregate media channel contribution across samples.
  sum_scaled_media_contribution_across_samples = scaled_media_contribution.sum(
      axis=0)
  # Aggregate media channel contribution across channels.
  sum_scaled_media_contribution_across_channels = scaled_media_contribution.sum(
      axis=2)

  # Calculate the baseline contribution.
  # Scaled prediction - sum of scaled contribution across channels.
  scaled_prediction = media_mix_model.trace["mu"]
  if media_mix_model.trace["media_transformed"].ndim > 3:
    # Sum up the scaled prediction across all the geos.
    scaled_prediction = scaled_prediction.sum(axis=-1)
  baseline_contribution = scaled_prediction - sum_scaled_media_contribution_across_channels

  # Sum up the scaled media, baseline contribution and predictio across samples.
  sum_scaled_media_contribution_across_channels_samples = sum_scaled_media_contribution_across_channels.sum(
      axis=0)
  sum_scaled_baseline_contribution_across_samples = baseline_contribution.sum(
      axis=0)

  # Adjust baseline contribution and prediction when there's any negative value.
  adjusted_sum_scaled_baseline_contribution_across_samples = np.where(
      sum_scaled_baseline_contribution_across_samples < 0, 0,
      sum_scaled_baseline_contribution_across_samples)
  adjusted_sum_scaled_prediction_across_samples = adjusted_sum_scaled_baseline_contribution_across_samples + sum_scaled_media_contribution_across_channels_samples

  # Calculate the media and baseline pct.
  # Media/baseline contribution across samples/total prediction across samples.
  media_contribution_pct_by_channel = (
      sum_scaled_media_contribution_across_samples /
      adjusted_sum_scaled_prediction_across_samples.reshape(-1, 1))
  # Adjust media pct contribution if the value is nan
  media_contribution_pct_by_channel = np.nan_to_num(
      media_contribution_pct_by_channel)

  baseline_contribution_pct = adjusted_sum_scaled_baseline_contribution_across_samples / adjusted_sum_scaled_prediction_across_samples
  # Adjust baseline pct contribution if the value is nan
  baseline_contribution_pct = np.nan_to_num(
      baseline_contribution_pct)

  # If the channel_names is none, then create naming covention for the channels.
  if channel_names is None:
    channel_names = media_mix_model.media_names

  # Create media/baseline contribution pct as dataframes.
  media_contribution_pct_by_channel_df = pd.DataFrame(
      media_contribution_pct_by_channel, columns=channel_names)
  baseline_contribution_pct_df = pd.DataFrame(
      baseline_contribution_pct, columns=["baseline"])
  contribution_pct_df = pd.merge(
      media_contribution_pct_by_channel_df,
      baseline_contribution_pct_df,
      left_index=True,
      right_index=True)

  # If there's target scaler then inverse transform the posterior prediction.
  posterior_pred = media_mix_model.trace["mu"]
  if target_scaler:
    posterior_pred = target_scaler.inverse_transform(posterior_pred)

  # Take the sum of posterior predictions across geos.
  if media_mix_model.trace["media_transformed"].ndim > 3:
    posterior_pred = posterior_pred.sum(axis=-1)

  # Take the average of the inverse transformed prediction across samples.
  posterior_pred_df = pd.DataFrame(
      posterior_pred.mean(axis=0), columns=["avg_prediction"])

  # Adjust prediction value when prediction is less than 0.
  posterior_pred_df["avg_prediction"] = np.where(
      posterior_pred_df["avg_prediction"] < 0, 0,
      posterior_pred_df["avg_prediction"])

  contribution_pct_df.columns = [
      "{}_percentage".format(col) for col in contribution_pct_df.columns
  ]
  contribution_df = pd.merge(
      contribution_pct_df, posterior_pred_df, left_index=True, right_index=True)

  # Create contribution by multiplying average prediction by media/baseline pct.
  for channel in channel_names:
    channel_contribution_col_name = "{} contribution".format(channel)
    channel_pct_col = "{}_percentage".format(channel)
    contribution_df.loc[:, channel_contribution_col_name] = contribution_df[
        channel_pct_col] * contribution_df["avg_prediction"]
    contribution_df.loc[:, channel_contribution_col_name] = contribution_df[
        channel_contribution_col_name].astype("float")
  contribution_df.loc[:, "baseline contribution"] = contribution_df[
      "baseline_percentage"] * contribution_df["avg_prediction"]

  period = np.arange(1, contribution_df.shape[0] + 1)
  contribution_df.loc[:, "period"] = period
  return contribution_df

def plot_response_curves(# jax-ndarray
    media_mix_model: lightweight_mmm.LightweightMMM,
    media_scaler: Optional[preprocessing.CustomScaler] = None,
    target_scaler: Optional[preprocessing.CustomScaler] = None,
    prices: jnp.ndarray = None,
    optimal_allocation_per_timeunit: Optional[jnp.ndarray] = None,
    steps: int = 50,
    percentage_add: float = 0.2,
    apply_log_scale: bool = False,
    figure_size: Tuple[int, int] = (8, 10),
    n_columns: int = 3,
    marker_size: int = 8,
    legend_fontsize: int = 8,
    seed: Optional[int] = None) -> matplotlib.figure.Figure:
  """Plots the response curves of each media channel based on the model.

  It plots an individual subplot for each media channel. If '
  optimal_allocation_per_timeunit is given it uses it to add markers based on
  historic average spend and the given optimal one on each of the individual
  subplots.

  It then plots a combined plot with all the response curves which can be
  changed to log scale if apply_log_scale is True.

  Args:
    media_mix_model: Media mix model to use for plotting the response curves.
    media_scaler: Scaler that was used to scale the media data before training.
    target_scaler: Scaler used for scaling the target, to unscaled values and
      plot in the original scale.
    prices: Prices to translate the media units to spend. If all your data is
      already in spend numbers you can leave this as None. If some of your data
      is media spend and others is media unit, leave the media spend with price
      1 and add the price to the media unit channels.
    optimal_allocation_per_timeunit: Optimal allocation per time unit per media
      channel. This can be obtained by running the optimization provided by
      LightweightMMM.
    steps: Number of steps to simulate.
    percentage_add: Percentage too exceed the maximum historic spend for the
      simulation of the response curve.
    apply_log_scale: Whether to apply the log scale to the predictions (Y axis).
      When some media channels have very large scale compare to others it might
      be useful to use apply_log_scale=True. Default is False.
    figure_size: Size of the plot figure.
    n_columns: Number of columns to display in the subplots grid. Modifying this
      parameter might require to adjust figure_size accordingly for the plot
      to still have reasonable structure.
    marker_size: Size of the marker for the optimization annotations. Only
      useful if optimal_allocation_per_timeunit is not None. Default is 8.
    legend_fontsize: Legend font size for individual subplots.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    Plots of response curves.
  """
  if not hasattr(media_mix_model, "trace"):
    raise lightweight_mmm.NotFittedModelError(
        "Model needs to be fit first before attempting to plot its response "
        "curves.")
  media = media_mix_model.media
  media_maxes = media.max(axis=0) * (1 + percentage_add)
  if media_mix_model._extra_features is not None:
    extra_features = jnp.expand_dims(
        media_mix_model._extra_features.mean(axis=0), axis=0)
  else:
    extra_features = None
  media_ranges = jnp.expand_dims(
      jnp.linspace(start=0, stop=media_maxes, num=steps), axis=0)

  make_predictions = jax.vmap(
      jax.vmap(_make_single_prediction,
               in_axes=(None, 0, None, None),
               out_axes=0),
      in_axes=(None, 0, None, None), out_axes=1)
  diagonal = jnp.repeat(
      jnp.eye(media_mix_model.n_media_channels), steps,
      axis=0).reshape(media_mix_model.n_media_channels, steps,
                      media_mix_model.n_media_channels)

  prediction_offset = media_mix_model.predict(
      media=jnp.zeros((1, *media.shape[1:])),
      extra_features=extra_features).mean(axis=0)

  if media.ndim == 3:
    diagonal = jnp.expand_dims(diagonal, axis=-1)
    prediction_offset = jnp.expand_dims(prediction_offset, axis=0)
  mock_media = media_ranges * diagonal
  predictions = jnp.squeeze(a=make_predictions(media_mix_model,
                                               mock_media,
                                               extra_features,
                                               seed))
  predictions = predictions - prediction_offset
  media_ranges = jnp.squeeze(media_ranges)
  if target_scaler:
    predictions = target_scaler.inverse_transform(predictions)

  if media_scaler:
    media_ranges = media_scaler.inverse_transform(media_ranges)

  if prices is not None:
    if media.ndim == 3:
      prices = jnp.expand_dims(prices, axis=-1)
    media_ranges *= prices

  if predictions.ndim == 3:
    media_ranges = jnp.sum(media_ranges, axis=-1)
    predictions = jnp.sum(predictions, axis=-1)

  if optimal_allocation_per_timeunit is not None:
    average_allocation = media_mix_model.media.mean(axis=0)
    average_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=average_allocation,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    optimal_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=optimal_allocation_per_timeunit,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    if media_scaler:
      average_allocation = media_scaler.inverse_transform(average_allocation)
      optimal_allocation_per_timeunit = media_scaler.inverse_transform(
          optimal_allocation_per_timeunit)
    if prices is not None:
      optimal_allocation_per_timeunit *= prices
      average_allocation *= prices
    if media.ndim == 3:
      average_allocation = jnp.sum(average_allocation, axis=-1)
      optimal_allocation_per_timeunit = jnp.sum(
          optimal_allocation_per_timeunit, axis=-1)

  kpi_label = "KPI" if target_scaler else "Normalized KPI"
  fig = plt.figure(media_mix_model.n_media_channels + 1,
                   figsize=figure_size,
                   tight_layout=True)
  n_rows = _calculate_number_rows_plot(
      n_media_channels=media_mix_model.n_media_channels, n_columns=n_columns)
  last_ax = fig.add_subplot(n_rows, 1, n_rows)
  for i in range(media_mix_model.n_media_channels):
    ax = fig.add_subplot(n_rows, n_columns, i + 1)
    sns.lineplot(
        x=media_ranges[:, i],
        y=predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=ax)
    sns.lineplot(
        x=media_ranges[:, i],
        y=jnp.log(predictions[:, i]) if apply_log_scale else predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=last_ax)
    if optimal_allocation_per_timeunit is not None:
      ax.plot(
          average_allocation[i],
          average_allocation_predictions[i],
          marker="o",
          markersize=marker_size,
          label="avg_spend",
          color=_PALETTE[i])
      ax.plot(
          optimal_allocation_per_timeunit[i],
          optimal_allocation_predictions[i],
          marker="x",
          markersize=marker_size + 2,
          label="optimal_spend",
          color=_PALETTE[i])
    ax.set_ylabel(kpi_label)
    ax.set_xlabel("Normalized Spend" if not media_scaler else "Spend")
    ax.legend(fontsize=legend_fontsize)

  fig.suptitle("Response curves", fontsize=20)
  last_ax.set_ylabel(kpi_label if not apply_log_scale else f"log({kpi_label})")
  last_ax.set_xlabel("Normalized spend per channel"
                     if not media_scaler else "Spend per channel")
  plt.close()
  return fig

def plot_cross_correlate(feature: jnp.ndarray,
                         target: jnp.ndarray,
                         maxlags: int = 10) -> Tuple[int, float]:
  """Plots the cross correlation coefficients between 2 vectors.

  In the chart look for positive peaks, this shows how the lags of the feature
  lead the target.

  Args:
    feature: Vector, the lags of which predict target.
    target: Vector, what is predicted.
    maxlags: Maximum number of lags.

  Returns:
    Lag index and corresponding correlation of the peak correlation.

  Raises:
    ValueError: If inputs don't have same length.
  """

  if len(feature) != len(target):
    raise ValueError("feature and target need to have the same length.")
  maxlags = jnp.minimum(len(feature) - 1, maxlags)
  mean_feature, mean_target = feature.mean(), target.mean()
  plot = plt.xcorr(
      x=feature - mean_feature, y=target - mean_target, maxlags=maxlags)
  plt.show()
  maxidx = plot[1][plot[0] <= 0].argmax()
  return plot[0][maxidx], plot[1][maxidx]


Prompt:
_channels_samples = sum_scaled_media_contribution_across_channels.sum(
      axis=0)
  sum_scaled_baseline_contribution_across_samples = baseline_contribution.sum(
      axis=0)

  # Adjust baseline contribution and prediction when there's any negative value.
  adjusted_sum_scaled_baseline_contribution_across_samples = np.where(
      sum_scaled_baseline_contribution_across_samples < 0, 0,
      sum_scaled_baseline_contribution_across_samples)
  adjusted_sum_scaled_prediction_across_samples = adjusted_sum_scaled_baseline_contribution_across_samples + sum_scaled_media_contribution_across_channels_samples

  # Calculate the media and baseline pct.
  # Media/baseline contribution across samples/total prediction across samples.
  media_contribution_pct_by_channel = (
      sum_scaled_media_contribution_across_samples /
      adjusted_sum_scaled_prediction_across_samples.reshape(-1, 1))
  # Adjust media pct contribution if the value is nan
  media_contribution_pct_by_channel = np.nan_to_num(
      media_contribution_pct_by_channel)

  baseline_contribution_pct = adjusted_sum_scaled_baseline_contribution_across_samples / adjusted_sum_scaled_prediction_across_samples
  # Adjust baseline pct contribution if the value is nan
  baseline_contribution_pct = np.nan_to_num(
      baseline_contribution_pct)

  # If the channel_names is none, then create naming covention for the channels.
  if channel_names is None:
    channel_names = media_mix_model.media_names

  # Create media/baseline contribution pct as dataframes.
  media_contribution_pct_by_channel_df = pd.DataFrame(
      media_contribution_pct_by_channel, columns=channel_names)
  baseline_contribution_pct_df = pd.DataFrame(
      baseline_contribution_pct, columns=["baseline"])
  contribution_pct_df = pd.merge(
      media_contribution_pct_by_channel_df,
      baseline_contribution_pct_df,
      left_index=True,
      right_index=True)

  # If there's target scaler then inverse transform the posterior prediction.
  posterior_pred = media_mix_model.trace["mu"]
  if target_scaler:
    posterior_pred = target_scaler.inverse_transform(posterior_pred)

  # Take the sum of posterior predictions across geos.
  if media_mix_model.trace["media_transformed"].ndim > 3:
    posterior_pred = posterior_pred.sum(axis=-1)

  # Take the average of the inverse transformed prediction across samples.
  posterior_pred_df = pd.DataFrame(
      posterior_pred.mean(axis=0), columns=["avg_prediction"])

  # Adjust prediction value when prediction is less than 0.
  posterior_pred_df["avg_prediction"] = np.where(
      posterior_pred_df["avg_prediction"] < 0, 0,
      posterior_pred_df["avg_prediction"])

  contribution_pct_df.columns = [
      "{}_percentage".format(col) for col in contribution_pct_df.columns
  ]
  contribution_df = pd.merge(
      contribution_pct_df, posterior_pred_df, left_index=True, right_index=True)

  # Create contribution by multiplying average prediction by media/baseline pct.
  for channel in channel_names:
    channel_contribution_col_name = "{} contribution".format(channel)
    channel_pct_col = "{}_percentage".format(channel)
    contribution_df.loc[:, channel_contribution_col_name] = contribution_df[
        channel_pct_col] * contribution_df["avg_prediction"]
    contribution_df.loc[:, channel_contribution_col_name] = contribution_df[
        channel_contribution_col_name].astype("float")
  contribution_df.loc[:, "baseline contribution"] = contribution_df[
      "baseline_percentage"] * contribution_df["avg_prediction"]

  period = np.arange(1, contribution_df.shape[0] + 1)
  contribution_df.loc[:, "period"] = period
  return contribution_df

def plot_response_curves(# jax-ndarray
    media_mix_model: lightweight_mmm.LightweightMMM,
    media_scaler: Optional[preprocessing.CustomScaler] = None,
    target_scaler: Optional[preprocessing.CustomScaler] = None,
    prices: jnp.ndarray = None,
    optimal_allocation_per_timeunit: Optional[jnp.ndarray] = None,
    steps: int = 50,
    percentage_add: float = 0.2,
    apply_log_scale: bool = False,
    figure_size: Tuple[int, int] = (8, 10),
    n_columns: int = 3,
    marker_size: int = 8,
    legend_fontsize: int = 8,
    seed: Optional[int] = None) -> matplotlib.figure.Figure:
  """Plots the response curves of each media channel based on the model.

  It plots an individual subplot for each media channel. If '
  optimal_allocation_per_timeunit is given it uses it to add markers based on
  historic average spend and the given optimal one on each of the individual
  subplots.

  It then plots a combined plot with all the response curves which can be
  changed to log scale if apply_log_scale is True.

  Args:
    media_mix_model: Media mix model to use for plotting the response curves.
    media_scaler: Scaler that was used to scale the media data before training.
    target_scaler: Scaler used for scaling the target, to unscaled values and
      plot in the original scale.
    prices: Prices to translate the media units to spend. If all your data is
      already in spend numbers you can leave this as None. If some of your data
      is media spend and others is media unit, leave the media spend with price
      1 and add the price to the media unit channels.
    optimal_allocation_per_timeunit: Optimal allocation per time unit per media
      channel. This can be obtained by running the optimization provided by
      LightweightMMM.
    steps: Number of steps to simulate.
    percentage_add: Percentage too exceed the maximum historic spend for the
      simulation of the response curve.
    apply_log_scale: Whether to apply the log scale to the predictions (Y axis).
      When some media channels have very large scale compare to others it might
      be useful to use apply_log_scale=True. Default is False.
    figure_size: Size of the plot figure.
    n_columns: Number of columns to display in the subplots grid. Modifying this
      parameter might require to adjust figure_size accordingly for the plot
      to still have reasonable structure.
    marker_size: Size of the marker for the optimization annotations. Only
      useful if optimal_allocation_per_timeunit is not None. Default is 8.
    legend_fontsize: Legend font size for individual subplots.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    Plots of response curves.
  """
  if not hasattr(media_mix_model, "trace"):
    raise lightweight_mmm.NotFittedModelError(
        "Model needs to be fit first before attempting to plot its response "
        "curves.")
  media = media_mix_model.media
  media_maxes = media.max(axis=0) * (1 + percentage_add)
  if media_mix_model._extra_features is not None:
    extra_features = jnp.expand_dims(
        media_mix_model._extra_features.mean(axis=0), axis=0)
  else:
    extra_features = None
  media_ranges = jnp.expand_dims(
      jnp.linspace(start=0, stop=media_maxes, num=steps), axis=0)

  make_predictions = jax.vmap(
      jax.vmap(_make_single_prediction,
               in_axes=(None, 0, None, None),
               out_axes=0),
      in_axes=(None, 0, None, None), out_axes=1)
  diagonal = jnp.repeat(
      jnp.eye(media_mix_model.n_media_channels), steps,
      axis=0).reshape(media_mix_model.n_media_channels, steps,
                      media_mix_model.n_media_channels)

  prediction_offset = media_mix_model.predict(
      media=jnp.zeros((1, *media.shape[1:])),
      extra_features=extra_features).mean(axis=0)

  if media.ndim == 3:
    diagonal = jnp.expand_dims(diagonal, axis=-1)
    prediction_offset = jnp.expand_dims(prediction_offset, axis=0)
  mock_media = media_ranges * diagonal
  predictions = jnp.squeeze(a=make_predictions(media_mix_model,
                                               mock_media,
                                               extra_features,
                                               seed))
  predictions = predictions - prediction_offset
  media_ranges = jnp.squeeze(media_ranges)
  if target_scaler:
    predictions = target_scaler.inverse_transform(predictions)

  if media_scaler:
    media_ranges = media_scaler.inverse_transform(media_ranges)

  if prices is not None:
    if media.ndim == 3:
      prices = jnp.expand_dims(prices, axis=-1)
    media_ranges *= prices

  if predictions.ndim == 3:
    media_ranges = jnp.sum(media_ranges, axis=-1)
    predictions = jnp.sum(predictions, axis=-1)

  if optimal_allocation_per_timeunit is not None:
    average_allocation = media_mix_model.media.mean(axis=0)
    average_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=average_allocation,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    optimal_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=optimal_allocation_per_timeunit,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    if media_scaler:
      average_allocation = media_scaler.inverse_transform(average_allocation)
      optimal_allocation_per_timeunit = media_scaler.inverse_transform(
          optimal_allocation_per_timeunit)
    if prices is not None:
      optimal_allocation_per_timeunit *= prices
      average_allocation *= prices
    if media.ndim == 3:
      average_allocation = jnp.sum(average_allocation, axis=-1)
      optimal_allocation_per_timeunit = jnp.sum(
          optimal_allocation_per_timeunit, axis=-1)

  kpi_label = "KPI" if target_scaler else "Normalized KPI"
  fig = plt.figure(media_mix_model.n_media_channels + 1,
                   figsize=figure_size,
                   tight_layout=True)
  n_rows = _calculate_number_rows_plot(
      n_media_channels=media_mix_model.n_media_channels, n_columns=n_columns)
  last_ax = fig.add_subplot(n_rows, 1, n_rows)
  for i in range(media_mix_model.n_media_channels):
    ax = fig.add_subplot(n_rows, n_columns, i + 1)
    sns.lineplot(
        x=media_ranges[:, i],
        y=predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=ax)
    sns.lineplot(
        x=media_ranges[:, i],
        y=jnp.log(predictions[:, i]) if apply_log_scale else predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=last_ax)
    if optimal_allocation_per_timeunit is not None:
      ax.plot(
          average_allocation[i],
          average_allocation_predictions[i],
          marker="o",
          markersize=marker_size,
          label="avg_spend",
          color=_PALETTE[i])
      ax.plot(
          optimal_allocation_per_timeunit[i],
          optimal_allocation_predictions[i],
          marker="x",
          markersize=marker_size + 2,
          label="optimal_spend",
          color=_PALETTE[i])
    ax.set_ylabel(kpi_label)
    ax.set_xlabel("Normalized Spend" if not media_scaler else "Spend")
    ax.legend(fontsize=legend_fontsize)

  fig.suptitle("Response curves", fontsize=20)
  last_ax.set_ylabel(kpi_label if not apply_log_scale else f"log({kpi_label})")
  last_ax.set_xlabel("Normalized spend per channel"
                     if not media_scaler else "Spend per channel")
  plt.close()
  return fig

def plot_cross_correlate(feature: jnp.ndarray,
                         target: jnp.ndarray,
                         maxlags: int = 10) -> Tuple[int, float]:
  """Plots the cross correlation coefficients between 2 vectors.

  In the chart look for positive peaks, this shows how the lags of the feature
  lead the target.

  Args:
    feature: Vector, the lags of which predict target.
    target: Vector, what is predicted.
    maxlags: Maximum number of lags.

  Returns:
    Lag index and corresponding correlation of the peak correlation.

  Raises:
    ValueError: If inputs don't have same length.
  """
  if len(feature) != len(target):
    raise ValueError("feature and target need to have the same length.")
  maxlags = jnp.minimum(len(feature) - 1, maxlags)
  mean_feature, mean_target = feature.mean(), target.mean()
  plot = plt.xcorr(
      x=feature - mean_feature, y=target - mean_target, maxlags=maxlags)
  plt.show()
  maxidx = plot[1][plot[0] <= 0].argmax()
  return plot[0][maxidx], plot[1][maxidx]

def plot_var_cost(media: jnp.ndarray, costs: jnp.ndarray,
                  names: List[str]) -> matplotlib.figure.Figure:
  """Plots a a chart between the coefficient of variation and cost.

  Args:
    media: Media matrix.
    costs: Cost vector.
    names: List of variable names.

  Returns:
    Plot of coefficient of variation and cost.

  Raises:
    ValueError if inputs don't conform to same length.
  """

  if media.shape[1] != len(costs):
    raise ValueError("media columns and costs needs to have same length.")
  if media.shape[1] != len(names):
    raise ValueError("media columns and names needs to have same length.")
  coef_of_variation = media.std(axis=0) / media.mean(axis=0)

  fig, ax = plt.subplots(1, 1)
  ax.scatter(x=costs, y=coef_of_variation)
  # https://queirozf.com/entries/add-labels-and-text-to-matplotlib-plots-annotation-examples.
  for i in range(len(costs)):
    x, y, label = costs[i], coef_of_variation[i], names[i]
    ax.annotate(text=label, xy=(x, y))
  ax.set_xlabel("Cost")
  ax.set_ylabel("Coef of Variation")
  plt.close()
  return fig


Prompt:
aler] = None,
    target_scaler: Optional[preprocessing.CustomScaler] = None,
    prices: jnp.ndarray = None,
    optimal_allocation_per_timeunit: Optional[jnp.ndarray] = None,
    steps: int = 50,
    percentage_add: float = 0.2,
    apply_log_scale: bool = False,
    figure_size: Tuple[int, int] = (8, 10),
    n_columns: int = 3,
    marker_size: int = 8,
    legend_fontsize: int = 8,
    seed: Optional[int] = None) -> matplotlib.figure.Figure:
  """Plots the response curves of each media channel based on the model.

  It plots an individual subplot for each media channel. If '
  optimal_allocation_per_timeunit is given it uses it to add markers based on
  historic average spend and the given optimal one on each of the individual
  subplots.

  It then plots a combined plot with all the response curves which can be
  changed to log scale if apply_log_scale is True.

  Args:
    media_mix_model: Media mix model to use for plotting the response curves.
    media_scaler: Scaler that was used to scale the media data before training.
    target_scaler: Scaler used for scaling the target, to unscaled values and
      plot in the original scale.
    prices: Prices to translate the media units to spend. If all your data is
      already in spend numbers you can leave this as None. If some of your data
      is media spend and others is media unit, leave the media spend with price
      1 and add the price to the media unit channels.
    optimal_allocation_per_timeunit: Optimal allocation per time unit per media
      channel. This can be obtained by running the optimization provided by
      LightweightMMM.
    steps: Number of steps to simulate.
    percentage_add: Percentage too exceed the maximum historic spend for the
      simulation of the response curve.
    apply_log_scale: Whether to apply the log scale to the predictions (Y axis).
      When some media channels have very large scale compare to others it might
      be useful to use apply_log_scale=True. Default is False.
    figure_size: Size of the plot figure.
    n_columns: Number of columns to display in the subplots grid. Modifying this
      parameter might require to adjust figure_size accordingly for the plot
      to still have reasonable structure.
    marker_size: Size of the marker for the optimization annotations. Only
      useful if optimal_allocation_per_timeunit is not None. Default is 8.
    legend_fontsize: Legend font size for individual subplots.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    Plots of response curves.
  """
  if not hasattr(media_mix_model, "trace"):
    raise lightweight_mmm.NotFittedModelError(
        "Model needs to be fit first before attempting to plot its response "
        "curves.")
  media = media_mix_model.media
  media_maxes = media.max(axis=0) * (1 + percentage_add)
  if media_mix_model._extra_features is not None:
    extra_features = jnp.expand_dims(
        media_mix_model._extra_features.mean(axis=0), axis=0)
  else:
    extra_features = None
  media_ranges = jnp.expand_dims(
      jnp.linspace(start=0, stop=media_maxes, num=steps), axis=0)

  make_predictions = jax.vmap(
      jax.vmap(_make_single_prediction,
               in_axes=(None, 0, None, None),
               out_axes=0),
      in_axes=(None, 0, None, None), out_axes=1)
  diagonal = jnp.repeat(
      jnp.eye(media_mix_model.n_media_channels), steps,
      axis=0).reshape(media_mix_model.n_media_channels, steps,
                      media_mix_model.n_media_channels)

  prediction_offset = media_mix_model.predict(
      media=jnp.zeros((1, *media.shape[1:])),
      extra_features=extra_features).mean(axis=0)

  if media.ndim == 3:
    diagonal = jnp.expand_dims(diagonal, axis=-1)
    prediction_offset = jnp.expand_dims(prediction_offset, axis=0)
  mock_media = media_ranges * diagonal
  predictions = jnp.squeeze(a=make_predictions(media_mix_model,
                                               mock_media,
                                               extra_features,
                                               seed))
  predictions = predictions - prediction_offset
  media_ranges = jnp.squeeze(media_ranges)
  if target_scaler:
    predictions = target_scaler.inverse_transform(predictions)

  if media_scaler:
    media_ranges = media_scaler.inverse_transform(media_ranges)

  if prices is not None:
    if media.ndim == 3:
      prices = jnp.expand_dims(prices, axis=-1)
    media_ranges *= prices

  if predictions.ndim == 3:
    media_ranges = jnp.sum(media_ranges, axis=-1)
    predictions = jnp.sum(predictions, axis=-1)

  if optimal_allocation_per_timeunit is not None:
    average_allocation = media_mix_model.media.mean(axis=0)
    average_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=average_allocation,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    optimal_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=optimal_allocation_per_timeunit,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    if media_scaler:
      average_allocation = media_scaler.inverse_transform(average_allocation)
      optimal_allocation_per_timeunit = media_scaler.inverse_transform(
          optimal_allocation_per_timeunit)
    if prices is not None:
      optimal_allocation_per_timeunit *= prices
      average_allocation *= prices
    if media.ndim == 3:
      average_allocation = jnp.sum(average_allocation, axis=-1)
      optimal_allocation_per_timeunit = jnp.sum(
          optimal_allocation_per_timeunit, axis=-1)

  kpi_label = "KPI" if target_scaler else "Normalized KPI"
  fig = plt.figure(media_mix_model.n_media_channels + 1,
                   figsize=figure_size,
                   tight_layout=True)
  n_rows = _calculate_number_rows_plot(
      n_media_channels=media_mix_model.n_media_channels, n_columns=n_columns)
  last_ax = fig.add_subplot(n_rows, 1, n_rows)
  for i in range(media_mix_model.n_media_channels):
    ax = fig.add_subplot(n_rows, n_columns, i + 1)
    sns.lineplot(
        x=media_ranges[:, i],
        y=predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=ax)
    sns.lineplot(
        x=media_ranges[:, i],
        y=jnp.log(predictions[:, i]) if apply_log_scale else predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=last_ax)
    if optimal_allocation_per_timeunit is not None:
      ax.plot(
          average_allocation[i],
          average_allocation_predictions[i],
          marker="o",
          markersize=marker_size,
          label="avg_spend",
          color=_PALETTE[i])
      ax.plot(
          optimal_allocation_per_timeunit[i],
          optimal_allocation_predictions[i],
          marker="x",
          markersize=marker_size + 2,
          label="optimal_spend",
          color=_PALETTE[i])
    ax.set_ylabel(kpi_label)
    ax.set_xlabel("Normalized Spend" if not media_scaler else "Spend")
    ax.legend(fontsize=legend_fontsize)

  fig.suptitle("Response curves", fontsize=20)
  last_ax.set_ylabel(kpi_label if not apply_log_scale else f"log({kpi_label})")
  last_ax.set_xlabel("Normalized spend per channel"
                     if not media_scaler else "Spend per channel")
  plt.close()
  return fig

def plot_cross_correlate(feature: jnp.ndarray,
                         target: jnp.ndarray,
                         maxlags: int = 10) -> Tuple[int, float]:
  """Plots the cross correlation coefficients between 2 vectors.

  In the chart look for positive peaks, this shows how the lags of the feature
  lead the target.

  Args:
    feature: Vector, the lags of which predict target.
    target: Vector, what is predicted.
    maxlags: Maximum number of lags.

  Returns:
    Lag index and corresponding correlation of the peak correlation.

  Raises:
    ValueError: If inputs don't have same length.
  """
  if len(feature) != len(target):
    raise ValueError("feature and target need to have the same length.")
  maxlags = jnp.minimum(len(feature) - 1, maxlags)
  mean_feature, mean_target = feature.mean(), target.mean()
  plot = plt.xcorr(
      x=feature - mean_feature, y=target - mean_target, maxlags=maxlags)
  plt.show()
  maxidx = plot[1][plot[0] <= 0].argmax()
  return plot[0][maxidx], plot[1][maxidx]

def plot_var_cost(media: jnp.ndarray, costs: jnp.ndarray,
                  names: List[str]) -> matplotlib.figure.Figure:
  """Plots a a chart between the coefficient of variation and cost.

  Args:
    media: Media matrix.
    costs: Cost vector.
    names: List of variable names.

  Returns:
    Plot of coefficient of variation and cost.

  Raises:
    ValueError if inputs don't conform to same length.
  """
  if media.shape[1] != len(costs):
    raise ValueError("media columns and costs needs to have same length.")
  if media.shape[1] != len(names):
    raise ValueError("media columns and names needs to have same length.")
  coef_of_variation = media.std(axis=0) / media.mean(axis=0)

  fig, ax = plt.subplots(1, 1)
  ax.scatter(x=costs, y=coef_of_variation)
  # https://queirozf.com/entries/add-labels-and-text-to-matplotlib-plots-annotation-examples.
  for i in range(len(costs)):
    x, y, label = costs[i], coef_of_variation[i], names[i]
    ax.annotate(text=label, xy=(x, y))
  ax.set_xlabel("Cost")
  ax.set_ylabel("Coef of Variation")
  plt.close()
  return fig

def _create_shaded_line_plot(predictions: jnp.ndarray,
                             target: jnp.ndarray,
                             axis: matplotlib.axes.Axes,
                             title_prefix: str = "",
                             interval_mid_range: float = .9,
                             digits: int = 3) -> None:
  """Creates a plot of ground truth, predicted value and credibility interval.

  Args:
    predictions: 2d array of predicted values.
    target: Array of true values. Must be same length as predictions.
    axis: Matplotlib axis in which to plot the data.
    title_prefix: Prefix to add as the label of the plot.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.
  """
  if predictions.shape[1] != len(target):
    raise ValueError(
        "Predicted data and ground-truth data must have same length.")
  upper_quantile = 1 - (1 - interval_mid_range) / 2
  lower_quantile = (1 - interval_mid_range) / 2
  upper_bound = jnp.quantile(a=predictions, q=upper_quantile, axis=0)
  lower_bound = jnp.quantile(a=predictions, q=lower_quantile, axis=0)

  r2, _ = arviz.r2_score(y_true=target, y_pred=predictions)
  mape = 100 * metrics.mean_absolute_percentage_error(
      y_true=target, y_pred=predictions.mean(axis=0))
  axis.plot(jnp.arange(target.shape[0]), target, c="grey", alpha=.9)
  axis.plot(
      jnp.arange(target.shape[0]),
      predictions.mean(axis=0),
      c="green",
      alpha=.9)
  axis.fill_between(
      x=jnp.arange(target.shape[0]),
      y1=lower_bound,
      y2=upper_bound,
      alpha=.35,
      color="green")
  axis.legend(["True KPI", "Predicted KPI"])
  axis.yaxis.grid(color="gray", linestyle="dashed", alpha=0.3)
  axis.xaxis.grid(color="gray", linestyle="dashed", alpha=0.3)
  title = " ".join([
      title_prefix,
      "True and predicted KPI.",
      "R2 = {r2:.{digits}f}".format(r2=r2, digits=digits),
      "MAPE = {mape:.{digits}f}%".format(mape=mape, digits=digits)
  ])
  axis.title.set_text(title)
  plt.close()

def _call_fit_plotter(
    predictions: jnp.array,
    target: jnp.array,
    interval_mid_range: float,
    digits: int) -> matplotlib.figure.Figure:
  """Calls the shaded line plot once for national and N times for geo models.

  Args:
    predictions: 2d array of predicted values.
    target: Array of true values. Must be same length as prediction.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.

  Returns:
    Figure of the plot.
  """
  # TODO(): Allow to pass geo names for fit plots

  if predictions.ndim == 3:  # Multiple plots for geo model
    figure, axes = plt.subplots(predictions.shape[-1],
                                figsize=(10, 5 * predictions.shape[-1]))
    for i, ax in enumerate(axes):
      _create_shaded_line_plot(predictions=predictions[..., i],
                               target=target[..., i],
                               axis=ax,
                               title_prefix=f"Geo {i}:",
                               interval_mid_range=interval_mid_range,
                               digits=digits)
  else:  # Single plot for national model
    figure, ax = plt.subplots(1, 1)
    _create_shaded_line_plot(predictions=predictions,
                             target=target,
                             axis=ax,
                             interval_mid_range=interval_mid_range,
                             digits=digits)
  return figure


Prompt:
 historic spend for the
      simulation of the response curve.
    apply_log_scale: Whether to apply the log scale to the predictions (Y axis).
      When some media channels have very large scale compare to others it might
      be useful to use apply_log_scale=True. Default is False.
    figure_size: Size of the plot figure.
    n_columns: Number of columns to display in the subplots grid. Modifying this
      parameter might require to adjust figure_size accordingly for the plot
      to still have reasonable structure.
    marker_size: Size of the marker for the optimization annotations. Only
      useful if optimal_allocation_per_timeunit is not None. Default is 8.
    legend_fontsize: Legend font size for individual subplots.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    Plots of response curves.
  """
  if not hasattr(media_mix_model, "trace"):
    raise lightweight_mmm.NotFittedModelError(
        "Model needs to be fit first before attempting to plot its response "
        "curves.")
  media = media_mix_model.media
  media_maxes = media.max(axis=0) * (1 + percentage_add)
  if media_mix_model._extra_features is not None:
    extra_features = jnp.expand_dims(
        media_mix_model._extra_features.mean(axis=0), axis=0)
  else:
    extra_features = None
  media_ranges = jnp.expand_dims(
      jnp.linspace(start=0, stop=media_maxes, num=steps), axis=0)

  make_predictions = jax.vmap(
      jax.vmap(_make_single_prediction,
               in_axes=(None, 0, None, None),
               out_axes=0),
      in_axes=(None, 0, None, None), out_axes=1)
  diagonal = jnp.repeat(
      jnp.eye(media_mix_model.n_media_channels), steps,
      axis=0).reshape(media_mix_model.n_media_channels, steps,
                      media_mix_model.n_media_channels)

  prediction_offset = media_mix_model.predict(
      media=jnp.zeros((1, *media.shape[1:])),
      extra_features=extra_features).mean(axis=0)

  if media.ndim == 3:
    diagonal = jnp.expand_dims(diagonal, axis=-1)
    prediction_offset = jnp.expand_dims(prediction_offset, axis=0)
  mock_media = media_ranges * diagonal
  predictions = jnp.squeeze(a=make_predictions(media_mix_model,
                                               mock_media,
                                               extra_features,
                                               seed))
  predictions = predictions - prediction_offset
  media_ranges = jnp.squeeze(media_ranges)
  if target_scaler:
    predictions = target_scaler.inverse_transform(predictions)

  if media_scaler:
    media_ranges = media_scaler.inverse_transform(media_ranges)

  if prices is not None:
    if media.ndim == 3:
      prices = jnp.expand_dims(prices, axis=-1)
    media_ranges *= prices

  if predictions.ndim == 3:
    media_ranges = jnp.sum(media_ranges, axis=-1)
    predictions = jnp.sum(predictions, axis=-1)

  if optimal_allocation_per_timeunit is not None:
    average_allocation = media_mix_model.media.mean(axis=0)
    average_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=average_allocation,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    optimal_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=optimal_allocation_per_timeunit,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    if media_scaler:
      average_allocation = media_scaler.inverse_transform(average_allocation)
      optimal_allocation_per_timeunit = media_scaler.inverse_transform(
          optimal_allocation_per_timeunit)
    if prices is not None:
      optimal_allocation_per_timeunit *= prices
      average_allocation *= prices
    if media.ndim == 3:
      average_allocation = jnp.sum(average_allocation, axis=-1)
      optimal_allocation_per_timeunit = jnp.sum(
          optimal_allocation_per_timeunit, axis=-1)

  kpi_label = "KPI" if target_scaler else "Normalized KPI"
  fig = plt.figure(media_mix_model.n_media_channels + 1,
                   figsize=figure_size,
                   tight_layout=True)
  n_rows = _calculate_number_rows_plot(
      n_media_channels=media_mix_model.n_media_channels, n_columns=n_columns)
  last_ax = fig.add_subplot(n_rows, 1, n_rows)
  for i in range(media_mix_model.n_media_channels):
    ax = fig.add_subplot(n_rows, n_columns, i + 1)
    sns.lineplot(
        x=media_ranges[:, i],
        y=predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=ax)
    sns.lineplot(
        x=media_ranges[:, i],
        y=jnp.log(predictions[:, i]) if apply_log_scale else predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=last_ax)
    if optimal_allocation_per_timeunit is not None:
      ax.plot(
          average_allocation[i],
          average_allocation_predictions[i],
          marker="o",
          markersize=marker_size,
          label="avg_spend",
          color=_PALETTE[i])
      ax.plot(
          optimal_allocation_per_timeunit[i],
          optimal_allocation_predictions[i],
          marker="x",
          markersize=marker_size + 2,
          label="optimal_spend",
          color=_PALETTE[i])
    ax.set_ylabel(kpi_label)
    ax.set_xlabel("Normalized Spend" if not media_scaler else "Spend")
    ax.legend(fontsize=legend_fontsize)

  fig.suptitle("Response curves", fontsize=20)
  last_ax.set_ylabel(kpi_label if not apply_log_scale else f"log({kpi_label})")
  last_ax.set_xlabel("Normalized spend per channel"
                     if not media_scaler else "Spend per channel")
  plt.close()
  return fig

def plot_cross_correlate(feature: jnp.ndarray,
                         target: jnp.ndarray,
                         maxlags: int = 10) -> Tuple[int, float]:
  """Plots the cross correlation coefficients between 2 vectors.

  In the chart look for positive peaks, this shows how the lags of the feature
  lead the target.

  Args:
    feature: Vector, the lags of which predict target.
    target: Vector, what is predicted.
    maxlags: Maximum number of lags.

  Returns:
    Lag index and corresponding correlation of the peak correlation.

  Raises:
    ValueError: If inputs don't have same length.
  """
  if len(feature) != len(target):
    raise ValueError("feature and target need to have the same length.")
  maxlags = jnp.minimum(len(feature) - 1, maxlags)
  mean_feature, mean_target = feature.mean(), target.mean()
  plot = plt.xcorr(
      x=feature - mean_feature, y=target - mean_target, maxlags=maxlags)
  plt.show()
  maxidx = plot[1][plot[0] <= 0].argmax()
  return plot[0][maxidx], plot[1][maxidx]

def plot_var_cost(media: jnp.ndarray, costs: jnp.ndarray,
                  names: List[str]) -> matplotlib.figure.Figure:
  """Plots a a chart between the coefficient of variation and cost.

  Args:
    media: Media matrix.
    costs: Cost vector.
    names: List of variable names.

  Returns:
    Plot of coefficient of variation and cost.

  Raises:
    ValueError if inputs don't conform to same length.
  """
  if media.shape[1] != len(costs):
    raise ValueError("media columns and costs needs to have same length.")
  if media.shape[1] != len(names):
    raise ValueError("media columns and names needs to have same length.")
  coef_of_variation = media.std(axis=0) / media.mean(axis=0)

  fig, ax = plt.subplots(1, 1)
  ax.scatter(x=costs, y=coef_of_variation)
  # https://queirozf.com/entries/add-labels-and-text-to-matplotlib-plots-annotation-examples.
  for i in range(len(costs)):
    x, y, label = costs[i], coef_of_variation[i], names[i]
    ax.annotate(text=label, xy=(x, y))
  ax.set_xlabel("Cost")
  ax.set_ylabel("Coef of Variation")
  plt.close()
  return fig

def _create_shaded_line_plot(predictions: jnp.ndarray,
                             target: jnp.ndarray,
                             axis: matplotlib.axes.Axes,
                             title_prefix: str = "",
                             interval_mid_range: float = .9,
                             digits: int = 3) -> None:
  """Creates a plot of ground truth, predicted value and credibility interval.

  Args:
    predictions: 2d array of predicted values.
    target: Array of true values. Must be same length as predictions.
    axis: Matplotlib axis in which to plot the data.
    title_prefix: Prefix to add as the label of the plot.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.
  """
  if predictions.shape[1] != len(target):
    raise ValueError(
        "Predicted data and ground-truth data must have same length.")
  upper_quantile = 1 - (1 - interval_mid_range) / 2
  lower_quantile = (1 - interval_mid_range) / 2
  upper_bound = jnp.quantile(a=predictions, q=upper_quantile, axis=0)
  lower_bound = jnp.quantile(a=predictions, q=lower_quantile, axis=0)

  r2, _ = arviz.r2_score(y_true=target, y_pred=predictions)
  mape = 100 * metrics.mean_absolute_percentage_error(
      y_true=target, y_pred=predictions.mean(axis=0))
  axis.plot(jnp.arange(target.shape[0]), target, c="grey", alpha=.9)
  axis.plot(
      jnp.arange(target.shape[0]),
      predictions.mean(axis=0),
      c="green",
      alpha=.9)
  axis.fill_between(
      x=jnp.arange(target.shape[0]),
      y1=lower_bound,
      y2=upper_bound,
      alpha=.35,
      color="green")
  axis.legend(["True KPI", "Predicted KPI"])
  axis.yaxis.grid(color="gray", linestyle="dashed", alpha=0.3)
  axis.xaxis.grid(color="gray", linestyle="dashed", alpha=0.3)
  title = " ".join([
      title_prefix,
      "True and predicted KPI.",
      "R2 = {r2:.{digits}f}".format(r2=r2, digits=digits),
      "MAPE = {mape:.{digits}f}%".format(mape=mape, digits=digits)
  ])
  axis.title.set_text(title)
  plt.close()

def _call_fit_plotter(
    predictions: jnp.array,
    target: jnp.array,
    interval_mid_range: float,
    digits: int) -> matplotlib.figure.Figure:
  """Calls the shaded line plot once for national and N times for geo models.

  Args:
    predictions: 2d array of predicted values.
    target: Array of true values. Must be same length as prediction.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.

  Returns:
    Figure of the plot.
  """
  # TODO(): Allow to pass geo names for fit plots
  if predictions.ndim == 3:  # Multiple plots for geo model
    figure, axes = plt.subplots(predictions.shape[-1],
                                figsize=(10, 5 * predictions.shape[-1]))
    for i, ax in enumerate(axes):
      _create_shaded_line_plot(predictions=predictions[..., i],
                               target=target[..., i],
                               axis=ax,
                               title_prefix=f"Geo {i}:",
                               interval_mid_range=interval_mid_range,
                               digits=digits)
  else:  # Single plot for national model
    figure, ax = plt.subplots(1, 1)
    _create_shaded_line_plot(predictions=predictions,
                             target=target,
                             axis=ax,
                             interval_mid_range=interval_mid_range,
                             digits=digits)
  return figure

def plot_model_fit(media_mix_model: lightweight_mmm.LightweightMMM,
                   target_scaler: Optional[preprocessing.CustomScaler] = None,
                   interval_mid_range: float = .9,
                   digits: int = 3) -> matplotlib.figure.Figure:
  """Plots the ground truth, predicted value and interval for the training data.

  Model needs to be fit before calling this function to plot.

  Args:
    media_mix_model: Media mix model.
    target_scaler: Scaler used for scaling the target, to unscaled values and
      plot in the original scale.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number.
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.

  Returns:
    Plot of model fit.
  """

  if not hasattr(media_mix_model, "trace"):
    raise lightweight_mmm.NotFittedModelError(
        "Model needs to be fit first before attempting to plot its fit.")
  target_train = media_mix_model._target
  posterior_pred = media_mix_model.trace["mu"]
  if target_scaler:
    posterior_pred = target_scaler.inverse_transform(posterior_pred)
    target_train = target_scaler.inverse_transform(target_train)

  return _call_fit_plotter(
      predictions=posterior_pred,
      target=target_train,
      interval_mid_range=interval_mid_range,
      digits=digits)


Prompt:
map(
      jax.vmap(_make_single_prediction,
               in_axes=(None, 0, None, None),
               out_axes=0),
      in_axes=(None, 0, None, None), out_axes=1)
  diagonal = jnp.repeat(
      jnp.eye(media_mix_model.n_media_channels), steps,
      axis=0).reshape(media_mix_model.n_media_channels, steps,
                      media_mix_model.n_media_channels)

  prediction_offset = media_mix_model.predict(
      media=jnp.zeros((1, *media.shape[1:])),
      extra_features=extra_features).mean(axis=0)

  if media.ndim == 3:
    diagonal = jnp.expand_dims(diagonal, axis=-1)
    prediction_offset = jnp.expand_dims(prediction_offset, axis=0)
  mock_media = media_ranges * diagonal
  predictions = jnp.squeeze(a=make_predictions(media_mix_model,
                                               mock_media,
                                               extra_features,
                                               seed))
  predictions = predictions - prediction_offset
  media_ranges = jnp.squeeze(media_ranges)
  if target_scaler:
    predictions = target_scaler.inverse_transform(predictions)

  if media_scaler:
    media_ranges = media_scaler.inverse_transform(media_ranges)

  if prices is not None:
    if media.ndim == 3:
      prices = jnp.expand_dims(prices, axis=-1)
    media_ranges *= prices

  if predictions.ndim == 3:
    media_ranges = jnp.sum(media_ranges, axis=-1)
    predictions = jnp.sum(predictions, axis=-1)

  if optimal_allocation_per_timeunit is not None:
    average_allocation = media_mix_model.media.mean(axis=0)
    average_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=average_allocation,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    optimal_allocation_predictions = _generate_diagonal_predictions(
        media_mix_model=media_mix_model,
        media_values=optimal_allocation_per_timeunit,
        extra_features=extra_features,
        target_scaler=target_scaler,
        prediction_offset=prediction_offset,
        seed=seed)
    if media_scaler:
      average_allocation = media_scaler.inverse_transform(average_allocation)
      optimal_allocation_per_timeunit = media_scaler.inverse_transform(
          optimal_allocation_per_timeunit)
    if prices is not None:
      optimal_allocation_per_timeunit *= prices
      average_allocation *= prices
    if media.ndim == 3:
      average_allocation = jnp.sum(average_allocation, axis=-1)
      optimal_allocation_per_timeunit = jnp.sum(
          optimal_allocation_per_timeunit, axis=-1)

  kpi_label = "KPI" if target_scaler else "Normalized KPI"
  fig = plt.figure(media_mix_model.n_media_channels + 1,
                   figsize=figure_size,
                   tight_layout=True)
  n_rows = _calculate_number_rows_plot(
      n_media_channels=media_mix_model.n_media_channels, n_columns=n_columns)
  last_ax = fig.add_subplot(n_rows, 1, n_rows)
  for i in range(media_mix_model.n_media_channels):
    ax = fig.add_subplot(n_rows, n_columns, i + 1)
    sns.lineplot(
        x=media_ranges[:, i],
        y=predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=ax)
    sns.lineplot(
        x=media_ranges[:, i],
        y=jnp.log(predictions[:, i]) if apply_log_scale else predictions[:, i],
        label=media_mix_model.media_names[i],
        color=_PALETTE[i],
        ax=last_ax)
    if optimal_allocation_per_timeunit is not None:
      ax.plot(
          average_allocation[i],
          average_allocation_predictions[i],
          marker="o",
          markersize=marker_size,
          label="avg_spend",
          color=_PALETTE[i])
      ax.plot(
          optimal_allocation_per_timeunit[i],
          optimal_allocation_predictions[i],
          marker="x",
          markersize=marker_size + 2,
          label="optimal_spend",
          color=_PALETTE[i])
    ax.set_ylabel(kpi_label)
    ax.set_xlabel("Normalized Spend" if not media_scaler else "Spend")
    ax.legend(fontsize=legend_fontsize)

  fig.suptitle("Response curves", fontsize=20)
  last_ax.set_ylabel(kpi_label if not apply_log_scale else f"log({kpi_label})")
  last_ax.set_xlabel("Normalized spend per channel"
                     if not media_scaler else "Spend per channel")
  plt.close()
  return fig

def plot_cross_correlate(feature: jnp.ndarray,
                         target: jnp.ndarray,
                         maxlags: int = 10) -> Tuple[int, float]:
  """Plots the cross correlation coefficients between 2 vectors.

  In the chart look for positive peaks, this shows how the lags of the feature
  lead the target.

  Args:
    feature: Vector, the lags of which predict target.
    target: Vector, what is predicted.
    maxlags: Maximum number of lags.

  Returns:
    Lag index and corresponding correlation of the peak correlation.

  Raises:
    ValueError: If inputs don't have same length.
  """
  if len(feature) != len(target):
    raise ValueError("feature and target need to have the same length.")
  maxlags = jnp.minimum(len(feature) - 1, maxlags)
  mean_feature, mean_target = feature.mean(), target.mean()
  plot = plt.xcorr(
      x=feature - mean_feature, y=target - mean_target, maxlags=maxlags)
  plt.show()
  maxidx = plot[1][plot[0] <= 0].argmax()
  return plot[0][maxidx], plot[1][maxidx]

def plot_var_cost(media: jnp.ndarray, costs: jnp.ndarray,
                  names: List[str]) -> matplotlib.figure.Figure:
  """Plots a a chart between the coefficient of variation and cost.

  Args:
    media: Media matrix.
    costs: Cost vector.
    names: List of variable names.

  Returns:
    Plot of coefficient of variation and cost.

  Raises:
    ValueError if inputs don't conform to same length.
  """
  if media.shape[1] != len(costs):
    raise ValueError("media columns and costs needs to have same length.")
  if media.shape[1] != len(names):
    raise ValueError("media columns and names needs to have same length.")
  coef_of_variation = media.std(axis=0) / media.mean(axis=0)

  fig, ax = plt.subplots(1, 1)
  ax.scatter(x=costs, y=coef_of_variation)
  # https://queirozf.com/entries/add-labels-and-text-to-matplotlib-plots-annotation-examples.
  for i in range(len(costs)):
    x, y, label = costs[i], coef_of_variation[i], names[i]
    ax.annotate(text=label, xy=(x, y))
  ax.set_xlabel("Cost")
  ax.set_ylabel("Coef of Variation")
  plt.close()
  return fig

def _create_shaded_line_plot(predictions: jnp.ndarray,
                             target: jnp.ndarray,
                             axis: matplotlib.axes.Axes,
                             title_prefix: str = "",
                             interval_mid_range: float = .9,
                             digits: int = 3) -> None:
  """Creates a plot of ground truth, predicted value and credibility interval.

  Args:
    predictions: 2d array of predicted values.
    target: Array of true values. Must be same length as predictions.
    axis: Matplotlib axis in which to plot the data.
    title_prefix: Prefix to add as the label of the plot.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.
  """
  if predictions.shape[1] != len(target):
    raise ValueError(
        "Predicted data and ground-truth data must have same length.")
  upper_quantile = 1 - (1 - interval_mid_range) / 2
  lower_quantile = (1 - interval_mid_range) / 2
  upper_bound = jnp.quantile(a=predictions, q=upper_quantile, axis=0)
  lower_bound = jnp.quantile(a=predictions, q=lower_quantile, axis=0)

  r2, _ = arviz.r2_score(y_true=target, y_pred=predictions)
  mape = 100 * metrics.mean_absolute_percentage_error(
      y_true=target, y_pred=predictions.mean(axis=0))
  axis.plot(jnp.arange(target.shape[0]), target, c="grey", alpha=.9)
  axis.plot(
      jnp.arange(target.shape[0]),
      predictions.mean(axis=0),
      c="green",
      alpha=.9)
  axis.fill_between(
      x=jnp.arange(target.shape[0]),
      y1=lower_bound,
      y2=upper_bound,
      alpha=.35,
      color="green")
  axis.legend(["True KPI", "Predicted KPI"])
  axis.yaxis.grid(color="gray", linestyle="dashed", alpha=0.3)
  axis.xaxis.grid(color="gray", linestyle="dashed", alpha=0.3)
  title = " ".join([
      title_prefix,
      "True and predicted KPI.",
      "R2 = {r2:.{digits}f}".format(r2=r2, digits=digits),
      "MAPE = {mape:.{digits}f}%".format(mape=mape, digits=digits)
  ])
  axis.title.set_text(title)
  plt.close()

def _call_fit_plotter(
    predictions: jnp.array,
    target: jnp.array,
    interval_mid_range: float,
    digits: int) -> matplotlib.figure.Figure:
  """Calls the shaded line plot once for national and N times for geo models.

  Args:
    predictions: 2d array of predicted values.
    target: Array of true values. Must be same length as prediction.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.

  Returns:
    Figure of the plot.
  """
  # TODO(): Allow to pass geo names for fit plots
  if predictions.ndim == 3:  # Multiple plots for geo model
    figure, axes = plt.subplots(predictions.shape[-1],
                                figsize=(10, 5 * predictions.shape[-1]))
    for i, ax in enumerate(axes):
      _create_shaded_line_plot(predictions=predictions[..., i],
                               target=target[..., i],
                               axis=ax,
                               title_prefix=f"Geo {i}:",
                               interval_mid_range=interval_mid_range,
                               digits=digits)
  else:  # Single plot for national model
    figure, ax = plt.subplots(1, 1)
    _create_shaded_line_plot(predictions=predictions,
                             target=target,
                             axis=ax,
                             interval_mid_range=interval_mid_range,
                             digits=digits)
  return figure

def plot_model_fit(media_mix_model: lightweight_mmm.LightweightMMM,
                   target_scaler: Optional[preprocessing.CustomScaler] = None,
                   interval_mid_range: float = .9,
                   digits: int = 3) -> matplotlib.figure.Figure:
  """Plots the ground truth, predicted value and interval for the training data.

  Model needs to be fit before calling this function to plot.

  Args:
    media_mix_model: Media mix model.
    target_scaler: Scaler used for scaling the target, to unscaled values and
      plot in the original scale.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number.
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.

  Returns:
    Plot of model fit.
  """
  if not hasattr(media_mix_model, "trace"):
    raise lightweight_mmm.NotFittedModelError(
        "Model needs to be fit first before attempting to plot its fit.")
  target_train = media_mix_model._target
  posterior_pred = media_mix_model.trace["mu"]
  if target_scaler:
    posterior_pred = target_scaler.inverse_transform(posterior_pred)
    target_train = target_scaler.inverse_transform(target_train)

  return _call_fit_plotter(
      predictions=posterior_pred,
      target=target_train,
      interval_mid_range=interval_mid_range,
      digits=digits)

def plot_out_of_sample_model_fit(out_of_sample_predictions: jnp.ndarray,
                                 out_of_sample_target: jnp.ndarray,
                                 interval_mid_range: float = .9,
                                 digits: int = 3) -> matplotlib.figure.Figure:
  """Plots the ground truth, predicted value and interval for the test data.

  Args:
    out_of_sample_predictions: Predictions for the out-of-sample period, as
      derived from mmm.predict.
    out_of_sample_target: Target for the out-of-sample period. Needs to be on
      the same scale as out_of_sample_predictions.
    interval_mid_range: Mid range interval to take for plotting. Eg. .9 will use
      .05 and .95 as the lower and upper quantiles. Must be a float number.
      between 0 and 1.
    digits: Number of decimals to display on metrics in the plot.

  Returns:
    Plot of model fit.
  """

  return _call_fit_plotter(
      predictions=out_of_sample_predictions,
      target=out_of_sample_target,
      interval_mid_range=interval_mid_range,
      digits=digits)


Prompt:
# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for seasonality."""

import jax
import jax.numpy as jnp
import numpyro
from numpyro import distributions as dist
from numpyro import handlers

from absl.testing import absltest
from absl.testing import parameterized
from lightweight_mmm.core import priors
from lightweight_mmm.core.time import seasonality

class SeasonalityTest(parameterized.TestCase):

  @parameterized.named_parameters([
      dict(
          testcase_name="2_degrees",
          seasonality_arange_value=150,
          degrees_arange_shape=5,
          gamma_seasonality_shape=(5, 2),
      ),
      dict(
          testcase_name="10_degree",
          seasonality_arange_value=150,
          degrees_arange_shape=10,
          gamma_seasonality_shape=(10, 2),
      ),
      dict(
          testcase_name="1_degree",
          seasonality_arange_value=200,
          degrees_arange_shape=1,
          gamma_seasonality_shape=(1, 2),
      ),
  ])
  def test_core_sinusoidal_seasonality_produces_correct_shape(
      self, seasonality_arange_value, degrees_arange_shape,
      gamma_seasonality_shape):
    seasonality_arange = jnp.expand_dims(
        jnp.arange(seasonality_arange_value), axis=-1)
    degrees_arange = jnp.arange(degrees_arange_shape)
    gamma_seasonality = jnp.ones(gamma_seasonality_shape)

    seasonality_values = seasonality._sinusoidal_seasonality(
        seasonality_arange=seasonality_arange,
        degrees_arange=degrees_arange,
        gamma_seasonality=gamma_seasonality,
        frequency=52,
    )
    self.assertEqual(seasonality_values.shape, (seasonality_arange_value,))

  @parameterized.named_parameters(
      dict(
          testcase_name="ten_degrees_national",
          data_shape=(500, 5),
          degrees_seasonality=10,
          expected_shape=(10, 500),
      ),
      dict(
          testcase_name="ten_degrees_geo",
          data_shape=(500, 5, 5),
          degrees_seasonality=10,
          expected_shape=(10, 500, 1),
      ),
      dict(
          testcase_name="one_degrees_national",
          data_shape=(500, 5),
          degrees_seasonality=1,
          expected_shape=(10, 500),
      ),
      dict(
          testcase_name="one_degrees_geo",
          data_shape=(500, 5, 5),
          degrees_seasonality=1,
          expected_shape=(10, 500, 1),
      ),
  )
  def test_model_sinusoidal_seasonality_produces_correct_shape(
      self, data_shape, degrees_seasonality, expected_shape):

    def mock_model_function(data, degrees_seasonality, frequency):

      numpyro.deterministic(
          "seasonality",
          seasonality.sinusoidal_seasonality(
              data=data,
              degrees_seasonality=degrees_seasonality,
              custom_priors={},
              frequency=frequency))


Prompt:
# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for seasonality."""

import jax
import jax.numpy as jnp
import numpyro
from numpyro import distributions as dist
from numpyro import handlers

from absl.testing import absltest
from absl.testing import parameterized
from lightweight_mmm.core import priors
from lightweight_mmm.core.time import seasonality

class SeasonalityTest(parameterized.TestCase):

  @parameterized.named_parameters([
      dict(
          testcase_name="2_degrees",
          seasonality_arange_value=150,
          degrees_arange_shape=5,
          gamma_seasonality_shape=(5, 2),
      ),
      dict(
          testcase_name="10_degree",
          seasonality_arange_value=150,
          degrees_arange_shape=10,
          gamma_seasonality_shape=(10, 2),
      ),
      dict(
          testcase_name="1_degree",
          seasonality_arange_value=200,
          degrees_arange_shape=1,
          gamma_seasonality_shape=(1, 2),
      ),
  ])
  def test_core_sinusoidal_seasonality_produces_correct_shape(
      self, seasonality_arange_value, degrees_arange_shape,
      gamma_seasonality_shape):
    seasonality_arange = jnp.expand_dims(
        jnp.arange(seasonality_arange_value), axis=-1)
    degrees_arange = jnp.arange(degrees_arange_shape)
    gamma_seasonality = jnp.ones(gamma_seasonality_shape)

    seasonality_values = seasonality._sinusoidal_seasonality(
        seasonality_arange=seasonality_arange,
        degrees_arange=degrees_arange,
        gamma_seasonality=gamma_seasonality,
        frequency=52,
    )
    self.assertEqual(seasonality_values.shape, (seasonality_arange_value,))

  @parameterized.named_parameters(
      dict(
          testcase_name="ten_degrees_national",
          data_shape=(500, 5),
          degrees_seasonality=10,
          expected_shape=(10, 500),
      ),
      dict(
          testcase_name="ten_degrees_geo",
          data_shape=(500, 5, 5),
          degrees_seasonality=10,
          expected_shape=(10, 500, 1),
      ),
      dict(
          testcase_name="one_degrees_national",
          data_shape=(500, 5),
          degrees_seasonality=1,
          expected_shape=(10, 500),
      ),
      dict(
          testcase_name="one_degrees_geo",
          data_shape=(500, 5, 5),
          degrees_seasonality=1,
          expected_shape=(10, 500, 1),
      ),
  )
  def test_model_sinusoidal_seasonality_produces_correct_shape(
      self, data_shape, degrees_seasonality, expected_shape):

    def mock_model_function(data, degrees_seasonality, frequency):
      numpyro.deterministic(
          "seasonality",
          seasonality.sinusoidal_seasonality(
              data=data,
              degrees_seasonality=degrees_seasonality,
              custom_priors={},
              frequency=frequency))

    num_samples = 10
    data = jnp.ones(data_shape)
    kernel = numpyro.infer.NUTS(model=mock_model_function)
    mcmc = numpyro.infer.MCMC(
        sampler=kernel, num_warmup=10, num_samples=num_samples, num_chains=1)
    rng_key = jax.random.PRNGKey(0)

    mcmc.run(
        rng_key,
        data=data,
        degrees_seasonality=degrees_seasonality,
        frequency=52,
    )
    seasonality_values = mcmc.get_samples()["seasonality"]

    self.assertEqual(seasonality_values.shape, expected_shape)

  def test_sinusoidal_seasonality_custom_priors_are_taken_correctly(self):
    prior_name = priors.GAMMA_SEASONALITY
    expected_value1, expected_value2 = 5.2, 7.56
    custom_priors = {
        prior_name:
            dist.Kumaraswamy(
                concentration1=expected_value1, concentration0=expected_value2)
    }
    media = jnp.ones((10, 5, 5))
    degrees_seasonality = 3
    frequency = 365

    trace_handler = handlers.trace(
        handlers.seed(seasonality.sinusoidal_seasonality, rng_seed=0))
    trace = trace_handler.get_trace(
        data=media,
        custom_priors=custom_priors,
        degrees_seasonality=degrees_seasonality,
        frequency=frequency,
    )
    values_and_dists = {
        name: site["fn"] for name, site in trace.items() if "fn" in site
    }

    used_distribution = values_and_dists[prior_name]
    if isinstance(used_distribution, dist.ExpandedDistribution):
      used_distribution = used_distribution.base_dist
    self.assertIsInstance(used_distribution, dist.Kumaraswamy)
    self.assertEqual(used_distribution.concentration0, expected_value2)
    self.assertEqual(used_distribution.concentration1, expected_value1)

  @parameterized.named_parameters(
      dict(
          testcase_name="ten_degrees",
          data_shape=(500, 3),
          expected_shape=(10, 500),
      ),
      dict(
          testcase_name="five_degrees",
          data_shape=(500, 3, 5),
          expected_shape=(10, 500, 1),
      ),
  )
  def test_intra_week_seasonality_produces_correct_shape(
      self, data_shape, expected_shape):

    def mock_model_function(data):

      numpyro.deterministic(
          "intra_week",
          seasonality.intra_week_seasonality(
              data=data,
              custom_priors={},
          ))


Prompt:
# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for trend."""

from absl.testing import absltest
from absl.testing import parameterized
import jax
import jax.numpy as jnp
import numpy as np
import numpyro
from numpyro import distributions as dist
from numpyro import handlers

from lightweight_mmm.core import core_utils
from lightweight_mmm.core import priors
from lightweight_mmm.core.time import trend

class TrendTest(parameterized.TestCase):

  @parameterized.named_parameters([
      dict(
          testcase_name="national",
          coef_trend_shape=(),
          trend_length=150,
          expo_trend_shape=(),
      ),
      dict(
          testcase_name="geo",
          coef_trend_shape=(5,),
          trend_length=150,
          expo_trend_shape=(),
      ),
  ])
  def test_core_trend_with_exponent_produces_correct_shape(
      self, coef_trend_shape, trend_length, expo_trend_shape):
    coef_trend = jnp.ones(coef_trend_shape)
    linear_trend = jnp.arange(trend_length)
    if coef_trend.ndim == 1:  # For geo model's case
      linear_trend = jnp.expand_dims(linear_trend, axis=-1)
    expo_trend = jnp.ones(expo_trend_shape)

    trend_values = trend._trend_with_exponent(
        coef_trend=coef_trend, trend=linear_trend, expo_trend=expo_trend)

    self.assertEqual(trend_values.shape,
                     (linear_trend.shape[0], *coef_trend_shape))

  @parameterized.named_parameters([
      dict(testcase_name="national", data_shape=(150, 3)),
      dict(testcase_name="geo", data_shape=(150, 3, 5)),
  ])
  def test_trend_with_exponent_produces_correct_shape(self, data_shape):

    def mock_model_function(data):

      numpyro.deterministic(
          "trend", trend.trend_with_exponent(
              data=data,
              custom_priors={},
          ))


Prompt:
# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for trend."""

from absl.testing import absltest
from absl.testing import parameterized
import jax
import jax.numpy as jnp
import numpy as np
import numpyro
from numpyro import distributions as dist
from numpyro import handlers

from lightweight_mmm.core import core_utils
from lightweight_mmm.core import priors
from lightweight_mmm.core.time import trend

class TrendTest(parameterized.TestCase):

  @parameterized.named_parameters([
      dict(
          testcase_name="national",
          coef_trend_shape=(),
          trend_length=150,
          expo_trend_shape=(),
      ),
      dict(
          testcase_name="geo",
          coef_trend_shape=(5,),
          trend_length=150,
          expo_trend_shape=(),
      ),
  ])
  def test_core_trend_with_exponent_produces_correct_shape(
      self, coef_trend_shape, trend_length, expo_trend_shape):
    coef_trend = jnp.ones(coef_trend_shape)
    linear_trend = jnp.arange(trend_length)
    if coef_trend.ndim == 1:  # For geo model's case
      linear_trend = jnp.expand_dims(linear_trend, axis=-1)
    expo_trend = jnp.ones(expo_trend_shape)

    trend_values = trend._trend_with_exponent(
        coef_trend=coef_trend, trend=linear_trend, expo_trend=expo_trend)

    self.assertEqual(trend_values.shape,
                     (linear_trend.shape[0], *coef_trend_shape))

  @parameterized.named_parameters([
      dict(testcase_name="national", data_shape=(150, 3)),
      dict(testcase_name="geo", data_shape=(150, 3, 5)),
  ])
  def test_trend_with_exponent_produces_correct_shape(self, data_shape):

    def mock_model_function(data):
      numpyro.deterministic(
          "trend", trend.trend_with_exponent(
              data=data,
              custom_priors={},
          ))

    num_samples = 10
    data = jnp.ones(data_shape)
    kernel = numpyro.infer.NUTS(model=mock_model_function)
    mcmc = numpyro.infer.MCMC(
        sampler=kernel, num_warmup=10, num_samples=num_samples, num_chains=1)
    rng_key = jax.random.PRNGKey(0)
    coef_expected_shape = () if data.ndim == 2 else (data.shape[2],)

    mcmc.run(rng_key, data=data)
    trend_values = mcmc.get_samples()["trend"]

    self.assertEqual(trend_values.shape,
                     (num_samples, data.shape[0], *coef_expected_shape))

  @parameterized.named_parameters(
      dict(
          testcase_name=f"model_{priors.COEF_TREND}",
          prior_name=priors.COEF_TREND,
      ),
      dict(
          testcase_name=f"model_{priors.EXPO_TREND}",
          prior_name=priors.EXPO_TREND,
      ),
  )
  def test_trend_with_exponent_custom_priors_are_taken_correctly(
      self, prior_name):
    expected_value1, expected_value2 = 5.2, 7.56
    custom_priors = {
        prior_name:
            dist.Kumaraswamy(
                concentration1=expected_value1, concentration0=expected_value2)
    }
    media = jnp.ones((10, 5, 5))

    trace_handler = handlers.trace(
        handlers.seed(trend.trend_with_exponent, rng_seed=0))
    trace = trace_handler.get_trace(
        data=media,
        custom_priors=custom_priors,
    )
    values_and_dists = {
        name: site["fn"] for name, site in trace.items() if "fn" in site
    }

    used_distribution = values_and_dists[prior_name]
    if isinstance(used_distribution, dist.ExpandedDistribution):
      used_distribution = used_distribution.base_dist
    self.assertIsInstance(used_distribution, dist.Kumaraswamy)
    self.assertEqual(used_distribution.concentration0, expected_value2)
    self.assertEqual(used_distribution.concentration1, expected_value1)

  @parameterized.named_parameters([
      dict(
          testcase_name="dynamic_trend_national_shape",
          number_periods=100,
          initial_level_shape=(),
          initial_slope_shape=(),
          variance_level_shape=(),
          variance_slope_shape=(),
      ),
      dict(
          testcase_name="dynamic_trend_geo_shape",
          number_periods=100,
          initial_level_shape=(2,),
          initial_slope_shape=(2,),
          variance_level_shape=(2,),
          variance_slope_shape=(2,),
      ),
  ])
  def test_core_dynamic_trend_produces_correct_shape(
      self, number_periods, initial_level_shape, initial_slope_shape,
      variance_level_shape, variance_slope_shape):
    initial_level = jnp.ones(initial_level_shape)
    initial_slope = jnp.ones(initial_slope_shape)
    variance_level = jnp.ones(variance_level_shape)
    variance_slope = jnp.ones(variance_slope_shape)
    random_walk_level = jnp.arange(number_periods)
    random_walk_slope = jnp.arange(number_periods)
    if initial_level.ndim == 1:  # For geo model's case
      random_walk_level = jnp.expand_dims(random_walk_level, axis=-1)
      random_walk_slope = jnp.expand_dims(random_walk_slope, axis=-1)

    dynamic_trend_values = trend._dynamic_trend(
        number_periods=number_periods,
        random_walk_level=random_walk_level,
        random_walk_slope=random_walk_slope,
        initial_level=initial_level,
        initial_slope=initial_slope,
        variance_level=variance_level,
        variance_slope=variance_slope,
    )

    self.assertEqual(dynamic_trend_values.shape,
                     (number_periods, *initial_level_shape))

  def test_core_dynamic_trend_produces_correct_value(self):
    number_periods = 5
    initial_level = jnp.ones(())
    initial_slope = jnp.ones(())
    variance_level = jnp.ones(())
    variance_slope = jnp.ones(())
    random_walk_level = jnp.arange(number_periods)
    random_walk_slope = jnp.arange(number_periods)
    dynamic_trend_expected_value = jnp.array([1, 3, 7, 14, 25])

    dynamic_trend_values = trend._dynamic_trend(
        number_periods=number_periods,
        random_walk_level=random_walk_level,
        random_walk_slope=random_walk_slope,
        initial_level=initial_level,
        initial_slope=initial_slope,
        variance_level=variance_level,
        variance_slope=variance_slope,
    )

    np.testing.assert_array_equal(x=dynamic_trend_values,
                                  y=dynamic_trend_expected_value)

  @parameterized.named_parameters([
      dict(
          testcase_name="national_with_prediction_is_true",
          data_shape=(100, 3),
          is_trend_prediction=True),
      dict(
          testcase_name="geo_with_prediction_is_true",
          data_shape=(150, 3, 5),
          is_trend_prediction=True),
      dict(
          testcase_name="national_with_prediction_is_false",
          data_shape=(100, 3),
          is_trend_prediction=False),
      dict(
          testcase_name="geo_with_prediction_is_false",
          data_shape=(150, 3, 5),
          is_trend_prediction=False),
  ])
  def test_dynamic_trend_produces_correct_shape(
      self, data_shape, is_trend_prediction):

    def mock_model_function(geo_size, data_size):

      numpyro.deterministic(
          "trend", trend.dynamic_trend(
              geo_size=geo_size,
              data_size=data_size,
              is_trend_prediction=is_trend_prediction,
              custom_priors={},
          ))


Prompt:
# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Core and modelling functions for seasonality."""

from typing import Mapping

import jax
import jax.numpy as jnp
import numpyro
from numpyro import distributions as dist

from lightweight_mmm.core import priors
from lightweight_mmm.core import core_utils

@jax.jit
def _sinusoidal_seasonality(
    seasonality_arange: jnp.ndarray,
    degrees_arange: jnp.ndarray,
    gamma_seasonality: jnp.ndarray,
    frequency: int,
) -> jnp.ndarray:
  """Core calculation of cyclic variation seasonality.

  Args:
    seasonality_arange: Array with range [0, N - 1] where N is the size of the
      data for which the seasonality is modelled.
    degrees_arange: Array with range [0, D - 1] where D is the number of degrees
      to use. Must be greater or equal than 1.
    gamma_seasonality: Factor to multiply to each degree calculation. Shape must
      be aligned with the number of degrees.
    frequency: Frecuency of the seasonality be in computed.

  Returns:
    An array with the seasonality values.
  """

  inner_value = seasonality_arange * 2 * jnp.pi * degrees_arange / frequency
  season_matrix_sin = jnp.sin(inner_value)
  season_matrix_cos = jnp.cos(inner_value)
  season_matrix = jnp.concatenate([
      jnp.expand_dims(a=season_matrix_sin, axis=-1),
      jnp.expand_dims(a=season_matrix_cos, axis=-1)
  ],
                                  axis=-1)
  return jnp.einsum("tds, ds -> t", season_matrix, gamma_seasonality)