Incongruent Response Curves and Contribution Calculations in terms of ROI about lightweight_mmm HOT 9 OPEN

def media_transform_hill_adstock(media_mix_model,
                                  media_data,
                           lag_weight,
                           half_max_effective_concentration,
                           slope, apply_adstock, normalise):
ransforms the input data with the adstock and hill functions.

  Args:
  media_data: Media data to be transformed. It is expected to have 2 dims for
    national models and 3 for geo models.
  normalise: Whether to normalise the output values.

  Returns:
  The transformed media data.
  """
  if media_mix_model.n_geos > 1:
      lag_weight=jnp.repeat(lag_weight,media_mix_model.n_geos).reshape(half_max_effective_concentration.shape)
      slope = jnp.squeeze(jnp.repeat(slope,media_mix_model.n_geos).reshape(half_max_effective_concentration.shape))

  
  if apply_adstock:
    return media_transforms.hill(
    data=media_transforms.adstock(
        data=media_data, lag_weight=lag_weight, normalise=normalise),
    half_max_effective_concentration=half_max_effective_concentration,
    slope=slope)

  else:
    return media_transforms.hill(
    data=media_data,
    half_max_effective_concentration=half_max_effective_concentration,
    slope=slope)

percentage_add=2
media_mix_model=mmm
steps=25
prices=average_cost_per_media
media = media_mix_model.media
media_maxes = media.max(axis=0) * (1 + percentage_add)
half_max_effective_concentration=media_mix_model._mcmc.get_samples()['half_max_effective_concentration']#.mean(axis=0)
lag_weight=media_mix_model._mcmc.get_samples()['lag_weight']#.mean(axis=0)
slope=media_mix_model._mcmc.get_samples()['slope']#.mean(axis=0)
beta_media=media_mix_model.trace['beta_media'].mean(axis=0)

media_ranges = jnp.expand_dims(
      jnp.linspace(start=0, stop=media_maxes, num=steps), axis=0)

media_ranges=jnp.repeat(media_ranges,len(beta_media)).reshape(steps,len(beta_media),media_mix_model.n_media_channels, media_mix_model.n_geos)


media_response=beta_media*media_transform_hill_adstock(media_mix_model,
                                  media_ranges,
                           lag_weight,
                           half_max_effective_concentration=half_max_effective_concentration,
                           slope=slope, normalise=True)

percentage_add=2
media_mix_model=mmm
steps=25
prices=average_cost_per_media
media = media_mix_model.media

media_maxes = media.max(axis=0) * (1 + percentage_add)
half_max_effective_concentration=media_mix_model._mcmc.get_samples()['half_max_effective_concentration'].mean(axis=0)
lag_weight=media_mix_model._mcmc.get_samples()['lag_weight'].mean(axis=0)
slope=media_mix_model._mcmc.get_samples()['slope'].mean(axis=0)
beta_media=media_mix_model.trace['beta_media'].mean(axis=0)

media_ranges = jnp.expand_dims(
      jnp.linspace(start=0, stop=media_maxes, num=steps), axis=0)

media_ranges=jnp.repeat(media_ranges,len(beta_media)).reshape(steps,len(beta_media),media_mix_model.n_media_channels, media_mix_model.n_geos)

media_response=beta_media*media_transform_hill_adstock(media_mix_model,
                                  media_ranges,
                           lag_weight,
                           half_max_effective_concentration=half_max_effective_concentration,
                           slope=slope, normalise=True)