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)

def sinusoidal_seasonality(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    degrees_seasonality: int = 2,
    frequency: int = 52,
) -> jnp.ndarray:
  """Calculates cyclic variation seasonality.

  For detailed info check:
    https://en.wikipedia.org/wiki/Seasonality#Modeling

  Args:
    data: Data for which the seasonality will be modelled for. It is used to
      obtain the length of the time dimension, axis 0.
    custom_priors: The custom priors we want the model to take instead of
      default ones.
    degrees_seasonality: Number of degrees to use. Must be greater or equal than
      1.
    frequency: Frecuency of the seasonality be in computed. By default is 52 for
      weekly data (52 weeks in a year).

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

  number_periods = data.shape[0]
  default_priors = priors.get_default_priors()
  n_geos = core_utils.get_number_geos(data=data)
  with numpyro.plate(name=f"{priors.GAMMA_SEASONALITY}_sin_cos_plate", size=2):
    with numpyro.plate(
        name=f"{priors.GAMMA_SEASONALITY}_plate", size=degrees_seasonality):
      gamma_seasonality = numpyro.sample(
          name=priors.GAMMA_SEASONALITY,
          fn=custom_priors.get(priors.GAMMA_SEASONALITY,
                               default_priors[priors.GAMMA_SEASONALITY]))
  seasonality_arange = jnp.expand_dims(a=jnp.arange(number_periods), axis=-1)
  degrees_arange = jnp.arange(degrees_seasonality)
  seasonality_values = _sinusoidal_seasonality(
      seasonality_arange=seasonality_arange,
      degrees_arange=degrees_arange,
      frequency=frequency,
      gamma_seasonality=gamma_seasonality,
  )
  if n_geos > 1:
    seasonality_values = jnp.expand_dims(seasonality_values, axis=-1)
  return seasonality_values


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)

def sinusoidal_seasonality(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    degrees_seasonality: int = 2,
    frequency: int = 52,
) -> jnp.ndarray:
  """Calculates cyclic variation seasonality.

  For detailed info check:
    https://en.wikipedia.org/wiki/Seasonality#Modeling

  Args:
    data: Data for which the seasonality will be modelled for. It is used to
      obtain the length of the time dimension, axis 0.
    custom_priors: The custom priors we want the model to take instead of
      default ones.
    degrees_seasonality: Number of degrees to use. Must be greater or equal than
      1.
    frequency: Frecuency of the seasonality be in computed. By default is 52 for
      weekly data (52 weeks in a year).

  Returns:
    An array with the seasonality values.
  """
  number_periods = data.shape[0]
  default_priors = priors.get_default_priors()
  n_geos = core_utils.get_number_geos(data=data)
  with numpyro.plate(name=f"{priors.GAMMA_SEASONALITY}_sin_cos_plate", size=2):
    with numpyro.plate(
        name=f"{priors.GAMMA_SEASONALITY}_plate", size=degrees_seasonality):
      gamma_seasonality = numpyro.sample(
          name=priors.GAMMA_SEASONALITY,
          fn=custom_priors.get(priors.GAMMA_SEASONALITY,
                               default_priors[priors.GAMMA_SEASONALITY]))
  seasonality_arange = jnp.expand_dims(a=jnp.arange(number_periods), axis=-1)
  degrees_arange = jnp.arange(degrees_seasonality)
  seasonality_values = _sinusoidal_seasonality(
      seasonality_arange=seasonality_arange,
      degrees_arange=degrees_arange,
      frequency=frequency,
      gamma_seasonality=gamma_seasonality,
  )
  if n_geos > 1:
    seasonality_values = jnp.expand_dims(seasonality_values, axis=-1)
  return seasonality_values

def _intra_week_seasonality(
    data: jnp.ndarray,
    weekday: jnp.ndarray,
) -> jnp.ndarray:
  data_size = data.shape[0]
  return weekday[jnp.arange(data_size) % 7]

def intra_week_seasonality(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
) -> jnp.ndarray:
  """Models intra week seasonality.

  Args:
    data: Data for which the seasonality will be modelled for. It is used to
      obtain the length of the time dimension, axis 0.
    custom_priors: The custom priors we want the model to take instead of
      default ones.

  Returns:
    The contribution of the weekday seasonality.
  """

  default_priors = priors.get_default_priors()
  with numpyro.plate(name=f"{priors.WEEKDAY}_plate", size=7):
    weekday = numpyro.sample(
        name=priors.WEEKDAY,
        fn=custom_priors.get(priors.WEEKDAY, default_priors[priors.WEEKDAY]))

  weekday_series = _intra_week_seasonality(data=data, weekday=weekday)

  if data.ndim == 3:  # For geo model's case
    weekday_series = jnp.expand_dims(weekday_series, axis=-1)

  return weekday_series


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 trend."""

import functools
from typing import Mapping

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

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

@jax.jit
def _trend_with_exponent(coef_trend: jnp.ndarray, trend: jnp.ndarray,
                         expo_trend: jnp.ndarray) -> jnp.ndarray:
  """Applies the coefficient and exponent to the trend to obtain trend values.

  Args:
    coef_trend: Coefficient to be multiplied by the trend.
    trend: Initial trend values.
    expo_trend: Exponent to be applied to the trend.

  Returns:
    The trend values generated.
  """
  return coef_trend * trend**expo_trend

def trend_with_exponent(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
) -> jnp.ndarray:
  """Trend with exponent for curvature.

  Args:
    data: Data for which trend will be created.
    custom_priors: The custom priors we want the model to take instead of the
      default ones. See our custom_priors documentation for details about the
      API and possible options.

  Returns:
    The values of the trend.
  """

  default_priors = priors.get_default_priors()
  n_geos = core_utils.get_number_geos(data=data)
  # TODO(): Force all geos to have the same trend sign.
  with numpyro.plate(name=f"{priors.COEF_TREND}_plate", size=n_geos):
    coef_trend = numpyro.sample(
        name=priors.COEF_TREND,
        fn=custom_priors.get(priors.COEF_TREND,
                             default_priors[priors.COEF_TREND]))

  expo_trend = numpyro.sample(
      name=priors.EXPO_TREND,
      fn=custom_priors.get(priors.EXPO_TREND,
                           default_priors[priors.EXPO_TREND]))
  linear_trend = jnp.arange(data.shape[0])
  if n_geos > 1:  # For geo model's case
    linear_trend = jnp.expand_dims(linear_trend, axis=-1)
  return _trend_with_exponent(
      coef_trend=coef_trend, trend=linear_trend, expo_trend=expo_trend)


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 trend."""

import functools
from typing import Mapping

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

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

@jax.jit
def _trend_with_exponent(coef_trend: jnp.ndarray, trend: jnp.ndarray,
                         expo_trend: jnp.ndarray) -> jnp.ndarray:
  """Applies the coefficient and exponent to the trend to obtain trend values.

  Args:
    coef_trend: Coefficient to be multiplied by the trend.
    trend: Initial trend values.
    expo_trend: Exponent to be applied to the trend.

  Returns:
    The trend values generated.
  """
  return coef_trend * trend**expo_trend

def trend_with_exponent(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
) -> jnp.ndarray:
  """Trend with exponent for curvature.

  Args:
    data: Data for which trend will be created.
    custom_priors: The custom priors we want the model to take instead of the
      default ones. See our custom_priors documentation for details about the
      API and possible options.

  Returns:
    The values of the trend.
  """
  default_priors = priors.get_default_priors()
  n_geos = core_utils.get_number_geos(data=data)
  # TODO(): Force all geos to have the same trend sign.
  with numpyro.plate(name=f"{priors.COEF_TREND}_plate", size=n_geos):
    coef_trend = numpyro.sample(
        name=priors.COEF_TREND,
        fn=custom_priors.get(priors.COEF_TREND,
                             default_priors[priors.COEF_TREND]))

  expo_trend = numpyro.sample(
      name=priors.EXPO_TREND,
      fn=custom_priors.get(priors.EXPO_TREND,
                           default_priors[priors.EXPO_TREND]))
  linear_trend = jnp.arange(data.shape[0])
  if n_geos > 1:  # For geo model's case
    linear_trend = jnp.expand_dims(linear_trend, axis=-1)
  return _trend_with_exponent(
      coef_trend=coef_trend, trend=linear_trend, expo_trend=expo_trend)

@functools.partial(jax.jit, static_argnames=("number_periods",))
def _dynamic_trend(
    number_periods: int,
    random_walk_level: jnp.ndarray,
    random_walk_slope: jnp.ndarray,
    initial_level: jnp.ndarray,
    initial_slope: jnp.ndarray,
    variance_level: jnp.ndarray,
    variance_slope: jnp.ndarray,
) -> jnp.ndarray:
  """Calculates dynamic trend using local linear trend method.

  More details about this function can be found in:
  https://storage.googleapis.com/pub-tools-public-publication-data/pdf/41854.pdf

  Args:
    number_periods: Number of time periods in the data.
    random_walk_level: Random walk of level from sample.
    random_walk_slope: Random walk of slope from sample.
    initial_level: The initial value for level in local linear trend model.
    initial_slope: The initial value for slope in local linear trend model.
    variance_level: The variance of the expected increase in level between time.
    variance_slope: The variance of the expected increase in slope between time.

  Returns:
    The dynamic trend values for the given data with the given parameters.
  """
  # Simulate gaussian random walk of level with initial level.

  random_level = variance_level * random_walk_level
  random_level_with_initial_level = jnp.concatenate(
      [jnp.array([random_level[0] + initial_level]), random_level[1:]])
  level_trend_t = jnp.cumsum(random_level_with_initial_level, axis=0)
  # Simulate gaussian random walk of slope with initial slope.
  random_slope = variance_slope * random_walk_slope
  random_slope_with_initial_slope = jnp.concatenate(
      [jnp.array([random_slope[0] + initial_slope]), random_slope[1:]])
  slope_trend_t = jnp.cumsum(random_slope_with_initial_slope, axis=0)
  # Accumulate sum of slope series to address latent variable slope in function
  # level_t = level_t-1 + slope_t-1.
  initial_zero_shape = [(1, 0)] if slope_trend_t.ndim == 1 else [(1, 0), (0, 0)]
  slope_trend_cumsum = jnp.pad(
      jnp.cumsum(slope_trend_t, axis=0)[:number_periods - 1],
      initial_zero_shape, mode="constant", constant_values=0)
  return level_trend_t + slope_trend_cumsum


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.

"""Module for modeling the intercept."""

from typing import Mapping

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

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

def simple_intercept(
    data: jnp.ndarray,
    custom_priors: Mapping[str,
                           dist.Distribution] = immutabledict.immutabledict(),
) -> jnp.ndarray:
  """Calculates a national or geo incercept.
  Note that this intercept is constant over time.

  Args:
    data: Media input data. Media data must have either 2 dims for national
      model or 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones. Refer to the full documentation on custom priors for
      details.

  Returns:
    The values of the intercept.
  """

  default_priors = priors.get_default_priors()
  n_geos = core_utils.get_number_geos(data=data)

  with numpyro.plate(name=f"{priors.INTERCEPT}_plate", size=n_geos):
    intercept = numpyro.sample(
        name=priors.INTERCEPT,
        fn=custom_priors.get(priors.INTERCEPT,
                             default_priors[priors.INTERCEPT]),
    )
  return intercept


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 intercept."""

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

from lightweight_mmm.core import core_utils
from lightweight_mmm.core import priors
from lightweight_mmm.core.baseline import intercept

class InterceptTest(parameterized.TestCase):

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

    def mock_model_function(data):

      numpyro.deterministic(
          "intercept_values",
          intercept.simple_intercept(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.

"""Set of core and modelling saturation functions."""

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

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

@jax.jit
def _hill(
    data: jnp.ndarray,
    half_max_effective_concentration: jnp.ndarray,
    slope: jnp.ndarray,
) -> jnp.ndarray:
  """Calculates the hill function for a given array of values.

  Refer to the following link for detailed information on this equation:
    https://en.wikipedia.org/wiki/Hill_equation_(biochemistry)

  Args:
    data: Input data.
    half_max_effective_concentration: ec50 value for the hill function.
    slope: Slope of the hill function.

  Returns:
    The hill values for the respective input data.
  """

  save_transform = core_utils.apply_exponent_safe(
      data=data / half_max_effective_concentration, exponent=-slope)
  return jnp.where(save_transform == 0, x=0, y=1. / (1 + save_transform))


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.

"""Set of core and modelling saturation functions."""

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

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

@jax.jit
def _hill(
    data: jnp.ndarray,
    half_max_effective_concentration: jnp.ndarray,
    slope: jnp.ndarray,
) -> jnp.ndarray:
  """Calculates the hill function for a given array of values.

  Refer to the following link for detailed information on this equation:
    https://en.wikipedia.org/wiki/Hill_equation_(biochemistry)

  Args:
    data: Input data.
    half_max_effective_concentration: ec50 value for the hill function.
    slope: Slope of the hill function.

  Returns:
    The hill values for the respective input data.
  """
  save_transform = core_utils.apply_exponent_safe(
      data=data / half_max_effective_concentration, exponent=-slope)
  return jnp.where(save_transform == 0, x=0, y=1. / (1 + save_transform))

def hill(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    prefix: str = "",
) -> jnp.ndarray:
  """Transforms the input data with the adstock and hill functions.

  Args:
    data: Media data to be transformed. It is expected to have 2 dims for
      national models and 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones. The possible names of parameters for hill_adstock and
      exponent are "lag_weight", "half_max_effective_concentration" and "slope".
    prefix: Prefix to use in the variable name for Numpyro.

  Returns:
    The transformed media data.
  """

  default_priors = priors.get_default_priors()

  with numpyro.plate(
      name=f"{prefix}{priors.HALF_MAX_EFFECTIVE_CONCENTRATION}_plate",
      size=data.shape[1]):
    half_max_effective_concentration = numpyro.sample(
        name=f"{prefix}{priors.HALF_MAX_EFFECTIVE_CONCENTRATION}",
        fn=custom_priors.get(
            priors.HALF_MAX_EFFECTIVE_CONCENTRATION,
            default_priors[priors.HALF_MAX_EFFECTIVE_CONCENTRATION]))

  with numpyro.plate(name=f"{prefix}{priors.SLOPE}_plate", size=data.shape[1]):
    slope = numpyro.sample(
        name=f"{prefix}{priors.SLOPE}",
        fn=custom_priors.get(priors.SLOPE, default_priors[priors.SLOPE]))

  if data.ndim == 3:
    half_max_effective_concentration = jnp.expand_dims(
        half_max_effective_concentration, axis=-1)
    slope = jnp.expand_dims(slope, axis=-1)

  return _hill(
      data=data,
      half_max_effective_concentration=half_max_effective_concentration,
      slope=slope)


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.

"""Set of core and modelling saturation functions."""

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

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

@jax.jit
def _hill(
    data: jnp.ndarray,
    half_max_effective_concentration: jnp.ndarray,
    slope: jnp.ndarray,
) -> jnp.ndarray:
  """Calculates the hill function for a given array of values.

  Refer to the following link for detailed information on this equation:
    https://en.wikipedia.org/wiki/Hill_equation_(biochemistry)

  Args:
    data: Input data.
    half_max_effective_concentration: ec50 value for the hill function.
    slope: Slope of the hill function.

  Returns:
    The hill values for the respective input data.
  """
  save_transform = core_utils.apply_exponent_safe(
      data=data / half_max_effective_concentration, exponent=-slope)
  return jnp.where(save_transform == 0, x=0, y=1. / (1 + save_transform))

def hill(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    prefix: str = "",
) -> jnp.ndarray:
  """Transforms the input data with the adstock and hill functions.

  Args:
    data: Media data to be transformed. It is expected to have 2 dims for
      national models and 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones. The possible names of parameters for hill_adstock and
      exponent are "lag_weight", "half_max_effective_concentration" and "slope".
    prefix: Prefix to use in the variable name for Numpyro.

  Returns:
    The transformed media data.
  """
  default_priors = priors.get_default_priors()

  with numpyro.plate(
      name=f"{prefix}{priors.HALF_MAX_EFFECTIVE_CONCENTRATION}_plate",
      size=data.shape[1]):
    half_max_effective_concentration = numpyro.sample(
        name=f"{prefix}{priors.HALF_MAX_EFFECTIVE_CONCENTRATION}",
        fn=custom_priors.get(
            priors.HALF_MAX_EFFECTIVE_CONCENTRATION,
            default_priors[priors.HALF_MAX_EFFECTIVE_CONCENTRATION]))

  with numpyro.plate(name=f"{prefix}{priors.SLOPE}_plate", size=data.shape[1]):
    slope = numpyro.sample(
        name=f"{prefix}{priors.SLOPE}",
        fn=custom_priors.get(priors.SLOPE, default_priors[priors.SLOPE]))

  if data.ndim == 3:
    half_max_effective_concentration = jnp.expand_dims(
        half_max_effective_concentration, axis=-1)
    slope = jnp.expand_dims(slope, axis=-1)

  return _hill(
      data=data,
      half_max_effective_concentration=half_max_effective_concentration,
      slope=slope)

def _exponent(data: jnp.ndarray, exponent_values: jnp.ndarray) -> jnp.ndarray:
  """Applies exponent to the given data."""
  return core_utils.apply_exponent_safe(data=data, exponent=exponent_values)

def exponent(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    prefix: str = "",
) -> jnp.ndarray:
  """Transforms the input data with the carryover function and exponent.

  Args:
    data: Media data to be transformed. It is expected to have 2 dims for
      national models and 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones.
    prefix: Prefix to use in the variable name for Numpyro.

  Returns:
    The transformed media data.
  """

  default_priors = priors.get_default_priors()

  with numpyro.plate(
      name=f"{prefix}{priors.EXPONENT}_plate", size=data.shape[1]):
    exponent_values = numpyro.sample(
        name=f"{prefix}{priors.EXPONENT}",
        fn=custom_priors.get(priors.EXPONENT, default_priors[priors.EXPONENT]))

  if data.ndim == 3:
    exponent_values = jnp.expand_dims(exponent_values, axis=-1)
  return _exponent(data=data, exponent_values=exponent_values)


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.

"""Set of core and modelling lagging functions."""

import functools
from typing import Mapping, Union

import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from lightweight_mmm.core import priors

@functools.partial(jax.vmap, in_axes=(1, 1, None), out_axes=1)
def _carryover_convolve(data: jnp.ndarray, weights: jnp.ndarray,
                        number_lags: int) -> jnp.ndarray:
  """Applies the convolution between the data and the weights for the carryover.

  Args:
    data: Input data.
    weights: Window weights for the carryover.
    number_lags: Number of lags the window has.

  Returns:
    The result values from convolving the data and the weights with padding.
  """
  window = jnp.concatenate([jnp.zeros(number_lags - 1), weights])
  return jax.scipy.signal.convolve(data, window, mode="same") / weights.sum()

@functools.partial(jax.jit, static_argnames=("number_lags",))
def _carryover(
    data: jnp.ndarray,
    ad_effect_retention_rate: jnp.ndarray,
    peak_effect_delay: jnp.ndarray,
    number_lags: int,
) -> jnp.ndarray:
  """Calculates media carryover.

  More details about this function can be found in:
  https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/46001.pdf

  Args:
    data: Input data. It is expected that data has either 2 dimensions for
      national models and 3 for geo models.
    ad_effect_retention_rate: Retention rate of the advertisement effect.
      Default is 0.5.
    peak_effect_delay: Delay of the peak effect in the carryover function.
      Default is 1.
    number_lags: Number of lags to include in the carryover calculation. Default
      is 13.

  Returns:
    The carryover values for the given data with the given parameters.
  """

  lags_arange = jnp.expand_dims(
      jnp.arange(number_lags, dtype=jnp.float32), axis=-1)
  convolve_func = _carryover_convolve
  if data.ndim == 3:
    # Since _carryover_convolve is already vmaped in the decorator we only need
    # to vmap it once here to handle the geo level data. We keep the windows bi
    # dimensional also for three dims data and vmap over only the extra data
    # dimension.
    convolve_func = jax.vmap(
        fun=_carryover_convolve, in_axes=(2, None, None), out_axes=2)
  weights = ad_effect_retention_rate**((lags_arange - peak_effect_delay)**2)
  return convolve_func(data, weights, number_lags)


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.

"""Set of core and modelling lagging functions."""

import functools
from typing import Mapping, Union

import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from lightweight_mmm.core import priors

@functools.partial(jax.vmap, in_axes=(1, 1, None), out_axes=1)
def _carryover_convolve(data: jnp.ndarray, weights: jnp.ndarray,
                        number_lags: int) -> jnp.ndarray:
  """Applies the convolution between the data and the weights for the carryover.

  Args:
    data: Input data.
    weights: Window weights for the carryover.
    number_lags: Number of lags the window has.

  Returns:
    The result values from convolving the data and the weights with padding.
  """
  window = jnp.concatenate([jnp.zeros(number_lags - 1), weights])
  return jax.scipy.signal.convolve(data, window, mode="same") / weights.sum()

@functools.partial(jax.jit, static_argnames=("number_lags",))
def _carryover(
    data: jnp.ndarray,
    ad_effect_retention_rate: jnp.ndarray,
    peak_effect_delay: jnp.ndarray,
    number_lags: int,
) -> jnp.ndarray:
  """Calculates media carryover.

  More details about this function can be found in:
  https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/46001.pdf

  Args:
    data: Input data. It is expected that data has either 2 dimensions for
      national models and 3 for geo models.
    ad_effect_retention_rate: Retention rate of the advertisement effect.
      Default is 0.5.
    peak_effect_delay: Delay of the peak effect in the carryover function.
      Default is 1.
    number_lags: Number of lags to include in the carryover calculation. Default
      is 13.

  Returns:
    The carryover values for the given data with the given parameters.
  """
  lags_arange = jnp.expand_dims(
      jnp.arange(number_lags, dtype=jnp.float32), axis=-1)
  convolve_func = _carryover_convolve
  if data.ndim == 3:
    # Since _carryover_convolve is already vmaped in the decorator we only need
    # to vmap it once here to handle the geo level data. We keep the windows bi
    # dimensional also for three dims data and vmap over only the extra data
    # dimension.
    convolve_func = jax.vmap(
        fun=_carryover_convolve, in_axes=(2, None, None), out_axes=2)
  weights = ad_effect_retention_rate**((lags_arange - peak_effect_delay)**2)
  return convolve_func(data, weights, number_lags)

def carryover(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    number_lags: int = 13,
    prefix: str = "",
) -> jnp.ndarray:
  """Transforms the input data with the carryover function.

  Args:
    data: Media data to be transformed. It is expected to have 2 dims for
      national models and 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones.
    number_lags: Number of lags for the carryover function.
    prefix: Prefix to use in the variable name for Numpyro.

  Returns:
    The transformed media data.
  """

  default_priors = priors.get_default_priors()
  with numpyro.plate(
      name=f"{prefix}{priors.AD_EFFECT_RETENTION_RATE}_plate",
      size=data.shape[1]):
    ad_effect_retention_rate = numpyro.sample(
        name=f"{prefix}{priors.AD_EFFECT_RETENTION_RATE}",
        fn=custom_priors.get(priors.AD_EFFECT_RETENTION_RATE,
                             default_priors[priors.AD_EFFECT_RETENTION_RATE]))

  with numpyro.plate(
      name=f"{prefix}{priors.PEAK_EFFECT_DELAY}_plate", size=data.shape[1]):
    peak_effect_delay = numpyro.sample(
        name=f"{prefix}{priors.PEAK_EFFECT_DELAY}",
        fn=custom_priors.get(priors.PEAK_EFFECT_DELAY,
                             default_priors[priors.PEAK_EFFECT_DELAY]))

  return _carryover(
      data=data,
      ad_effect_retention_rate=ad_effect_retention_rate,
      peak_effect_delay=peak_effect_delay,
      number_lags=number_lags)


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.

"""Set of core and modelling lagging functions."""

import functools
from typing import Mapping, Union

import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from lightweight_mmm.core import priors

@functools.partial(jax.vmap, in_axes=(1, 1, None), out_axes=1)
def _carryover_convolve(data: jnp.ndarray, weights: jnp.ndarray,
                        number_lags: int) -> jnp.ndarray:
  """Applies the convolution between the data and the weights for the carryover.

  Args:
    data: Input data.
    weights: Window weights for the carryover.
    number_lags: Number of lags the window has.

  Returns:
    The result values from convolving the data and the weights with padding.
  """
  window = jnp.concatenate([jnp.zeros(number_lags - 1), weights])
  return jax.scipy.signal.convolve(data, window, mode="same") / weights.sum()

@functools.partial(jax.jit, static_argnames=("number_lags",))
def _carryover(
    data: jnp.ndarray,
    ad_effect_retention_rate: jnp.ndarray,
    peak_effect_delay: jnp.ndarray,
    number_lags: int,
) -> jnp.ndarray:
  """Calculates media carryover.

  More details about this function can be found in:
  https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/46001.pdf

  Args:
    data: Input data. It is expected that data has either 2 dimensions for
      national models and 3 for geo models.
    ad_effect_retention_rate: Retention rate of the advertisement effect.
      Default is 0.5.
    peak_effect_delay: Delay of the peak effect in the carryover function.
      Default is 1.
    number_lags: Number of lags to include in the carryover calculation. Default
      is 13.

  Returns:
    The carryover values for the given data with the given parameters.
  """
  lags_arange = jnp.expand_dims(
      jnp.arange(number_lags, dtype=jnp.float32), axis=-1)
  convolve_func = _carryover_convolve
  if data.ndim == 3:
    # Since _carryover_convolve is already vmaped in the decorator we only need
    # to vmap it once here to handle the geo level data. We keep the windows bi
    # dimensional also for three dims data and vmap over only the extra data
    # dimension.
    convolve_func = jax.vmap(
        fun=_carryover_convolve, in_axes=(2, None, None), out_axes=2)
  weights = ad_effect_retention_rate**((lags_arange - peak_effect_delay)**2)
  return convolve_func(data, weights, number_lags)

def carryover(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    number_lags: int = 13,
    prefix: str = "",
) -> jnp.ndarray:
  """Transforms the input data with the carryover function.

  Args:
    data: Media data to be transformed. It is expected to have 2 dims for
      national models and 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones.
    number_lags: Number of lags for the carryover function.
    prefix: Prefix to use in the variable name for Numpyro.

  Returns:
    The transformed media data.
  """
  default_priors = priors.get_default_priors()
  with numpyro.plate(
      name=f"{prefix}{priors.AD_EFFECT_RETENTION_RATE}_plate",
      size=data.shape[1]):
    ad_effect_retention_rate = numpyro.sample(
        name=f"{prefix}{priors.AD_EFFECT_RETENTION_RATE}",
        fn=custom_priors.get(priors.AD_EFFECT_RETENTION_RATE,
                             default_priors[priors.AD_EFFECT_RETENTION_RATE]))

  with numpyro.plate(
      name=f"{prefix}{priors.PEAK_EFFECT_DELAY}_plate", size=data.shape[1]):
    peak_effect_delay = numpyro.sample(
        name=f"{prefix}{priors.PEAK_EFFECT_DELAY}",
        fn=custom_priors.get(priors.PEAK_EFFECT_DELAY,
                             default_priors[priors.PEAK_EFFECT_DELAY]))

  return _carryover(
      data=data,
      ad_effect_retention_rate=ad_effect_retention_rate,
      peak_effect_delay=peak_effect_delay,
      number_lags=number_lags)

@jax.jit
def _adstock(
    data: jnp.ndarray,
    lag_weight: Union[float, jnp.ndarray] = .9,
    normalise: bool = True,
) -> jnp.ndarray:
  """Calculates the adstock value of a given array.

  To learn more about advertising lag:
  https://en.wikipedia.org/wiki/Advertising_adstock

  Args:
    data: Input array.
    lag_weight: lag_weight effect of the adstock function. Default is 0.9.
    normalise: Whether to normalise the output value. This normalization will
      divide the output values by (1 / (1 - lag_weight)).

  Returns:
    The adstock output of the input array.
  """

  def adstock_internal(
      prev_adstock: jnp.ndarray,
      data: jnp.ndarray,
      lag_weight: Union[float, jnp.ndarray] = lag_weight,
  ) -> jnp.ndarray:
    adstock_value = prev_adstock * lag_weight + data
    return adstock_value, adstock_value# jax-ndarray

  _, adstock_values = jax.lax.scan(
      f=adstock_internal, init=data[0, ...], xs=data[1:, ...])
  adstock_values = jnp.concatenate([jnp.array([data[0, ...]]), adstock_values])
  return jax.lax.cond(
      normalise,
      lambda adstock_values: adstock_values / (1. / (1 - lag_weight)),
      lambda adstock_values: adstock_values,
      operand=adstock_values)


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.

"""Set of core and modelling lagging functions."""

import functools
from typing import Mapping, Union

import jax
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from lightweight_mmm.core import priors

@functools.partial(jax.vmap, in_axes=(1, 1, None), out_axes=1)
def _carryover_convolve(data: jnp.ndarray, weights: jnp.ndarray,
                        number_lags: int) -> jnp.ndarray:
  """Applies the convolution between the data and the weights for the carryover.

  Args:
    data: Input data.
    weights: Window weights for the carryover.
    number_lags: Number of lags the window has.

  Returns:
    The result values from convolving the data and the weights with padding.
  """
  window = jnp.concatenate([jnp.zeros(number_lags - 1), weights])
  return jax.scipy.signal.convolve(data, window, mode="same") / weights.sum()

@functools.partial(jax.jit, static_argnames=("number_lags",))
def _carryover(
    data: jnp.ndarray,
    ad_effect_retention_rate: jnp.ndarray,
    peak_effect_delay: jnp.ndarray,
    number_lags: int,
) -> jnp.ndarray:
  """Calculates media carryover.

  More details about this function can be found in:
  https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/46001.pdf

  Args:
    data: Input data. It is expected that data has either 2 dimensions for
      national models and 3 for geo models.
    ad_effect_retention_rate: Retention rate of the advertisement effect.
      Default is 0.5.
    peak_effect_delay: Delay of the peak effect in the carryover function.
      Default is 1.
    number_lags: Number of lags to include in the carryover calculation. Default
      is 13.

  Returns:
    The carryover values for the given data with the given parameters.
  """
  lags_arange = jnp.expand_dims(
      jnp.arange(number_lags, dtype=jnp.float32), axis=-1)
  convolve_func = _carryover_convolve
  if data.ndim == 3:
    # Since _carryover_convolve is already vmaped in the decorator we only need
    # to vmap it once here to handle the geo level data. We keep the windows bi
    # dimensional also for three dims data and vmap over only the extra data
    # dimension.
    convolve_func = jax.vmap(
        fun=_carryover_convolve, in_axes=(2, None, None), out_axes=2)
  weights = ad_effect_retention_rate**((lags_arange - peak_effect_delay)**2)
  return convolve_func(data, weights, number_lags)

def carryover(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    number_lags: int = 13,
    prefix: str = "",
) -> jnp.ndarray:
  """Transforms the input data with the carryover function.

  Args:
    data: Media data to be transformed. It is expected to have 2 dims for
      national models and 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones.
    number_lags: Number of lags for the carryover function.
    prefix: Prefix to use in the variable name for Numpyro.

  Returns:
    The transformed media data.
  """
  default_priors = priors.get_default_priors()
  with numpyro.plate(
      name=f"{prefix}{priors.AD_EFFECT_RETENTION_RATE}_plate",
      size=data.shape[1]):
    ad_effect_retention_rate = numpyro.sample(
        name=f"{prefix}{priors.AD_EFFECT_RETENTION_RATE}",
        fn=custom_priors.get(priors.AD_EFFECT_RETENTION_RATE,
                             default_priors[priors.AD_EFFECT_RETENTION_RATE]))

  with numpyro.plate(
      name=f"{prefix}{priors.PEAK_EFFECT_DELAY}_plate", size=data.shape[1]):
    peak_effect_delay = numpyro.sample(
        name=f"{prefix}{priors.PEAK_EFFECT_DELAY}",
        fn=custom_priors.get(priors.PEAK_EFFECT_DELAY,
                             default_priors[priors.PEAK_EFFECT_DELAY]))

  return _carryover(
      data=data,
      ad_effect_retention_rate=ad_effect_retention_rate,
      peak_effect_delay=peak_effect_delay,
      number_lags=number_lags)

@jax.jit
def _adstock(
    data: jnp.ndarray,
    lag_weight: Union[float, jnp.ndarray] = .9,
    normalise: bool = True,
) -> jnp.ndarray:
  """Calculates the adstock value of a given array.

  To learn more about advertising lag:
  https://en.wikipedia.org/wiki/Advertising_adstock

  Args:
    data: Input array.
    lag_weight: lag_weight effect of the adstock function. Default is 0.9.
    normalise: Whether to normalise the output value. This normalization will
      divide the output values by (1 / (1 - lag_weight)).

  Returns:
    The adstock output of the input array.
  """

  def adstock_internal(
      prev_adstock: jnp.ndarray,
      data: jnp.ndarray,
      lag_weight: Union[float, jnp.ndarray] = lag_weight,
  ) -> jnp.ndarray:
    adstock_value = prev_adstock * lag_weight + data
    return adstock_value, adstock_value# jax-ndarray

  _, adstock_values = jax.lax.scan(
      f=adstock_internal, init=data[0, ...], xs=data[1:, ...])
  adstock_values = jnp.concatenate([jnp.array([data[0, ...]]), adstock_values])
  return jax.lax.cond(
      normalise,
      lambda adstock_values: adstock_values / (1. / (1 - lag_weight)),
      lambda adstock_values: adstock_values,
      operand=adstock_values)

def adstock(
    data: jnp.ndarray,
    custom_priors: Mapping[str, dist.Distribution],
    *,
    normalise: bool = True,
    prefix: str = "",
) -> jnp.ndarray:
  """Transforms the input data with the adstock function and exponent.

  Args:
    data: Media data to be transformed. It is expected to have 2 dims for
      national models and 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones. The possible names of parameters for adstock and exponent
      are "lag_weight" and "exponent".
    normalise: Whether to normalise the output values.
    prefix: Prefix to use in the variable name for Numpyro.

  Returns:
    The transformed media data.
  """

  default_priors = priors.get_default_priors()
  with numpyro.plate(
      name=f"{prefix}{priors.LAG_WEIGHT}_plate", size=data.shape[1]):
    lag_weight = numpyro.sample(
        name=f"{prefix}{priors.LAG_WEIGHT}",
        fn=custom_priors.get(priors.LAG_WEIGHT,
                             default_priors[priors.LAG_WEIGHT]))

  if data.ndim == 3:
    lag_weight = jnp.expand_dims(lag_weight, axis=-1)

  return _adstock(data=data, lag_weight=lag_weight, normalise=normalise)


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 lagging."""

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 handlers
import numpyro.distributions as dist

from lightweight_mmm.core import priors
from lightweight_mmm.core.transformations import lagging

class LaggingTest(parameterized.TestCase):

  @parameterized.named_parameters(
      dict(
          testcase_name="national",
          data_shape=(150, 3),
          ad_effect_retention_rate_shape=(3,),
          peak_effect_delay_shape=(3,),
          number_lags=13,
      ),
      dict(
          testcase_name="geo",
          data_shape=(150, 3, 5),
          ad_effect_retention_rate_shape=(3,),
          peak_effect_delay_shape=(3,),
          number_lags=13,
      ),
  )
  def test_core_carryover_produces_correct_shape(
      self,
      data_shape,
      ad_effect_retention_rate_shape,
      peak_effect_delay_shape,
      number_lags,
  ):
    data = jnp.ones(data_shape)
    ad_effect_retention_rate = jnp.ones(ad_effect_retention_rate_shape)
    peak_effect_delay = jnp.ones(peak_effect_delay_shape)

    output = lagging._carryover(
        data=data,
        ad_effect_retention_rate=ad_effect_retention_rate,
        peak_effect_delay=peak_effect_delay,
        number_lags=number_lags,
    )

    self.assertEqual(output.shape, data_shape)

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

    def mock_model_function(data, number_lags):

      numpyro.deterministic(
          "carryover",
          lagging.carryover(
              data=data, custom_priors={}, number_lags=number_lags))


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 lagging."""

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 handlers
import numpyro.distributions as dist

from lightweight_mmm.core import priors
from lightweight_mmm.core.transformations import lagging

class LaggingTest(parameterized.TestCase):

  @parameterized.named_parameters(
      dict(
          testcase_name="national",
          data_shape=(150, 3),
          ad_effect_retention_rate_shape=(3,),
          peak_effect_delay_shape=(3,),
          number_lags=13,
      ),
      dict(
          testcase_name="geo",
          data_shape=(150, 3, 5),
          ad_effect_retention_rate_shape=(3,),
          peak_effect_delay_shape=(3,),
          number_lags=13,
      ),
  )
  def test_core_carryover_produces_correct_shape(
      self,
      data_shape,
      ad_effect_retention_rate_shape,
      peak_effect_delay_shape,
      number_lags,
  ):
    data = jnp.ones(data_shape)
    ad_effect_retention_rate = jnp.ones(ad_effect_retention_rate_shape)
    peak_effect_delay = jnp.ones(peak_effect_delay_shape)

    output = lagging._carryover(
        data=data,
        ad_effect_retention_rate=ad_effect_retention_rate,
        peak_effect_delay=peak_effect_delay,
        number_lags=number_lags,
    )

    self.assertEqual(output.shape, data_shape)

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

    def mock_model_function(data, number_lags):
      numpyro.deterministic(
          "carryover",
          lagging.carryover(
              data=data, custom_priors={}, number_lags=number_lags))

    num_samples = 10
    data = jnp.ones(data_shape)
    number_lags = 15
    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, number_lags=number_lags)
    carryover_values = mcmc.get_samples()["carryover"]

    self.assertEqual(carryover_values.shape, (num_samples, *data.shape))

  @parameterized.named_parameters(
      dict(
          testcase_name="ad_effect_retention_rate",
          prior_name=priors.AD_EFFECT_RETENTION_RATE,
      ),
      dict(
          testcase_name="peak_effect_delay",
          prior_name=priors.PEAK_EFFECT_DELAY,
      ),
  )
  def test_carryover_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))
    number_lags = 13

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

    used_distribution = values_and_dists[prior_name]
    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="national",
          data_shape=(150, 3),
          lag_weight_shape=(3,),
      ),
      dict(
          testcase_name="geo",
          data_shape=(150, 3, 5),
          lag_weight_shape=(3, 1),
      ),
  )
  def test_core_adstock_produces_correct_shape(self, data_shape,
                                               lag_weight_shape):
    data = jnp.ones(data_shape)
    lag_weight = jnp.ones(lag_weight_shape)

    output = lagging._adstock(data=data, lag_weight=lag_weight)

    self.assertEqual(output.shape, data_shape)

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

    def mock_model_function(data, normalise):

      numpyro.deterministic(
          "adstock",
          lagging.adstock(data=data, custom_priors={}, normalise=normalise))


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.

"""Module containing the different models available in the lightweightMMM lib.

Currently this file contains a main model with three possible options for
processing the media data. Which essentially grants the possibility of building
three different models.
  - Adstock
  - Hill-Adstock
  - Carryover
"""
import sys
#  pylint: disable=g-import-not-at-top
if sys.version_info >= (3, 8):
  from typing import Protocol
else:
  from typing_extensions import Protocol

from typing import Any, Dict, Mapping, MutableMapping, Optional, Sequence, Union

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

from lightweight_mmm import media_transforms

Prior = Union[
    dist.Distribution,
    Dict[str, float],
    Sequence[float],
    float
]

class TransformFunction(Protocol):

  def __call__(
      self,
      media_data: jnp.ndarray,
      custom_priors: MutableMapping[str, Prior],
      **kwargs: Any) -> jnp.ndarray:
    ...

_INTERCEPT = "intercept"
_COEF_TREND = "coef_trend"
_EXPO_TREND = "expo_trend"
_SIGMA = "sigma"
_GAMMA_SEASONALITY = "gamma_seasonality"
_WEEKDAY = "weekday"
_COEF_EXTRA_FEATURES = "coef_extra_features"
_COEF_SEASONALITY = "coef_seasonality"

MODEL_PRIORS_NAMES = frozenset((
    _INTERCEPT,
    _COEF_TREND,
    _EXPO_TREND,
    _SIGMA,
    _GAMMA_SEASONALITY,
    _WEEKDAY,
    _COEF_EXTRA_FEATURES,
    _COEF_SEASONALITY))

_EXPONENT = "exponent"
_LAG_WEIGHT = "lag_weight"
_HALF_MAX_EFFECTIVE_CONCENTRATION = "half_max_effective_concentration"
_SLOPE = "slope"
_AD_EFFECT_RETENTION_RATE = "ad_effect_retention_rate"
_PEAK_EFFECT_DELAY = "peak_effect_delay"

TRANSFORM_PRIORS_NAMES = immutabledict.immutabledict({
    "carryover":
        frozenset((_AD_EFFECT_RETENTION_RATE, _PEAK_EFFECT_DELAY, _EXPONENT)),
    "adstock":
        frozenset((_EXPONENT, _LAG_WEIGHT)),
    "hill_adstock":
        frozenset((_LAG_WEIGHT, _HALF_MAX_EFFECTIVE_CONCENTRATION, _SLOPE))
})

GEO_ONLY_PRIORS = frozenset((_COEF_SEASONALITY,))

def _get_default_priors() -> Mapping[str, Prior]:
  # Since JAX cannot be called before absl.app.run in tests we get default
  # priors from a function.

  return immutabledict.immutabledict({
      _INTERCEPT: dist.HalfNormal(scale=2.),
      _COEF_TREND: dist.Normal(loc=0., scale=1.),
      _EXPO_TREND: dist.Uniform(low=0.5, high=1.5),
      _SIGMA: dist.Gamma(concentration=1., rate=1.),
      _GAMMA_SEASONALITY: dist.Normal(loc=0., scale=1.),
      _WEEKDAY: dist.Normal(loc=0., scale=.5),
      _COEF_EXTRA_FEATURES: dist.Normal(loc=0., scale=1.),
      _COEF_SEASONALITY: dist.HalfNormal(scale=.5)
  })


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.

"""Module containing the different models available in the lightweightMMM lib.

Currently this file contains a main model with three possible options for
processing the media data. Which essentially grants the possibility of building
three different models.
  - Adstock
  - Hill-Adstock
  - Carryover
"""
import sys
#  pylint: disable=g-import-not-at-top
if sys.version_info >= (3, 8):
  from typing import Protocol
else:
  from typing_extensions import Protocol

from typing import Any, Dict, Mapping, MutableMapping, Optional, Sequence, Union

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

from lightweight_mmm import media_transforms

Prior = Union[
    dist.Distribution,
    Dict[str, float],
    Sequence[float],
    float
]

class TransformFunction(Protocol):

  def __call__(
      self,
      media_data: jnp.ndarray,
      custom_priors: MutableMapping[str, Prior],
      **kwargs: Any) -> jnp.ndarray:
    ...

_INTERCEPT = "intercept"
_COEF_TREND = "coef_trend"
_EXPO_TREND = "expo_trend"
_SIGMA = "sigma"
_GAMMA_SEASONALITY = "gamma_seasonality"
_WEEKDAY = "weekday"
_COEF_EXTRA_FEATURES = "coef_extra_features"
_COEF_SEASONALITY = "coef_seasonality"

MODEL_PRIORS_NAMES = frozenset((
    _INTERCEPT,
    _COEF_TREND,
    _EXPO_TREND,
    _SIGMA,
    _GAMMA_SEASONALITY,
    _WEEKDAY,
    _COEF_EXTRA_FEATURES,
    _COEF_SEASONALITY))

_EXPONENT = "exponent"
_LAG_WEIGHT = "lag_weight"
_HALF_MAX_EFFECTIVE_CONCENTRATION = "half_max_effective_concentration"
_SLOPE = "slope"
_AD_EFFECT_RETENTION_RATE = "ad_effect_retention_rate"
_PEAK_EFFECT_DELAY = "peak_effect_delay"

TRANSFORM_PRIORS_NAMES = immutabledict.immutabledict({
    "carryover":
        frozenset((_AD_EFFECT_RETENTION_RATE, _PEAK_EFFECT_DELAY, _EXPONENT)),
    "adstock":
        frozenset((_EXPONENT, _LAG_WEIGHT)),
    "hill_adstock":
        frozenset((_LAG_WEIGHT, _HALF_MAX_EFFECTIVE_CONCENTRATION, _SLOPE))
})

GEO_ONLY_PRIORS = frozenset((_COEF_SEASONALITY,))

def _get_default_priors() -> Mapping[str, Prior]:
  # Since JAX cannot be called before absl.app.run in tests we get default
  # priors from a function.
  return immutabledict.immutabledict({
      _INTERCEPT: dist.HalfNormal(scale=2.),
      _COEF_TREND: dist.Normal(loc=0., scale=1.),
      _EXPO_TREND: dist.Uniform(low=0.5, high=1.5),
      _SIGMA: dist.Gamma(concentration=1., rate=1.),
      _GAMMA_SEASONALITY: dist.Normal(loc=0., scale=1.),
      _WEEKDAY: dist.Normal(loc=0., scale=.5),
      _COEF_EXTRA_FEATURES: dist.Normal(loc=0., scale=1.),
      _COEF_SEASONALITY: dist.HalfNormal(scale=.5)
  })

def _get_transform_default_priors() -> Mapping[str, Prior]:
  # Since JAX cannot be called before absl.app.run in tests we get default
  # priors from a function.

  return immutabledict.immutabledict({
      "carryover":
          immutabledict.immutabledict({
              _AD_EFFECT_RETENTION_RATE:
                  dist.Beta(concentration1=1., concentration0=1.),
              _PEAK_EFFECT_DELAY:
                  dist.HalfNormal(scale=2.),
              _EXPONENT:
                  dist.Beta(concentration1=9., concentration0=1.)
          }),
      "adstock":
          immutabledict.immutabledict({
              _EXPONENT: dist.Beta(concentration1=9., concentration0=1.),
              _LAG_WEIGHT: dist.Beta(concentration1=2., concentration0=1.)
          }),
      "hill_adstock":
          immutabledict.immutabledict({
              _LAG_WEIGHT:
                  dist.Beta(concentration1=2., concentration0=1.),
              _HALF_MAX_EFFECTIVE_CONCENTRATION:
                  dist.Gamma(concentration=1., rate=1.),
              _SLOPE:
                  dist.Gamma(concentration=1., rate=1.)
          })
  })


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.

"""Module containing the different models available in the lightweightMMM lib.

Currently this file contains a main model with three possible options for
processing the media data. Which essentially grants the possibility of building
three different models.
  - Adstock
  - Hill-Adstock
  - Carryover
"""
import sys
#  pylint: disable=g-import-not-at-top
if sys.version_info >= (3, 8):
  from typing import Protocol
else:
  from typing_extensions import Protocol

from typing import Any, Dict, Mapping, MutableMapping, Optional, Sequence, Union

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

from lightweight_mmm import media_transforms

Prior = Union[
    dist.Distribution,
    Dict[str, float],
    Sequence[float],
    float
]

class TransformFunction(Protocol):

  def __call__(
      self,
      media_data: jnp.ndarray,
      custom_priors: MutableMapping[str, Prior],
      **kwargs: Any) -> jnp.ndarray:
    ...

_INTERCEPT = "intercept"
_COEF_TREND = "coef_trend"
_EXPO_TREND = "expo_trend"
_SIGMA = "sigma"
_GAMMA_SEASONALITY = "gamma_seasonality"
_WEEKDAY = "weekday"
_COEF_EXTRA_FEATURES = "coef_extra_features"
_COEF_SEASONALITY = "coef_seasonality"

MODEL_PRIORS_NAMES = frozenset((
    _INTERCEPT,
    _COEF_TREND,
    _EXPO_TREND,
    _SIGMA,
    _GAMMA_SEASONALITY,
    _WEEKDAY,
    _COEF_EXTRA_FEATURES,
    _COEF_SEASONALITY))

_EXPONENT = "exponent"
_LAG_WEIGHT = "lag_weight"
_HALF_MAX_EFFECTIVE_CONCENTRATION = "half_max_effective_concentration"
_SLOPE = "slope"
_AD_EFFECT_RETENTION_RATE = "ad_effect_retention_rate"
_PEAK_EFFECT_DELAY = "peak_effect_delay"

TRANSFORM_PRIORS_NAMES = immutabledict.immutabledict({
    "carryover":
        frozenset((_AD_EFFECT_RETENTION_RATE, _PEAK_EFFECT_DELAY, _EXPONENT)),
    "adstock":
        frozenset((_EXPONENT, _LAG_WEIGHT)),
    "hill_adstock":
        frozenset((_LAG_WEIGHT, _HALF_MAX_EFFECTIVE_CONCENTRATION, _SLOPE))
})

GEO_ONLY_PRIORS = frozenset((_COEF_SEASONALITY,))

def _get_default_priors() -> Mapping[str, Prior]:
  # Since JAX cannot be called before absl.app.run in tests we get default
  # priors from a function.
  return immutabledict.immutabledict({
      _INTERCEPT: dist.HalfNormal(scale=2.),
      _COEF_TREND: dist.Normal(loc=0., scale=1.),
      _EXPO_TREND: dist.Uniform(low=0.5, high=1.5),
      _SIGMA: dist.Gamma(concentration=1., rate=1.),
      _GAMMA_SEASONALITY: dist.Normal(loc=0., scale=1.),
      _WEEKDAY: dist.Normal(loc=0., scale=.5),
      _COEF_EXTRA_FEATURES: dist.Normal(loc=0., scale=1.),
      _COEF_SEASONALITY: dist.HalfNormal(scale=.5)
  })

def _get_transform_default_priors() -> Mapping[str, Prior]:
  # Since JAX cannot be called before absl.app.run in tests we get default
  # priors from a function.
  return immutabledict.immutabledict({
      "carryover":
          immutabledict.immutabledict({
              _AD_EFFECT_RETENTION_RATE:
                  dist.Beta(concentration1=1., concentration0=1.),
              _PEAK_EFFECT_DELAY:
                  dist.HalfNormal(scale=2.),
              _EXPONENT:
                  dist.Beta(concentration1=9., concentration0=1.)
          }),
      "adstock":
          immutabledict.immutabledict({
              _EXPONENT: dist.Beta(concentration1=9., concentration0=1.),
              _LAG_WEIGHT: dist.Beta(concentration1=2., concentration0=1.)
          }),
      "hill_adstock":
          immutabledict.immutabledict({
              _LAG_WEIGHT:
                  dist.Beta(concentration1=2., concentration0=1.),
              _HALF_MAX_EFFECTIVE_CONCENTRATION:
                  dist.Gamma(concentration=1., rate=1.),
              _SLOPE:
                  dist.Gamma(concentration=1., rate=1.)
          })
  })

def transform_adstock(media_data: jnp.ndarray,
                      custom_priors: MutableMapping[str, Prior],
                      normalise: bool = True) -> jnp.ndarray:
  """Transforms the input data with the adstock function and exponent.

  Args:
    media_data: Media data to be transformed. It is expected to have 2 dims for
      national models and 3 for geo models.
    custom_priors: The custom priors we want the model to take instead of the
      default ones. The possible names of parameters for adstock and exponent
      are "lag_weight" and "exponent".
    normalise: Whether to normalise the output values.

  Returns:
    The transformed media data.
  """

  transform_default_priors = _get_transform_default_priors()["adstock"]
  with numpyro.plate(name=f"{_LAG_WEIGHT}_plate",
                     size=media_data.shape[1]):
    lag_weight = numpyro.sample(
        name=_LAG_WEIGHT,
        fn=custom_priors.get(_LAG_WEIGHT,
                             transform_default_priors[_LAG_WEIGHT]))

  with numpyro.plate(name=f"{_EXPONENT}_plate",
                     size=media_data.shape[1]):
    exponent = numpyro.sample(
        name=_EXPONENT,
        fn=custom_priors.get(_EXPONENT,
                             transform_default_priors[_EXPONENT]))

  if media_data.ndim == 3:
    lag_weight = jnp.expand_dims(lag_weight, axis=-1)
    exponent = jnp.expand_dims(exponent, axis=-1)

  adstock = media_transforms.adstock(
      data=media_data, lag_weight=lag_weight, normalise=normalise)

  return media_transforms.apply_exponent_safe(data=adstock, exponent=exponent)


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 plot."""

from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import jax.numpy as jnp
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpyro.distributions as dist
import pandas as pd

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import models
from lightweight_mmm import plot
from lightweight_mmm import preprocessing

MOCK_NATIONAL_TRACE = {
    "coef_extra_features": np.ones([10, 2]),
    "coef_media": np.ones([10, 5]),
    "coef_trend": np.ones([10, 1]),
    "expo_trend": np.ones([10, 1]),
    "gamma_seasonality": np.ones([10, 3, 2]),
    "intercept": np.ones([10, 1]),
    "media_transformed": np.ones([10, 50, 5,]),
    "mu": np.ones([10, 50]),
    "sigma": np.ones([10, 1]),
    "ad_effect_retention_rate": np.ones([10, 5]),
    "exponent": np.ones([10, 5]),
    "half_max_effective_concentration": np.ones([10, 5]),
    "lag_weight": np.ones([10, 5]),
    "slope": np.ones([10, 5]),
    "peak_effect_delay": np.ones([10, 5]),
    }

MOCK_GEO_TRACE = {
    "channel_coef_media": np.ones([10, 5, 1]),
    "coef_extra_features": np.ones([10, 2, 3]),
    "coef_media": np.ones([10, 5, 3]),
    "coef_seasonality": np.ones([10, 3]),
    "coef_trend": np.ones([10, 3]),
    "expo_trend": np.ones([10, 1]),
    "gamma_seasonality": np.ones([10, 3, 2]),
    "intercept": np.ones([10, 3]),
    "media_transformed": np.ones([10, 50, 5, 3]),
    "mu": np.ones([10, 50, 3]),
    "sigma": np.ones([10, 3]),
    "ad_effect_retention_rate": np.ones([10, 5]),
    "exponent": np.ones([10, 5]),
    "half_max_effective_concentration": np.ones([10, 5]),
    "lag_weight": np.ones([10, 5]),
    "peak_effect_delay": np.ones([10, 5]),
    "slope": np.ones([10, 5]),
}

def _set_up_mock_mmm(model_name: str,
                     is_geo_model: bool) -> lightweight_mmm.LightweightMMM:
  """Creates a mock LightweightMMM instance that acts like a fitted model.

  These instances are used when we want to run tests on more diverse ranges of
  models than the two standard national_mmm and geo_mmm defined below but don't
  need the unit tests to spend time actually running the model fits.

  Args:
    model_name: One of ["adstock", "carryover", or "hill_adstock"], specifying
      which model type should be used in the mock LightweightMMM.
    is_geo_model: Whether to create a geo-level model (True) or a national-level
      model (False).

  Returns:
    mmm: A LightweightMMM object that can be treated like a fitted model
    for plotting-related unit tests.
  """

  initial_mock_trace = MOCK_GEO_TRACE if is_geo_model else MOCK_NATIONAL_TRACE
  all_model_names = {"adstock", "carryover", "hill_adstock"}
  model_items_to_delete = frozenset.union(*[
      models.TRANSFORM_PRIORS_NAMES[x]
      for x in all_model_names - {model_name}
  ]) - models.TRANSFORM_PRIORS_NAMES[model_name]
  mock_trace = {
      key: initial_mock_trace[key]
      for key in initial_mock_trace
      if key not in model_items_to_delete
  }
  mmm = lightweight_mmm.LightweightMMM(model_name=model_name)
  mmm.n_media_channels = 5
  mmm.n_geos = 3 if is_geo_model else 1
  mmm._media_prior = jnp.ones(5)
  mmm._weekday_seasonality = False
  mmm._degrees_seasonality = 3
  mmm.custom_priors = {}
  mmm._extra_features = None
  mmm.trace = mock_trace
  mmm.media = jnp.ones_like(mock_trace["media_transformed"][0])
  mmm.media_names = [f"channel_{i}" for i in range(5)]
  return mmm


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 plot."""

from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import jax.numpy as jnp
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpyro.distributions as dist
import pandas as pd

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import models
from lightweight_mmm import plot
from lightweight_mmm import preprocessing

MOCK_NATIONAL_TRACE = {
    "coef_extra_features": np.ones([10, 2]),
    "coef_media": np.ones([10, 5]),
    "coef_trend": np.ones([10, 1]),
    "expo_trend": np.ones([10, 1]),
    "gamma_seasonality": np.ones([10, 3, 2]),
    "intercept": np.ones([10, 1]),
    "media_transformed": np.ones([10, 50, 5,]),
    "mu": np.ones([10, 50]),
    "sigma": np.ones([10, 1]),
    "ad_effect_retention_rate": np.ones([10, 5]),
    "exponent": np.ones([10, 5]),
    "half_max_effective_concentration": np.ones([10, 5]),
    "lag_weight": np.ones([10, 5]),
    "slope": np.ones([10, 5]),
    "peak_effect_delay": np.ones([10, 5]),
    }

MOCK_GEO_TRACE = {
    "channel_coef_media": np.ones([10, 5, 1]),
    "coef_extra_features": np.ones([10, 2, 3]),
    "coef_media": np.ones([10, 5, 3]),
    "coef_seasonality": np.ones([10, 3]),
    "coef_trend": np.ones([10, 3]),
    "expo_trend": np.ones([10, 1]),
    "gamma_seasonality": np.ones([10, 3, 2]),
    "intercept": np.ones([10, 3]),
    "media_transformed": np.ones([10, 50, 5, 3]),
    "mu": np.ones([10, 50, 3]),
    "sigma": np.ones([10, 3]),
    "ad_effect_retention_rate": np.ones([10, 5]),
    "exponent": np.ones([10, 5]),
    "half_max_effective_concentration": np.ones([10, 5]),
    "lag_weight": np.ones([10, 5]),
    "peak_effect_delay": np.ones([10, 5]),
    "slope": np.ones([10, 5]),
}

def _set_up_mock_mmm(model_name: str,
                     is_geo_model: bool) -> lightweight_mmm.LightweightMMM:
  """Creates a mock LightweightMMM instance that acts like a fitted model.

  These instances are used when we want to run tests on more diverse ranges of
  models than the two standard national_mmm and geo_mmm defined below but don't
  need the unit tests to spend time actually running the model fits.

  Args:
    model_name: One of ["adstock", "carryover", or "hill_adstock"], specifying
      which model type should be used in the mock LightweightMMM.
    is_geo_model: Whether to create a geo-level model (True) or a national-level
      model (False).

  Returns:
    mmm: A LightweightMMM object that can be treated like a fitted model
    for plotting-related unit tests.
  """
  initial_mock_trace = MOCK_GEO_TRACE if is_geo_model else MOCK_NATIONAL_TRACE
  all_model_names = {"adstock", "carryover", "hill_adstock"}
  model_items_to_delete = frozenset.union(*[
      models.TRANSFORM_PRIORS_NAMES[x]
      for x in all_model_names - {model_name}
  ]) - models.TRANSFORM_PRIORS_NAMES[model_name]
  mock_trace = {
      key: initial_mock_trace[key]
      for key in initial_mock_trace
      if key not in model_items_to_delete
  }
  mmm = lightweight_mmm.LightweightMMM(model_name=model_name)
  mmm.n_media_channels = 5
  mmm.n_geos = 3 if is_geo_model else 1
  mmm._media_prior = jnp.ones(5)
  mmm._weekday_seasonality = False
  mmm._degrees_seasonality = 3
  mmm.custom_priors = {}
  mmm._extra_features = None
  mmm.trace = mock_trace
  mmm.media = jnp.ones_like(mock_trace["media_transformed"][0])
  mmm.media_names = [f"channel_{i}" for i in range(5)]
  return mmm

class PlotTest(parameterized.TestCase):

  @classmethod
  def setUpClass(cls):

    super(PlotTest, cls).setUpClass()
    cls.national_mmm = lightweight_mmm.LightweightMMM()
    cls.national_mmm.fit(
        media=jnp.ones((50, 5)),
        target=jnp.ones(50),
        media_prior=jnp.ones(5) * 50,
        number_warmup=2,
        number_samples=2,
        number_chains=1)
    cls.geo_mmm = lightweight_mmm.LightweightMMM()
    cls.geo_mmm.fit(
        media=jnp.ones((50, 5, 3)),
        target=jnp.ones((50, 3)),
        media_prior=jnp.ones(5) * 50,
        number_warmup=2,
        number_samples=2,
        number_chains=1)
    cls.not_fitted_mmm = lightweight_mmm.LightweightMMM()


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 plot."""

from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import jax.numpy as jnp
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import numpyro.distributions as dist
import pandas as pd

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import models
from lightweight_mmm import plot
from lightweight_mmm import preprocessing

MOCK_NATIONAL_TRACE = {
    "coef_extra_features": np.ones([10, 2]),
    "coef_media": np.ones([10, 5]),
    "coef_trend": np.ones([10, 1]),
    "expo_trend": np.ones([10, 1]),
    "gamma_seasonality": np.ones([10, 3, 2]),
    "intercept": np.ones([10, 1]),
    "media_transformed": np.ones([10, 50, 5,]),
    "mu": np.ones([10, 50]),
    "sigma": np.ones([10, 1]),
    "ad_effect_retention_rate": np.ones([10, 5]),
    "exponent": np.ones([10, 5]),
    "half_max_effective_concentration": np.ones([10, 5]),
    "lag_weight": np.ones([10, 5]),
    "slope": np.ones([10, 5]),
    "peak_effect_delay": np.ones([10, 5]),
    }

MOCK_GEO_TRACE = {
    "channel_coef_media": np.ones([10, 5, 1]),
    "coef_extra_features": np.ones([10, 2, 3]),
    "coef_media": np.ones([10, 5, 3]),
    "coef_seasonality": np.ones([10, 3]),
    "coef_trend": np.ones([10, 3]),
    "expo_trend": np.ones([10, 1]),
    "gamma_seasonality": np.ones([10, 3, 2]),
    "intercept": np.ones([10, 3]),
    "media_transformed": np.ones([10, 50, 5, 3]),
    "mu": np.ones([10, 50, 3]),
    "sigma": np.ones([10, 3]),
    "ad_effect_retention_rate": np.ones([10, 5]),
    "exponent": np.ones([10, 5]),
    "half_max_effective_concentration": np.ones([10, 5]),
    "lag_weight": np.ones([10, 5]),
    "peak_effect_delay": np.ones([10, 5]),
    "slope": np.ones([10, 5]),
}

def _set_up_mock_mmm(model_name: str,
                     is_geo_model: bool) -> lightweight_mmm.LightweightMMM:
  """Creates a mock LightweightMMM instance that acts like a fitted model.

  These instances are used when we want to run tests on more diverse ranges of
  models than the two standard national_mmm and geo_mmm defined below but don't
  need the unit tests to spend time actually running the model fits.

  Args:
    model_name: One of ["adstock", "carryover", or "hill_adstock"], specifying
      which model type should be used in the mock LightweightMMM.
    is_geo_model: Whether to create a geo-level model (True) or a national-level
      model (False).

  Returns:
    mmm: A LightweightMMM object that can be treated like a fitted model
    for plotting-related unit tests.
  """
  initial_mock_trace = MOCK_GEO_TRACE if is_geo_model else MOCK_NATIONAL_TRACE
  all_model_names = {"adstock", "carryover", "hill_adstock"}
  model_items_to_delete = frozenset.union(*[
      models.TRANSFORM_PRIORS_NAMES[x]
      for x in all_model_names - {model_name}
  ]) - models.TRANSFORM_PRIORS_NAMES[model_name]
  mock_trace = {
      key: initial_mock_trace[key]
      for key in initial_mock_trace
      if key not in model_items_to_delete
  }
  mmm = lightweight_mmm.LightweightMMM(model_name=model_name)
  mmm.n_media_channels = 5
  mmm.n_geos = 3 if is_geo_model else 1
  mmm._media_prior = jnp.ones(5)
  mmm._weekday_seasonality = False
  mmm._degrees_seasonality = 3
  mmm.custom_priors = {}
  mmm._extra_features = None
  mmm.trace = mock_trace
  mmm.media = jnp.ones_like(mock_trace["media_transformed"][0])
  mmm.media_names = [f"channel_{i}" for i in range(5)]
  return mmm

class PlotTest(parameterized.TestCase):

  @classmethod
  def setUpClass(cls):
    super(PlotTest, cls).setUpClass()
    cls.national_mmm = lightweight_mmm.LightweightMMM()
    cls.national_mmm.fit(
        media=jnp.ones((50, 5)),
        target=jnp.ones(50),
        media_prior=jnp.ones(5) * 50,
        number_warmup=2,
        number_samples=2,
        number_chains=1)
    cls.geo_mmm = lightweight_mmm.LightweightMMM()
    cls.geo_mmm.fit(
        media=jnp.ones((50, 5, 3)),
        target=jnp.ones((50, 3)),
        media_prior=jnp.ones(5) * 50,
        number_warmup=2,
        number_samples=2,
        number_chains=1)
    cls.not_fitted_mmm = lightweight_mmm.LightweightMMM()

  def setUp(self):

    super().setUp()
    self.mock_ax_scatter = self.enter_context(
        mock.patch.object(plot.plt.Axes, "scatter", autospec=True))
    self.mock_sns_lineplot = self.enter_context(
        mock.patch.object(plot.sns, "lineplot", autospec=True))
    self.mock_plt_plot = self.enter_context(
        mock.patch.object(plot.plt.Axes, "plot", autospec=True))
    self.mock_plt_barplot = self.enter_context(
        mock.patch.object(plot.plt.Axes, "bar", autospec=True))
    self.mock_pd_area_plot = self.enter_context(
        mock.patch.object(plot.pd.DataFrame.plot, "area", autospec=True))
    self.mock_sns_kdeplot = self.enter_context(
        mock.patch.object(plot.sns, "kdeplot", autospec=True))
    self.mock_plt_ax_legend = self.enter_context(
        mock.patch.object(plot.plt.Axes, "legend", autospec=True))


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.

"""Utilities for optimizing your media based on media mix models."""
import functools
from typing import Optional, Tuple, Union
from absl import logging
import jax
import jax.numpy as jnp
from scipy import optimize

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import preprocessing

@functools.partial(
    jax.jit,
    static_argnames=("media_mix_model", "media_input_shape", "target_scaler",
                     "media_scaler"))
def _objective_function(extra_features: jnp.ndarray,
                        media_mix_model: lightweight_mmm.LightweightMMM,
                        media_input_shape: Tuple[int,
                                                 int], media_gap: Optional[int],
                        target_scaler: Optional[preprocessing.CustomScaler],
                        media_scaler: preprocessing.CustomScaler,
                        geo_ratio: jnp.array,
                        seed: Optional[int],
                        media_values: jnp.ndarray) -> jnp.float64:
  """Objective function to calculate the sum of all predictions of the model.

  Args:
    extra_features: Extra features the model requires for prediction.
    media_mix_model: Media mix model to use. Must have a predict method to be
      used.
    media_input_shape: Input shape of the data required by the model to get
      predictions. This is needed since optimization might flatten some arrays
      and they need to be reshaped before running new predictions.
    media_gap: Media data gap between the end of training data and the start of
      the out of sample media given. Eg. if 100 weeks of data were used for
      training and prediction starts 2 months after training data finished we
      need to provide the 8 weeks missing between the training data and the
      prediction data so data transformations (adstock, carryover, ...) can take
      place correctly.
    target_scaler: Scaler that was used to scale the target before training.
    media_scaler: Scaler that was used to scale the media data before training.
    geo_ratio: The ratio to split channel media across geo. Should sum up to 1
      for each channel and should have shape (c, g).
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.
    media_values: Media values required by the model to run predictions.

  Returns:
    The negative value of the sum of all predictions.
  """

  if hasattr(media_mix_model, "n_geos") and media_mix_model.n_geos > 1:
    media_values = geo_ratio * jnp.expand_dims(media_values, axis=-1)
  media_values = jnp.tile(
      media_values / media_input_shape[0], reps=media_input_shape[0])
  # Distribute budget of each channels across time.
  media_values = jnp.reshape(a=media_values, newshape=media_input_shape)
  media_values = media_scaler.transform(media_values)
  return -jnp.sum(
      media_mix_model.predict(
          media=media_values.reshape(media_input_shape),
          extra_features=extra_features,
          media_gap=media_gap,
          target_scaler=target_scaler,
          seed=seed).mean(axis=0))


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.

"""Utilities for optimizing your media based on media mix models."""
import functools
from typing import Optional, Tuple, Union
from absl import logging
import jax
import jax.numpy as jnp
from scipy import optimize

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import preprocessing

@functools.partial(
    jax.jit,
    static_argnames=("media_mix_model", "media_input_shape", "target_scaler",
                     "media_scaler"))
def _objective_function(extra_features: jnp.ndarray,
                        media_mix_model: lightweight_mmm.LightweightMMM,
                        media_input_shape: Tuple[int,
                                                 int], media_gap: Optional[int],
                        target_scaler: Optional[preprocessing.CustomScaler],
                        media_scaler: preprocessing.CustomScaler,
                        geo_ratio: jnp.array,
                        seed: Optional[int],
                        media_values: jnp.ndarray) -> jnp.float64:
  """Objective function to calculate the sum of all predictions of the model.

  Args:
    extra_features: Extra features the model requires for prediction.
    media_mix_model: Media mix model to use. Must have a predict method to be
      used.
    media_input_shape: Input shape of the data required by the model to get
      predictions. This is needed since optimization might flatten some arrays
      and they need to be reshaped before running new predictions.
    media_gap: Media data gap between the end of training data and the start of
      the out of sample media given. Eg. if 100 weeks of data were used for
      training and prediction starts 2 months after training data finished we
      need to provide the 8 weeks missing between the training data and the
      prediction data so data transformations (adstock, carryover, ...) can take
      place correctly.
    target_scaler: Scaler that was used to scale the target before training.
    media_scaler: Scaler that was used to scale the media data before training.
    geo_ratio: The ratio to split channel media across geo. Should sum up to 1
      for each channel and should have shape (c, g).
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.
    media_values: Media values required by the model to run predictions.

  Returns:
    The negative value of the sum of all predictions.
  """
  if hasattr(media_mix_model, "n_geos") and media_mix_model.n_geos > 1:
    media_values = geo_ratio * jnp.expand_dims(media_values, axis=-1)
  media_values = jnp.tile(
      media_values / media_input_shape[0], reps=media_input_shape[0])
  # Distribute budget of each channels across time.
  media_values = jnp.reshape(a=media_values, newshape=media_input_shape)
  media_values = media_scaler.transform(media_values)
  return -jnp.sum(
      media_mix_model.predict(
          media=media_values.reshape(media_input_shape),
          extra_features=extra_features,
          media_gap=media_gap,
          target_scaler=target_scaler,
          seed=seed).mean(axis=0))

@jax.jit
def _budget_constraint(media: jnp.ndarray,
                       prices: jnp.ndarray,
                       budget: jnp.ndarray) -> jnp.float64:
  """Calculates optimization constraint to keep spend equal to the budget.

  Args:
    media: Array with the values of the media for this iteration.
    prices: Prices of each media channel at any given time.
    budget: Total budget of the optimization.

  Returns:
    The result from substracting the total spending and the budget.
  """
  media = media.reshape((-1, len(prices)))
  return jnp.sum(media * prices) - budget

def _get_lower_and_upper_bounds(
    media: jnp.ndarray,
    n_time_periods: int,
    lower_pct: jnp.ndarray,
    upper_pct: jnp.ndarray,
    media_scaler: Optional[preprocessing.CustomScaler] = None
) -> optimize.Bounds:
  """Gets the lower and upper bounds for optimisation based on historic data.

  It creates an upper bound based on a percentage above the mean value on
  each channel and a lower bound based on a relative decrease of the mean
  value.

  Args:
    media: Media data to get historic mean.
    n_time_periods: Number of time periods to optimize for. If model is built on
      weekly data, this would be the number of weeks ahead to optimize.
    lower_pct: Relative percentage decrease from the mean value to consider as
      new lower bound.
    upper_pct: Relative percentage increase from the mean value to consider as
      new upper bound.
    media_scaler: Scaler that was used to scale the media data before training.

  Returns:
    A list of tuples with the lower and upper bound for each media channel.
  """

  if media.ndim == 3:
    lower_pct = jnp.expand_dims(lower_pct, axis=-1)
    upper_pct = jnp.expand_dims(upper_pct, axis=-1)

  mean_data = media.mean(axis=0)
  lower_bounds = jnp.maximum(mean_data * (1 - lower_pct), 0)
  upper_bounds = mean_data * (1 + upper_pct)

  if media_scaler:
    lower_bounds = media_scaler.inverse_transform(lower_bounds)
    upper_bounds = media_scaler.inverse_transform(upper_bounds)

  if media.ndim == 3:
    lower_bounds = lower_bounds.sum(axis=-1)
    upper_bounds = upper_bounds.sum(axis=-1)

  return optimize.Bounds(lb=lower_bounds * n_time_periods,
                         ub=upper_bounds * n_time_periods)