Prompt:
def save_model(
    media_mix_model: Any,
    file_path: str
    ) -> None:
  """Saves the given model in the given path.

  Args:
    media_mix_model: Model to save on disk.
    file_path: File path where the model should be placed.
  """
  with gfile.GFile(file_path, "wb") as file:
    pickle.dump(obj=media_mix_model, file=file)

def load_model(file_path: str) -> Any:
  """Loads a model given a string path.

  Args:
    file_path: Path of the file containing the model.

  Returns:
    The LightweightMMM object that was stored in the given path.
  """
  with gfile.GFile(file_path, "rb") as file:
    media_mix_model = pickle.load(file=file)

  for attr in dir(media_mix_model):
    if attr.startswith("__"):
      continue
    attr_value = getattr(media_mix_model, attr)
    if isinstance(attr_value, np.ndarray):
      setattr(media_mix_model, attr, jnp.array(attr_value))

  return media_mix_model

def get_time_seed() -> int:
  """Generates an integer using the last decimals of time.time().

  Returns:
    Integer to be used as seed.
  """
  # time.time() has the following format: 1645174953.0429401
  return int(str(time.time()).split(".")[1])

def simulate_dummy_data(
    data_size: int,
    n_media_channels: int,
    n_extra_features: int,
    geos: int = 1,
    seed: int = 5
    ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
  """Simulates dummy data needed for media mix modelling.

  This function's goal is to be super simple and not have many parameters,
  although it does not generate a fully realistic dataset is only meant to be
  used for demos/tutorial purposes. Uses carryover for lagging but has no
  saturation and no trend.

  The data simulated includes the media data, extra features, a target/KPI and
  costs.

  Args:
    data_size: Number of rows to generate.
    n_media_channels: Number of media channels to generate.
    n_extra_features: Number of extra features to generate.
    geos: Number of geos for geo level data (default = 1 for national).
    seed: Random seed.

  Returns:
    The simulated media, extra features, target and costs.
  """
  if data_size < 1 or n_media_channels < 1 or n_extra_features < 1:
    raise ValueError(
        "Data size, n_media_channels and n_extra_features must be greater than"
        " 0. Please check the values introduced are greater than zero.")
  data_offset = int(data_size * 0.2)
  data_size += data_offset
  key = random.PRNGKey(seed)
  sub_keys = random.split(key=key, num=7)
  media_data = random.normal(key=sub_keys[0],
                             shape=(data_size, n_media_channels)) * 1.5 + 20

  extra_features = random.normal(key=sub_keys[1],
                                 shape=(data_size, n_extra_features)) + 5
  # Reduce the costs to make ROI realistic.
  costs = media_data[data_offset:].sum(axis=0) * .1

  seasonality = media_transforms.calculate_seasonality(
      number_periods=data_size,
      degrees=2,
      frequency=52,
      gamma_seasonality=1)
  target_noise = random.normal(key=sub_keys[2], shape=(data_size,)) + 3

  # media_data_transformed = media_transforms.adstock(media_data)
  media_data_transformed = media_transforms.carryover(
      data=media_data,
      ad_effect_retention_rate=jnp.full((n_media_channels,), fill_value=.5),
      peak_effect_delay=jnp.full((n_media_channels,), fill_value=1.))
  beta_media = random.normal(key=sub_keys[3], shape=(n_media_channels,)) + 1
  beta_extra_features = random.normal(key=sub_keys[4],
                                      shape=(n_extra_features,))
  # There is no trend to keep this very simple.
  target = 15 + seasonality + media_data_transformed.dot(
      beta_media) + extra_features.dot(beta_extra_features) + target_noise

  logging.info("Correlation between transformed media and target")
  logging.info([
      np.corrcoef(target[data_offset:], media_data_transformed[data_offset:,
                                                               i])[0, 1]
      for i in range(n_media_channels)
  ])

  logging.info("True ROI for media channels")
  logging.info([
      sum(media_data_transformed[data_offset:, i] * beta_media[i]) / costs[i]
      for i in range(n_media_channels)
  ])

  if geos > 1:
    # Distribute national data to geo and add some more noise.
    weights = random.uniform(key=sub_keys[5], shape=(1, geos))
    weights /= sum(weights)
    target_noise = random.normal(key=sub_keys[6], shape=(data_size, geos)) * .5
    target = target[:, np.newaxis].dot(weights) + target_noise
    media_data = media_data[:, :, np.newaxis].dot(weights)
    extra_features = extra_features[:, :, np.newaxis].dot(weights)

  return (media_data[data_offset:], extra_features[data_offset:],
          target[data_offset:], costs)

def _split_array_into_list(
    dataframe: pd.DataFrame,
    split_level_feature: str,
    features: List[str],
    national_model_flag: bool = True) -> List[np.ndarray]:
  """Splits data frame into list of jax arrays.

  Args:
    dataframe: Dataframe with all the modeling feature.
    split_level_feature: Feature that will be used to split.
    features: List of feature to export from data frame.
    national_model_flag: Whether the data frame is used for national model.

  Returns:
    List of jax arrays.
  """
  split_level = dataframe[split_level_feature].unique()
  array_list_by_level = [
      dataframe.loc[dataframe[split_level_feature] == level, features].values.T
      for level in split_level
  ]
  feature_array = jnp.stack(array_list_by_level)
  if national_model_flag:
    feature_array = jnp.squeeze(feature_array, axis=2)
  return feature_array

def dataframe_to_jax(
    dataframe: pd.DataFrame,
    media_features: List[str],
    extra_features: List[str],
    date_feature: str,
    target: str,
    geo_feature: Optional[str] = None,
    cost_features: Optional[List[str]] = None
    ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
  """Converts pandas dataframe to right data format for media mix model.

  This function's goal is to convert dataframe which is most familar with data
  scientists to jax arrays to help the users who are not familar with array to
  use the lightweight MMM library easier.

  Args:
    dataframe: Dataframe with geo, KPI, media and non-media features.
    media_features: List of media feature names.
    extra_features: List of non media feature names.
    date_feature: Date feature name.
    target: Target variables name.
    geo_feature: Geo feature name and it is optional if the data is at national
      level.
    cost_features: List of media cost variables and it is optional if user
      use actual media cost as their media features in the model.

  Returns:
    Media, extra features, target and costs arrays.

  Raises:
    ValueError: If each geo has unequal number of weeks or there is only one
    value in the geo feature.
  """
  if geo_feature is not None:
    if dataframe[geo_feature].nunique() == 1:
      raise ValueError(
          "Geo feature has at least two geos or keep default for national model"
          )
    count_by_geo = dataframe.groupby(
        geo_feature)[date_feature].count().reset_index()
    unique_date_count = count_by_geo[date_feature].nunique()
    if unique_date_count != 1:
      raise ValueError("Not all the geos have same number of weeks.")
    national_model_flag = False
    features_to_sort = [date_feature, geo_feature]
  else:
    national_model_flag = True
    features_to_sort = [date_feature]

  df_sorted = dataframe.sort_values(by=features_to_sort)
  media_features_data = _split_array_into_list(
      dataframe=df_sorted,
      split_level_feature=date_feature,
      features=media_features,
      national_model_flag=national_model_flag)

  extra_features_data = _split_array_into_list(
      dataframe=df_sorted,
      split_level_feature=date_feature,
      features=extra_features,
      national_model_flag=national_model_flag)

  target_data = _split_array_into_list(
      dataframe=df_sorted,
      split_level_feature=date_feature,
      features=[target],
      national_model_flag=national_model_flag)
  target_data = jnp.squeeze(target_data)

  if cost_features:
    cost_data = jnp.dot(
        jnp.full(len(dataframe), 1), dataframe[cost_features].values)
  else:
    cost_data = jnp.dot(
        jnp.full(len(dataframe), 1), dataframe[media_features].values)
  return (media_features_data, extra_features_data, target_data, cost_data)# jax-ndarray

def get_halfnormal_mean_from_scale(scale: float) -> float:
  """Returns the mean of the half-normal distribition."""
  # https://en.wikipedia.org/wiki/Half-normal_distribution
  return scale * np.sqrt(2) / np.sqrt(np.pi)

def get_halfnormal_scale_from_mean(mean: float) -> float:
  """Returns the scale of the half-normal distribution."""
  # https://en.wikipedia.org/wiki/Half-normal_distribution
  return mean * np.sqrt(np.pi) / np.sqrt(2)

def get_beta_params_from_mu_sigma(mu: float,
                                  sigma: float,
                                  bracket: Tuple[float, float] = (.5, 100.)
                                  ) -> Tuple[float, float]:
  """Deterministically estimates (a, b) from (mu, sigma) of a beta variable.

  https://en.wikipedia.org/wiki/Beta_distribution

  Args:
    mu: The sample mean of the beta distributed variable.
    sigma: The sample standard deviation of the beta distributed variable.
    bracket: Search bracket for b.

  Returns:
    Tuple of the (a, b) parameters.
  """
  # Assume a = 1 to find b.
  def _f(x):
    return x ** 2 + 4 * x + 5 + 2 / x - 1 / sigma ** 2
  b = optimize.root_scalar(_f, bracket=bracket, method="brentq").root
  # Given b, now find a better a.
  a = b / (1 / mu - 1)
  return a, b

def _estimate_pdf(p: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
  """Estimates smooth pdf with Gaussian kernel.

  Args:
    p: Samples.
    x: The continuous x space (sorted).

  Returns:
    A density vector.
  """
  density = sum(stats.norm(xi).pdf(x) for xi in p)
  return density / density.sum()

def _pmf(p: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
  """Estimates discrete pmf.

  Args:
    p: Samples.
    x: The discrete x space (sorted).

  Returns:
    A pmf vector.
  """
  p_cdf = jnp.array([jnp.sum(p <= x[i]) for i in range(len(x))])
  p_pmf = np.concatenate([[p_cdf[0]], jnp.diff(p_cdf)])
  return p_pmf / p_pmf.sum()

def distance_pior_posterior(p: jnp.ndarray, q: jnp.ndarray, method: str = "KS",
                            discrete: bool = True) -> float:
  """Quantifies the distance between two distributions.

  Note we do not use KL divergence because it's not defined when a probability
  is 0.

  https://en.wikipedia.org/wiki/Hellinger_distance

  Args:
    p: Samples for distribution 1.
    q: Samples for distribution 2.
    method: We can have four methods: KS, Hellinger, JS and min.
    discrete: Whether input data is discrete or continuous.

  Returns:
    The distance metric (between 0 and 1).
  """

  if method == "KS":
    # https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.ks_2samp.html
    return stats.ks_2samp(p, q).statistic
  elif method in ["Hellinger", "JS", "min"]:
    if discrete:
      x = jnp.unique(jnp.concatenate((p, q)))
      p_pdf = _pmf(p, x)
      q_pdf = _pmf(q, x)
    else:
      minx, maxx = min(p.min(), q.min()), max(p.max(), q.max())
      x = np.linspace(minx, maxx, 100)
      p_pdf = _estimate_pdf(p, x)
      q_pdf = _estimate_pdf(q, x)
  if method == "Hellinger":
    return np.sqrt(jnp.sum((np.sqrt(p_pdf) - np.sqrt(q_pdf)) ** 2)) / np.sqrt(2)
  elif method == "JS":
    # https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.jensenshannon.html
    return spatial.distance.jensenshannon(p_pdf, q_pdf)
  else:
    return 1 - np.minimum(p_pdf, q_pdf).sum()

def interpolate_outliers(x: jnp.ndarray,
                         outlier_idx: jnp.ndarray) -> jnp.ndarray:
  """Overwrites outliers in x with interpolated values.

  Args:
    x: The original univariate variable with outliers.
    outlier_idx: Indices of the outliers in x.

  Returns:
    A cleaned x with outliers overwritten.

  """

  time_idx = jnp.arange(len(x))
  inverse_idx = jnp.array([i for i in range(len(x)) if i not in outlier_idx])
  interp_func = interpolate.interp1d(
      time_idx[inverse_idx], x[inverse_idx], kind="linear")
  x = x.at[outlier_idx].set(interp_func(time_idx[outlier_idx]))
  return x


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

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

from lightweight_mmm import models

class ModelsTest(parameterized.TestCase):

  @parameterized.named_parameters(
      dict(testcase_name="one_channel", shape=(10, 1)),
      dict(testcase_name="five_channel", shape=(10, 5)),
      dict(testcase_name="same_channels_as_rows", shape=(10, 10)),
      dict(testcase_name="geo_shape_1", shape=(10, 10, 5)),
      dict(testcase_name="geo_shape_2", shape=(10, 5, 2)),
      dict(testcase_name="one_channel_one_row", shape=(1, 1)))
  def test_transform_adstock_produces_correct_output_shape(self, shape):

    def mock_model_function(media_data):

      numpyro.deterministic(
          "transformed_media",
          models.transform_adstock(media_data, custom_priors={}))


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

"""Tests for models."""

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

from lightweight_mmm import models

class ModelsTest(parameterized.TestCase):

  @parameterized.named_parameters(
      dict(testcase_name="one_channel", shape=(10, 1)),
      dict(testcase_name="five_channel", shape=(10, 5)),
      dict(testcase_name="same_channels_as_rows", shape=(10, 10)),
      dict(testcase_name="geo_shape_1", shape=(10, 10, 5)),
      dict(testcase_name="geo_shape_2", shape=(10, 5, 2)),
      dict(testcase_name="one_channel_one_row", shape=(1, 1)))
  def test_transform_adstock_produces_correct_output_shape(self, shape):

    def mock_model_function(media_data):
      numpyro.deterministic(
          "transformed_media",
          models.transform_adstock(media_data, custom_priors={}))

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

    mcmc.run(rng_key, media_data=media)
    transformed_media = mcmc.get_samples()["transformed_media"].mean(axis=0)

    self.assertEqual(media.shape, transformed_media.shape)

  @parameterized.named_parameters(
      dict(testcase_name="one_channel", shape=(10, 1)),
      dict(testcase_name="five_channel", shape=(10, 5)),
      dict(testcase_name="same_channels_as_rows", shape=(10, 10)),
      dict(testcase_name="geo_shape_1", shape=(10, 10, 5)),
      dict(testcase_name="geo_shape_2", shape=(10, 5, 2)),
      dict(testcase_name="one_channel_one_row", shape=(1, 1)))
  def test_transform_hill_adstock_produces_correct_output_shape(self, shape):

    def mock_model_function(media_data):

      numpyro.deterministic(
          "transformed_media",
          models.transform_hill_adstock(media_data, custom_priors={}))


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

"""Tests for models."""

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

from lightweight_mmm import models

class ModelsTest(parameterized.TestCase):

  @parameterized.named_parameters(
      dict(testcase_name="one_channel", shape=(10, 1)),
      dict(testcase_name="five_channel", shape=(10, 5)),
      dict(testcase_name="same_channels_as_rows", shape=(10, 10)),
      dict(testcase_name="geo_shape_1", shape=(10, 10, 5)),
      dict(testcase_name="geo_shape_2", shape=(10, 5, 2)),
      dict(testcase_name="one_channel_one_row", shape=(1, 1)))
  def test_transform_adstock_produces_correct_output_shape(self, shape):

    def mock_model_function(media_data):
      numpyro.deterministic(
          "transformed_media",
          models.transform_adstock(media_data, custom_priors={}))

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

    mcmc.run(rng_key, media_data=media)
    transformed_media = mcmc.get_samples()["transformed_media"].mean(axis=0)

    self.assertEqual(media.shape, transformed_media.shape)

  @parameterized.named_parameters(
      dict(testcase_name="one_channel", shape=(10, 1)),
      dict(testcase_name="five_channel", shape=(10, 5)),
      dict(testcase_name="same_channels_as_rows", shape=(10, 10)),
      dict(testcase_name="geo_shape_1", shape=(10, 10, 5)),
      dict(testcase_name="geo_shape_2", shape=(10, 5, 2)),
      dict(testcase_name="one_channel_one_row", shape=(1, 1)))
  def test_transform_hill_adstock_produces_correct_output_shape(self, shape):

    def mock_model_function(media_data):
      numpyro.deterministic(
          "transformed_media",
          models.transform_hill_adstock(media_data, custom_priors={}))

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

    mcmc.run(rng_key, media_data=media)
    transformed_media = mcmc.get_samples()["transformed_media"].mean(axis=0)

    self.assertEqual(media.shape, transformed_media.shape)

  @parameterized.named_parameters(
      dict(testcase_name="one_channel", shape=(10, 1)),
      dict(testcase_name="five_channel", shape=(10, 5)),
      dict(testcase_name="same_channels_as_rows", shape=(10, 10)),
      dict(testcase_name="geo_shape_1", shape=(10, 10, 5)),
      dict(testcase_name="geo_shape_2", shape=(10, 5, 2)),
      dict(testcase_name="one_channel_one_row", shape=(1, 1)))
  def test_transform_carryover_produces_correct_output_shape(self, shape):

    def mock_model_function(media_data):

      numpyro.deterministic(
          "transformed_media",
          models.transform_carryover(media_data, custom_priors={}))


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

"""Tests for optimize_media."""
from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import jax
import jax.numpy as jnp
import numpy as np

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import optimize_media
from lightweight_mmm import preprocessing

class OptimizeMediaTest(parameterized.TestCase):

  @classmethod
  def setUpClass(cls):

    super(OptimizeMediaTest, 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)


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 optimize_media."""
from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import jax
import jax.numpy as jnp
import numpy as np

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import optimize_media
from lightweight_mmm import preprocessing

class OptimizeMediaTest(parameterized.TestCase):

  @classmethod
  def setUpClass(cls):
    super(OptimizeMediaTest, 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)

  def setUp(self):

    super().setUp()
    self.mock_minimize = self.enter_context(
        mock.patch.object(optimize_media.optimize, "minimize", 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.

"""Media transformations for accounting for lagging or media effects."""

import functools
from typing import Union

import jax
import jax.numpy as jnp

@functools.partial(jax.jit, static_argnums=[0, 1])
def calculate_seasonality(
    number_periods: int,
    degrees: int,
    gamma_seasonality: Union[int, float, jnp.ndarray],
    frequency: int = 52,
) -> jnp.ndarray:
  """Calculates cyclic variation seasonality using Fourier terms.

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

  Args:
    number_periods: Number of seasonal periods in the data. Eg. for 1 year of
      seasonal data it will be 52, for 3 years of the same kind 156.
    degrees: 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: Frequency of the seasonality being computed. By default is 52 for
      weekly data (52 weeks in a year).

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

  seasonality_range = jnp.expand_dims(a=jnp.arange(number_periods), axis=-1)
  degrees_range = jnp.arange(1, degrees+1)
  inner_value = seasonality_range * 2 * jnp.pi * degrees_range / 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 (season_matrix * gamma_seasonality).sum(axis=2).sum(axis=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.

"""Media transformations for accounting for lagging or media effects."""

import functools
from typing import Union

import jax
import jax.numpy as jnp

@functools.partial(jax.jit, static_argnums=[0, 1])
def calculate_seasonality(
    number_periods: int,
    degrees: int,
    gamma_seasonality: Union[int, float, jnp.ndarray],
    frequency: int = 52,
) -> jnp.ndarray:
  """Calculates cyclic variation seasonality using Fourier terms.

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

  Args:
    number_periods: Number of seasonal periods in the data. Eg. for 1 year of
      seasonal data it will be 52, for 3 years of the same kind 156.
    degrees: 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: Frequency of the seasonality being computed. By default is 52 for
      weekly data (52 weeks in a year).

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

  seasonality_range = jnp.expand_dims(a=jnp.arange(number_periods), axis=-1)
  degrees_range = jnp.arange(1, degrees+1)
  inner_value = seasonality_range * 2 * jnp.pi * degrees_range / 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 (season_matrix * gamma_seasonality).sum(axis=2).sum(axis=1)

@jax.jit
def adstock(data: jnp.ndarray,
            lag_weight: float = .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: float = 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.

"""Media transformations for accounting for lagging or media effects."""

import functools
from typing import Union

import jax
import jax.numpy as jnp

@functools.partial(jax.jit, static_argnums=[0, 1])
def calculate_seasonality(
    number_periods: int,
    degrees: int,
    gamma_seasonality: Union[int, float, jnp.ndarray],
    frequency: int = 52,
) -> jnp.ndarray:
  """Calculates cyclic variation seasonality using Fourier terms.

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

  Args:
    number_periods: Number of seasonal periods in the data. Eg. for 1 year of
      seasonal data it will be 52, for 3 years of the same kind 156.
    degrees: 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: Frequency of the seasonality being computed. By default is 52 for
      weekly data (52 weeks in a year).

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

  seasonality_range = jnp.expand_dims(a=jnp.arange(number_periods), axis=-1)
  degrees_range = jnp.arange(1, degrees+1)
  inner_value = seasonality_range * 2 * jnp.pi * degrees_range / 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 (season_matrix * gamma_seasonality).sum(axis=2).sum(axis=1)

@jax.jit
def adstock(data: jnp.ndarray,
            lag_weight: float = .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: float = 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)

@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 = 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.

"""Media transformations for accounting for lagging or media effects."""

import functools
from typing import Union

import jax
import jax.numpy as jnp

@functools.partial(jax.jit, static_argnums=[0, 1])
def calculate_seasonality(
    number_periods: int,
    degrees: int,
    gamma_seasonality: Union[int, float, jnp.ndarray],
    frequency: int = 52,
) -> jnp.ndarray:
  """Calculates cyclic variation seasonality using Fourier terms.

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

  Args:
    number_periods: Number of seasonal periods in the data. Eg. for 1 year of
      seasonal data it will be 52, for 3 years of the same kind 156.
    degrees: 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: Frequency of the seasonality being computed. By default is 52 for
      weekly data (52 weeks in a year).

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

  seasonality_range = jnp.expand_dims(a=jnp.arange(number_periods), axis=-1)
  degrees_range = jnp.arange(1, degrees+1)
  inner_value = seasonality_range * 2 * jnp.pi * degrees_range / 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 (season_matrix * gamma_seasonality).sum(axis=2).sum(axis=1)

@jax.jit
def adstock(data: jnp.ndarray,
            lag_weight: float = .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: float = 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)

@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 = apply_exponent_safe(
      data=data / half_max_effective_concentration, exponent=-slope)
  return jnp.where(save_transform == 0, x=0, y=1. / (1 + save_transform))

@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 = 13) -> 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.

"""Tests for lightweight_mmm."""

import copy

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

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import models

class LightweightMmmTest(parameterized.TestCase):

  @classmethod
  def setUpClass(cls):

    super(LightweightMmmTest, 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,
        extra_features=jnp.ones((50, 2)),
        number_warmup=2,
        number_samples=4,
        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,
        extra_features=jnp.ones((50, 2, 3)),
        number_warmup=2,
        number_samples=4,
        number_chains=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.

"""A simple and lightweight library for Media Mix Modelling.

Simple usage of this class goes as following:

```
mmm = lightweight_mmm.LightweightMMM()
mmm.fit(media=media_data,
        extra_features=extra_features,
        media_prior=costs,
        target=target,
        number_samples=1000,
        number_chains=2)

# For obtaining media contribution percentage and ROI
predictions, media_contribution_hat_pct, roi_hat = mmm.get_posterior_metrics()

# For running predictions on unseen data
mmm.predict(media=media_data_test, extra_features=extra_features_test)
```
"""

import collections
import dataclasses
import functools
import itertools
import logging
import numbers
from typing import Any, Callable, Dict, Mapping, MutableMapping, Optional, Sequence, Tuple, Union

from absl import logging
import immutabledict
import jax
import jax.numpy as jnp
import numpy as np
import numpyro
from numpyro import distributions as dist
from numpyro import infer

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

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

_NAMES_TO_MODEL_TRANSFORMS = immutabledict.immutabledict({
    "hill_adstock": models.transform_hill_adstock,
    "adstock": models.transform_adstock,
    "carryover": models.transform_carryover
})
_MODEL_FUNCTION = models.media_mix_model

def _compare_equality_for_lmmm(item_1: Any, item_2: Any) -> bool:
  """Compares two items for equality.

  Helper function for the __eq__ method of LightweightmMM. First checks if items
  are strings or lists of strings (it's okay if empty lists compare True), then
  uses jnp.array_equal if the items are jax.numpy.DeviceArray or other related
  sequences, and uses items' __eq__ otherwise.

  Note: this implementation does not cover every possible data structure, but
  it does cover all the data structures seen in attributes used by
  LightweightMMM. Sometimes the DeviceArray is hidden in the value of a
  MutableMapping, hence the recursion.

  Args:
    item_1: First item to be compared.
    item_2: Second item to be compared.

  Returns:
    Boolean for whether item_1 equals item_2.
  """

  # This is pretty strict but LMMM classes don't need to compare equal unless
  # they are exact copies.

  if type(item_1) != type(item_2):
    is_equal = False
  elif isinstance(item_1, str):
    is_equal = item_1 == item_2
  elif isinstance(item_1, (jax.Array, np.ndarray, Sequence)):
    if all(isinstance(x, str) for x in item_1) and all(
        isinstance(x, str) for x in item_2):
      is_equal = item_1 == item_2
    else:
      is_equal = np.array_equal(item_1, item_2, equal_nan=True)
  elif isinstance(item_1, MutableMapping):
    is_equal = all([
        _compare_equality_for_lmmm(item_1[x], item_2[x])
        for x in item_1.keys() | item_2.keys()
    ])
  else:
    is_equal = item_1 == item_2

  return is_equal


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.

"""A simple and lightweight library for Media Mix Modelling.

Simple usage of this class goes as following:

```
mmm = lightweight_mmm.LightweightMMM()
mmm.fit(media=media_data,
        extra_features=extra_features,
        media_prior=costs,
        target=target,
        number_samples=1000,
        number_chains=2)

# For obtaining media contribution percentage and ROI
predictions, media_contribution_hat_pct, roi_hat = mmm.get_posterior_metrics()

# For running predictions on unseen data
mmm.predict(media=media_data_test, extra_features=extra_features_test)
```
"""

import collections
import dataclasses
import functools
import itertools
import logging
import numbers
from typing import Any, Callable, Dict, Mapping, MutableMapping, Optional, Sequence, Tuple, Union

from absl import logging
import immutabledict
import jax
import jax.numpy as jnp
import numpy as np
import numpyro
from numpyro import distributions as dist
from numpyro import infer

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

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

_NAMES_TO_MODEL_TRANSFORMS = immutabledict.immutabledict({
    "hill_adstock": models.transform_hill_adstock,
    "adstock": models.transform_adstock,
    "carryover": models.transform_carryover
})
_MODEL_FUNCTION = models.media_mix_model

def _compare_equality_for_lmmm(item_1: Any, item_2: Any) -> bool:
  """Compares two items for equality.

  Helper function for the __eq__ method of LightweightmMM. First checks if items
  are strings or lists of strings (it's okay if empty lists compare True), then
  uses jnp.array_equal if the items are jax.numpy.DeviceArray or other related
  sequences, and uses items' __eq__ otherwise.

  Note: this implementation does not cover every possible data structure, but
  it does cover all the data structures seen in attributes used by
  LightweightMMM. Sometimes the DeviceArray is hidden in the value of a
  MutableMapping, hence the recursion.

  Args:
    item_1: First item to be compared.
    item_2: Second item to be compared.

  Returns:
    Boolean for whether item_1 equals item_2.
  """

  # This is pretty strict but LMMM classes don't need to compare equal unless
  # they are exact copies.
  if type(item_1) != type(item_2):
    is_equal = False
  elif isinstance(item_1, str):
    is_equal = item_1 == item_2
  elif isinstance(item_1, (jax.Array, np.ndarray, Sequence)):
    if all(isinstance(x, str) for x in item_1) and all(
        isinstance(x, str) for x in item_2):
      is_equal = item_1 == item_2
    else:
      is_equal = np.array_equal(item_1, item_2, equal_nan=True)
  elif isinstance(item_1, MutableMapping):
    is_equal = all([
        _compare_equality_for_lmmm(item_1[x], item_2[x])
        for x in item_1.keys() | item_2.keys()
    ])
  else:
    is_equal = item_1 == item_2

  return is_equal

class NotFittedModelError(Exception):
  pass

@dataclasses.dataclass(unsafe_hash=True, eq=False)
class LightweightMMM:
  """Lightweight Media Mix Modelling wrapper for bayesian models.

  The currently available models are the following:
   - hill_adstock
   - adstock
   - carryover

  It also offers the necessary utilities for calculating media contribution and
  media ROI based on models' results.

  Attributes:
    trace: Sampling trace of the bayesian model once fitted.
    n_media_channels: Number of media channels the model was trained with.
    n_geos: Number of geos for geo models or 1 for national models.
    model_name: Name of the model.
    media: The media data the model is trained on. Usefull for a variety of
      insights post model fitting.
    media_names: Names of the media channels passed at fitting time.
    custom_priors: The set of custom priors the model was trained with. An empty
      dictionary if none were passed.
  """
  model_name: str = "hill_adstock"
  n_media_channels: int = dataclasses.field(init=False, repr=False)
  n_geos: int = dataclasses.field(init=False, repr=False)
  media: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  media_names: Sequence[str] = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  trace: Dict[str, jnp.DeviceArray] = dataclasses.field(
      init=False, repr=False, hash=False, compare=False)
  custom_priors: MutableMapping[str, Prior] = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _degrees_seasonality: int = dataclasses.field(init=False, repr=False)
  _weekday_seasonality: bool = dataclasses.field(init=False, repr=False)
  _media_prior: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _extra_features: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _target: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _train_media_size: int = dataclasses.field(
      init=False, repr=False, hash=True, compare=False)
  _mcmc: numpyro.infer.MCMC = dataclasses.field(
      init=False, repr=False, hash=False, compare=False)

  def __post_init__(self):

    if self.model_name not in _NAMES_TO_MODEL_TRANSFORMS:
      raise ValueError("Model name passed not valid. Please use any of the"
                       "following: 'hill_adstock', 'adstock', 'carryover'.")
    self._model_function = _MODEL_FUNCTION
    self._model_transform_function = _NAMES_TO_MODEL_TRANSFORMS[self.model_name]
    self._prior_names = models.MODEL_PRIORS_NAMES.union(
        models.TRANSFORM_PRIORS_NAMES[self.model_name])


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.

"""A simple and lightweight library for Media Mix Modelling.

Simple usage of this class goes as following:

```
mmm = lightweight_mmm.LightweightMMM()
mmm.fit(media=media_data,
        extra_features=extra_features,
        media_prior=costs,
        target=target,
        number_samples=1000,
        number_chains=2)

# For obtaining media contribution percentage and ROI
predictions, media_contribution_hat_pct, roi_hat = mmm.get_posterior_metrics()

# For running predictions on unseen data
mmm.predict(media=media_data_test, extra_features=extra_features_test)
```
"""

import collections
import dataclasses
import functools
import itertools
import logging
import numbers
from typing import Any, Callable, Dict, Mapping, MutableMapping, Optional, Sequence, Tuple, Union

from absl import logging
import immutabledict
import jax
import jax.numpy as jnp
import numpy as np
import numpyro
from numpyro import distributions as dist
from numpyro import infer

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

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

_NAMES_TO_MODEL_TRANSFORMS = immutabledict.immutabledict({
    "hill_adstock": models.transform_hill_adstock,
    "adstock": models.transform_adstock,
    "carryover": models.transform_carryover
})
_MODEL_FUNCTION = models.media_mix_model

def _compare_equality_for_lmmm(item_1: Any, item_2: Any) -> bool:
  """Compares two items for equality.

  Helper function for the __eq__ method of LightweightmMM. First checks if items
  are strings or lists of strings (it's okay if empty lists compare True), then
  uses jnp.array_equal if the items are jax.numpy.DeviceArray or other related
  sequences, and uses items' __eq__ otherwise.

  Note: this implementation does not cover every possible data structure, but
  it does cover all the data structures seen in attributes used by
  LightweightMMM. Sometimes the DeviceArray is hidden in the value of a
  MutableMapping, hence the recursion.

  Args:
    item_1: First item to be compared.
    item_2: Second item to be compared.

  Returns:
    Boolean for whether item_1 equals item_2.
  """

  # This is pretty strict but LMMM classes don't need to compare equal unless
  # they are exact copies.
  if type(item_1) != type(item_2):
    is_equal = False
  elif isinstance(item_1, str):
    is_equal = item_1 == item_2
  elif isinstance(item_1, (jax.Array, np.ndarray, Sequence)):
    if all(isinstance(x, str) for x in item_1) and all(
        isinstance(x, str) for x in item_2):
      is_equal = item_1 == item_2
    else:
      is_equal = np.array_equal(item_1, item_2, equal_nan=True)
  elif isinstance(item_1, MutableMapping):
    is_equal = all([
        _compare_equality_for_lmmm(item_1[x], item_2[x])
        for x in item_1.keys() | item_2.keys()
    ])
  else:
    is_equal = item_1 == item_2

  return is_equal

class NotFittedModelError(Exception):
  pass

@dataclasses.dataclass(unsafe_hash=True, eq=False)
class LightweightMMM:
  """Lightweight Media Mix Modelling wrapper for bayesian models.

  The currently available models are the following:
   - hill_adstock
   - adstock
   - carryover

  It also offers the necessary utilities for calculating media contribution and
  media ROI based on models' results.

  Attributes:
    trace: Sampling trace of the bayesian model once fitted.
    n_media_channels: Number of media channels the model was trained with.
    n_geos: Number of geos for geo models or 1 for national models.
    model_name: Name of the model.
    media: The media data the model is trained on. Usefull for a variety of
      insights post model fitting.
    media_names: Names of the media channels passed at fitting time.
    custom_priors: The set of custom priors the model was trained with. An empty
      dictionary if none were passed.
  """
  model_name: str = "hill_adstock"
  n_media_channels: int = dataclasses.field(init=False, repr=False)
  n_geos: int = dataclasses.field(init=False, repr=False)
  media: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  media_names: Sequence[str] = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  trace: Dict[str, jnp.DeviceArray] = dataclasses.field(
      init=False, repr=False, hash=False, compare=False)
  custom_priors: MutableMapping[str, Prior] = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _degrees_seasonality: int = dataclasses.field(init=False, repr=False)
  _weekday_seasonality: bool = dataclasses.field(init=False, repr=False)
  _media_prior: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _extra_features: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _target: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _train_media_size: int = dataclasses.field(
      init=False, repr=False, hash=True, compare=False)
  _mcmc: numpyro.infer.MCMC = dataclasses.field(
      init=False, repr=False, hash=False, compare=False)

  def __post_init__(self):
    if self.model_name not in _NAMES_TO_MODEL_TRANSFORMS:
      raise ValueError("Model name passed not valid. Please use any of the"
                       "following: 'hill_adstock', 'adstock', 'carryover'.")
    self._model_function = _MODEL_FUNCTION
    self._model_transform_function = _NAMES_TO_MODEL_TRANSFORMS[self.model_name]
    self._prior_names = models.MODEL_PRIORS_NAMES.union(
        models.TRANSFORM_PRIORS_NAMES[self.model_name])

  def __eq__(self, other: Any) -> bool:
    """Equality method for LightweightMMMM.

    We need a special method here to handle a couple of issues. First, some of
    the attributes for LightweightMMM are arrays, which contain multiple values
    and cannot be evaluated with the default __eq__ method. Second, some
    attributes are initially undefined and only get values after fitting a
    model. The latter is dealt with within this function, and the former within
    the helper function _compare_equality_for_lmmm().

    Args:
      other: Dataclass to compare against.

    Returns:
      Boolean for whether self == other; NotImplemented if other is not a
      LightweightMMM.
    """

    if not isinstance(other, LightweightMMM):
      return NotImplemented

    def _create_list_of_attributes_to_compare(
        mmm_instance: Any) -> Sequence[str]:
      all_attributes_that_can_be_compared = sorted(
          [x.name for x in dataclasses.fields(mmm_instance) if x.compare])
      attributes_which_have_been_instantiated = [
          x for x in all_attributes_that_can_be_compared
          if hasattr(mmm_instance, x)
      ]
      return attributes_which_have_been_instantiated

    self_attributes = _create_list_of_attributes_to_compare(self)
    other_attributes = _create_list_of_attributes_to_compare(other)

    return all(
        _compare_equality_for_lmmm(getattr(self, a1), getattr(other, a2))
        for a1, a2 in itertools.zip_longest(self_attributes, other_attributes))


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.

"""A simple and lightweight library for Media Mix Modelling.

Simple usage of this class goes as following:

```
mmm = lightweight_mmm.LightweightMMM()
mmm.fit(media=media_data,
        extra_features=extra_features,
        media_prior=costs,
        target=target,
        number_samples=1000,
        number_chains=2)

# For obtaining media contribution percentage and ROI
predictions, media_contribution_hat_pct, roi_hat = mmm.get_posterior_metrics()

# For running predictions on unseen data
mmm.predict(media=media_data_test, extra_features=extra_features_test)
```
"""

import collections
import dataclasses
import functools
import itertools
import logging
import numbers
from typing import Any, Callable, Dict, Mapping, MutableMapping, Optional, Sequence, Tuple, Union

from absl import logging
import immutabledict
import jax
import jax.numpy as jnp
import numpy as np
import numpyro
from numpyro import distributions as dist
from numpyro import infer

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

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

_NAMES_TO_MODEL_TRANSFORMS = immutabledict.immutabledict({
    "hill_adstock": models.transform_hill_adstock,
    "adstock": models.transform_adstock,
    "carryover": models.transform_carryover
})
_MODEL_FUNCTION = models.media_mix_model

def _compare_equality_for_lmmm(item_1: Any, item_2: Any) -> bool:
  """Compares two items for equality.

  Helper function for the __eq__ method of LightweightmMM. First checks if items
  are strings or lists of strings (it's okay if empty lists compare True), then
  uses jnp.array_equal if the items are jax.numpy.DeviceArray or other related
  sequences, and uses items' __eq__ otherwise.

  Note: this implementation does not cover every possible data structure, but
  it does cover all the data structures seen in attributes used by
  LightweightMMM. Sometimes the DeviceArray is hidden in the value of a
  MutableMapping, hence the recursion.

  Args:
    item_1: First item to be compared.
    item_2: Second item to be compared.

  Returns:
    Boolean for whether item_1 equals item_2.
  """

  # This is pretty strict but LMMM classes don't need to compare equal unless
  # they are exact copies.
  if type(item_1) != type(item_2):
    is_equal = False
  elif isinstance(item_1, str):
    is_equal = item_1 == item_2
  elif isinstance(item_1, (jax.Array, np.ndarray, Sequence)):
    if all(isinstance(x, str) for x in item_1) and all(
        isinstance(x, str) for x in item_2):
      is_equal = item_1 == item_2
    else:
      is_equal = np.array_equal(item_1, item_2, equal_nan=True)
  elif isinstance(item_1, MutableMapping):
    is_equal = all([
        _compare_equality_for_lmmm(item_1[x], item_2[x])
        for x in item_1.keys() | item_2.keys()
    ])
  else:
    is_equal = item_1 == item_2

  return is_equal

class NotFittedModelError(Exception):
  pass

@dataclasses.dataclass(unsafe_hash=True, eq=False)
class LightweightMMM:
  """Lightweight Media Mix Modelling wrapper for bayesian models.

  The currently available models are the following:
   - hill_adstock
   - adstock
   - carryover

  It also offers the necessary utilities for calculating media contribution and
  media ROI based on models' results.

  Attributes:
    trace: Sampling trace of the bayesian model once fitted.
    n_media_channels: Number of media channels the model was trained with.
    n_geos: Number of geos for geo models or 1 for national models.
    model_name: Name of the model.
    media: The media data the model is trained on. Usefull for a variety of
      insights post model fitting.
    media_names: Names of the media channels passed at fitting time.
    custom_priors: The set of custom priors the model was trained with. An empty
      dictionary if none were passed.
  """
  model_name: str = "hill_adstock"
  n_media_channels: int = dataclasses.field(init=False, repr=False)
  n_geos: int = dataclasses.field(init=False, repr=False)
  media: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  media_names: Sequence[str] = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  trace: Dict[str, jnp.DeviceArray] = dataclasses.field(
      init=False, repr=False, hash=False, compare=False)
  custom_priors: MutableMapping[str, Prior] = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _degrees_seasonality: int = dataclasses.field(init=False, repr=False)
  _weekday_seasonality: bool = dataclasses.field(init=False, repr=False)
  _media_prior: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _extra_features: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _target: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _train_media_size: int = dataclasses.field(
      init=False, repr=False, hash=True, compare=False)
  _mcmc: numpyro.infer.MCMC = dataclasses.field(
      init=False, repr=False, hash=False, compare=False)

  def __post_init__(self):
    if self.model_name not in _NAMES_TO_MODEL_TRANSFORMS:
      raise ValueError("Model name passed not valid. Please use any of the"
                       "following: 'hill_adstock', 'adstock', 'carryover'.")
    self._model_function = _MODEL_FUNCTION
    self._model_transform_function = _NAMES_TO_MODEL_TRANSFORMS[self.model_name]
    self._prior_names = models.MODEL_PRIORS_NAMES.union(
        models.TRANSFORM_PRIORS_NAMES[self.model_name])

  def __eq__(self, other: Any) -> bool:
    """Equality method for LightweightMMMM.

    We need a special method here to handle a couple of issues. First, some of
    the attributes for LightweightMMM are arrays, which contain multiple values
    and cannot be evaluated with the default __eq__ method. Second, some
    attributes are initially undefined and only get values after fitting a
    model. The latter is dealt with within this function, and the former within
    the helper function _compare_equality_for_lmmm().

    Args:
      other: Dataclass to compare against.

    Returns:
      Boolean for whether self == other; NotImplemented if other is not a
      LightweightMMM.
    """
    if not isinstance(other, LightweightMMM):
      return NotImplemented

    def _create_list_of_attributes_to_compare(
        mmm_instance: Any) -> Sequence[str]:

      all_attributes_that_can_be_compared = sorted(
          [x.name for x in dataclasses.fields(mmm_instance) if x.compare])
      attributes_which_have_been_instantiated = [
          x for x in all_attributes_that_can_be_compared
          if hasattr(mmm_instance, x)
      ]
      return attributes_which_have_been_instantiated


Prompt:
 trace: Sampling trace of the bayesian model once fitted.
    n_media_channels: Number of media channels the model was trained with.
    n_geos: Number of geos for geo models or 1 for national models.
    model_name: Name of the model.
    media: The media data the model is trained on. Usefull for a variety of
      insights post model fitting.
    media_names: Names of the media channels passed at fitting time.
    custom_priors: The set of custom priors the model was trained with. An empty
      dictionary if none were passed.
  """
  model_name: str = "hill_adstock"
  n_media_channels: int = dataclasses.field(init=False, repr=False)
  n_geos: int = dataclasses.field(init=False, repr=False)
  media: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  media_names: Sequence[str] = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  trace: Dict[str, jnp.DeviceArray] = dataclasses.field(
      init=False, repr=False, hash=False, compare=False)
  custom_priors: MutableMapping[str, Prior] = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _degrees_seasonality: int = dataclasses.field(init=False, repr=False)
  _weekday_seasonality: bool = dataclasses.field(init=False, repr=False)
  _media_prior: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _extra_features: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _target: jnp.DeviceArray = dataclasses.field(
      init=False, repr=False, hash=False, compare=True)
  _train_media_size: int = dataclasses.field(
      init=False, repr=False, hash=True, compare=False)
  _mcmc: numpyro.infer.MCMC = dataclasses.field(
      init=False, repr=False, hash=False, compare=False)

  def __post_init__(self):
    if self.model_name not in _NAMES_TO_MODEL_TRANSFORMS:
      raise ValueError("Model name passed not valid. Please use any of the"
                       "following: 'hill_adstock', 'adstock', 'carryover'.")
    self._model_function = _MODEL_FUNCTION
    self._model_transform_function = _NAMES_TO_MODEL_TRANSFORMS[self.model_name]
    self._prior_names = models.MODEL_PRIORS_NAMES.union(
        models.TRANSFORM_PRIORS_NAMES[self.model_name])

  def __eq__(self, other: Any) -> bool:
    """Equality method for LightweightMMMM.

    We need a special method here to handle a couple of issues. First, some of
    the attributes for LightweightMMM are arrays, which contain multiple values
    and cannot be evaluated with the default __eq__ method. Second, some
    attributes are initially undefined and only get values after fitting a
    model. The latter is dealt with within this function, and the former within
    the helper function _compare_equality_for_lmmm().

    Args:
      other: Dataclass to compare against.

    Returns:
      Boolean for whether self == other; NotImplemented if other is not a
      LightweightMMM.
    """
    if not isinstance(other, LightweightMMM):
      return NotImplemented

    def _create_list_of_attributes_to_compare(
        mmm_instance: Any) -> Sequence[str]:
      all_attributes_that_can_be_compared = sorted(
          [x.name for x in dataclasses.fields(mmm_instance) if x.compare])
      attributes_which_have_been_instantiated = [
          x for x in all_attributes_that_can_be_compared
          if hasattr(mmm_instance, x)
      ]
      return attributes_which_have_been_instantiated

    self_attributes = _create_list_of_attributes_to_compare(self)
    other_attributes = _create_list_of_attributes_to_compare(other)

    return all(
        _compare_equality_for_lmmm(getattr(self, a1), getattr(other, a2))
        for a1, a2 in itertools.zip_longest(self_attributes, other_attributes))

  def _preprocess_custom_priors(
      self,
      custom_priors: Dict[str, Prior]) -> MutableMapping[str, Prior]:
    """Preprocesses the user input custom priors to Numpyro distributions.

    If numpyro distributions are given they remains untouched, however if any
    other option is passed, it is passed to the default distribution to alter
    its constructor values.

    Args:
      custom_priors: Mapping of the name of the prior to its custom value.

    Returns:
      A mapping of names to numpyro distributions based on user input and
        default values.
    """
    default_priors = {
        **models._get_default_priors(),
        **models._get_transform_default_priors()[self.model_name]
    }
    # Checking that the key is contained in custom_priors has already been done
    # at this point in the fit function.
    for prior_name in custom_priors:
      if isinstance(custom_priors[prior_name], numbers.Number):
        custom_priors[prior_name] = default_priors[prior_name].__class__(
            custom_priors[prior_name])
      elif (isinstance(custom_priors[prior_name], collections.abc.Sequence) and
            not isinstance(custom_priors[prior_name], str)):
        custom_priors[prior_name] = default_priors[prior_name].__class__(
            *custom_priors[prior_name])
      elif isinstance(custom_priors[prior_name], dict):
        custom_priors[prior_name] = default_priors[prior_name].__class__(
            **custom_priors[prior_name])
      elif not isinstance(custom_priors[prior_name], dist.Distribution):
        raise ValueError(
            "Priors given must be a Numpyro distribution or one of the "
            "following to fit in the constructor of our default Numpyro "
            "distribution. It could be given as args or kwargs as long as it "
            "is the correct format for such object. Please refer to our "
            "documentation on custom priors to know more.")
    return custom_priors

  def fit(
      self,
      media: jnp.ndarray,
      media_prior: jnp.ndarray,
      target: jnp.ndarray,
      extra_features: Optional[jnp.ndarray] = None,
      degrees_seasonality: int = 2,
      seasonality_frequency: int = 52,
      weekday_seasonality: bool = False,
      media_names: Optional[Sequence[str]] = None,
      number_warmup: int = 1000,
      number_samples: int = 1000,
      number_chains: int = 2,
      target_accept_prob: float = .85,
      init_strategy: Callable[[Mapping[Any, Any], Any],
                              jnp.ndarray] = numpyro.infer.init_to_median,
      custom_priors: Optional[Dict[str, Prior]] = None,
      seed: Optional[int] = None) -> None:
    """Fits MMM given the media data, extra features, costs and sales/KPI.

    For detailed information on the selected model please refer to its
    respective function in the models.py file.

    Args:
      media: Media input data. Media data must have either 2 dims for national
        model or 3 for geo models.
      media_prior: Costs of each media channel. The number of cost values must
        be equal to the number of media channels.
      target: Target KPI to use, like for example sales.
      extra_features: Other variables to add to the model.
      degrees_seasonality: Number of degrees to use for seasonality. Default is
        2.
      seasonality_frequency: Frequency of the time period used. Default is 52 as
        in 52 weeks per year.
      weekday_seasonality: In case of daily data, also estimate seven weekday
        parameters.
      media_names: Names of the media channels passed.
      number_warmup: Number of warm up samples. Default is 1000.
      number_samples: Number of samples during sampling. Default is 1000.
      number_chains: Number of chains to sample. Default is 2.
      target_accept_prob: Target acceptance probability for step size in the
        NUTS sampler. Default is .85.
      init_strategy: Initialization function for numpyro NUTS. The available
        options can be found in
        https://num.pyro.ai/en/stable/utilities.html#initialization-strategies.
        Default is numpyro.infer.init_to_median.
      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.
      seed: Seed to use for PRNGKey during training. For better replicability
        run all different trainings with the same seed.
    """
    if media.ndim not in (2, 3):
      raise ValueError(
          "Media data must have either 2 dims for national model or 3 for geo "
          "models.")
    if media.ndim == 3 and media_prior.ndim == 1:
      media_prior = jnp.expand_dims(media_prior, axis=-1)

    if media.shape[1] != len(media_prior):
      raise ValueError("The number of data channels provided must match the "
                       "number of cost values.")
    if media.min() < 0:
      raise ValueError("Media values must be greater or equal to zero.")

    if custom_priors:
      not_used_custom_priors = set(custom_priors.keys()).difference(
          self._prior_names)
      if not_used_custom_priors:
        raise ValueError(
            "The following passed custom priors dont have a match in the model."
            " Please double check the names have been written correctly: %s" %
            not_used_custom_priors)
      custom_priors = self._preprocess_custom_priors(
          custom_priors=custom_priors)
      geo_custom_priors = set(custom_priors.keys()).intersection(
          models.GEO_ONLY_PRIORS)
      if media.ndim == 2 and geo_custom_priors:
        raise ValueError(
            "The given data is for national models but custom_prior contains "
            "priors for the geo version of the model. Please either remove geo "
            "priors for national model or pass media data with geo dimension.")
    else:
      custom_priors = {}

    if weekday_seasonality and seasonality_frequency == 52:
      logging.warn("You have chosen daily seasonality and frequency 52 "
                   "(weekly), please check you made the right seasonality "
                   "choices.")

    if extra_features is not None:
      extra_features = jnp.array(extra_features)

    if seed is None:
      seed = utils.get_time_seed()

    train_media_size = media.shape[0]
    kernel = numpyro.infer.NUTS(
        model=self._model_function,
        target_accept_prob=target_accept_prob,
        init_strategy=init_strategy)

    mcmc = numpyro.infer.MCMC(
        sampler=kernel,
        num_warmup=number_warmup,
        num_samples=number_samples,
        num_chains=number_chains)
    mcmc.run(
        rng_key=jax.random.PRNGKey(seed),
        media_data=jnp.array(media),
        extra_features=extra_features,
        target_data=jnp.array(target),
        media_prior=jnp.array(media_prior),
        degrees_seasonality=degrees_seasonality,
        frequency=seasonality_frequency,
        transform_function=self._model_transform_function,
        weekday_seasonality=weekday_seasonality,
        custom_priors=custom_priors)

    self.custom_priors = custom_priors
    if media_names is not None:
      self.media_names = media_names
    else:
      self.media_names = [f"channel_{i}" for i in range(media.shape[1])]
    self.n_media_channels = media.shape[1]
    self.n_geos = media.shape[2] if media.ndim == 3 else 1
    self._media_prior = media_prior
    self.trace = mcmc.get_samples()
    self._number_warmup = number_warmup
    self._number_samples = number_samples
    self._number_chains = number_chains
    self._target = target
    self._train_media_size = train_media_size
    self._degrees_seasonality = degrees_seasonality
    self._seasonality_frequency = seasonality_frequency
    self._weekday_seasonality = weekday_seasonality
    self.media = media
    self._extra_features = extra_features# jax-devicearray
    self._mcmc = mcmc
    logging.info("Model has been fitted")

  def print_summary(self) -> None:
    """Calls print_summary function from numpyro to print parameters summary.
    """
    # TODO(): add name selection for print.
    self._mcmc.print_summary()

  @functools.partial(
      jax.jit,
      static_argnums=(0,),
      static_argnames=("degrees_seasonality", "weekday_seasonality",
                       "transform_function", "model"))
  def _predict(
      self,
      rng_key: jnp.ndarray,
      media_data: jnp.ndarray,
      extra_features: Optional[jnp.ndarray],
      media_prior: jnp.ndarray,
      degrees_seasonality: int, frequency: int,
      transform_function: Callable[[Any], jnp.ndarray],
      weekday_seasonality: bool,
      model: Callable[[Any], None],
      posterior_samples: Dict[str, jnp.ndarray],
      custom_priors: Dict[str, Prior]
      ) -> Dict[str, jnp.ndarray]:
    """Encapsulates the numpyro.infer.Predictive function for predict method.

    It serves as a helper jitted function for running predictions.

    Args:
      rng_key: A jax.random.PRNGKey.
      media_data: Media array for needed for the model to run predictions.
      extra_features: Extra features for needed for the model to run.
      media_prior: Cost prior used for training the model.
      degrees_seasonality: Number of degrees for the seasonality.
      frequency: Frequency of the seasonality.
      transform_function: Media transform function to use within the model.
      weekday_seasonality: Allow daily weekday estimation.
      model: Numpyro model to use for numpyro.infer.Predictive.
      posterior_samples: Mapping of the posterior samples.
      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 predictions for the given data.
    """

    return infer.Predictive(
        model=model, posterior_samples=posterior_samples)(
            rng_key=rng_key,
            media_data=media_data,
            extra_features=extra_features,
            media_prior=media_prior,
            target_data=None,
            degrees_seasonality=degrees_seasonality,
            frequency=frequency,
            transform_function=transform_function,
            custom_priors=custom_priors,
            weekday_seasonality=weekday_seasonality)


Prompt:
  def _preprocess_custom_priors(
      self,
      custom_priors: Dict[str, Prior]) -> MutableMapping[str, Prior]:
    """Preprocesses the user input custom priors to Numpyro distributions.

    If numpyro distributions are given they remains untouched, however if any
    other option is passed, it is passed to the default distribution to alter
    its constructor values.

    Args:
      custom_priors: Mapping of the name of the prior to its custom value.

    Returns:
      A mapping of names to numpyro distributions based on user input and
        default values.
    """
    default_priors = {
        **models._get_default_priors(),
        **models._get_transform_default_priors()[self.model_name]
    }
    # Checking that the key is contained in custom_priors has already been done
    # at this point in the fit function.
    for prior_name in custom_priors:
      if isinstance(custom_priors[prior_name], numbers.Number):
        custom_priors[prior_name] = default_priors[prior_name].__class__(
            custom_priors[prior_name])
      elif (isinstance(custom_priors[prior_name], collections.abc.Sequence) and
            not isinstance(custom_priors[prior_name], str)):
        custom_priors[prior_name] = default_priors[prior_name].__class__(
            *custom_priors[prior_name])
      elif isinstance(custom_priors[prior_name], dict):
        custom_priors[prior_name] = default_priors[prior_name].__class__(
            **custom_priors[prior_name])
      elif not isinstance(custom_priors[prior_name], dist.Distribution):
        raise ValueError(
            "Priors given must be a Numpyro distribution or one of the "
            "following to fit in the constructor of our default Numpyro "
            "distribution. It could be given as args or kwargs as long as it "
            "is the correct format for such object. Please refer to our "
            "documentation on custom priors to know more.")
    return custom_priors

  def fit(
      self,
      media: jnp.ndarray,
      media_prior: jnp.ndarray,
      target: jnp.ndarray,
      extra_features: Optional[jnp.ndarray] = None,
      degrees_seasonality: int = 2,
      seasonality_frequency: int = 52,
      weekday_seasonality: bool = False,
      media_names: Optional[Sequence[str]] = None,
      number_warmup: int = 1000,
      number_samples: int = 1000,
      number_chains: int = 2,
      target_accept_prob: float = .85,
      init_strategy: Callable[[Mapping[Any, Any], Any],
                              jnp.ndarray] = numpyro.infer.init_to_median,
      custom_priors: Optional[Dict[str, Prior]] = None,
      seed: Optional[int] = None) -> None:
    """Fits MMM given the media data, extra features, costs and sales/KPI.

    For detailed information on the selected model please refer to its
    respective function in the models.py file.

    Args:
      media: Media input data. Media data must have either 2 dims for national
        model or 3 for geo models.
      media_prior: Costs of each media channel. The number of cost values must
        be equal to the number of media channels.
      target: Target KPI to use, like for example sales.
      extra_features: Other variables to add to the model.
      degrees_seasonality: Number of degrees to use for seasonality. Default is
        2.
      seasonality_frequency: Frequency of the time period used. Default is 52 as
        in 52 weeks per year.
      weekday_seasonality: In case of daily data, also estimate seven weekday
        parameters.
      media_names: Names of the media channels passed.
      number_warmup: Number of warm up samples. Default is 1000.
      number_samples: Number of samples during sampling. Default is 1000.
      number_chains: Number of chains to sample. Default is 2.
      target_accept_prob: Target acceptance probability for step size in the
        NUTS sampler. Default is .85.
      init_strategy: Initialization function for numpyro NUTS. The available
        options can be found in
        https://num.pyro.ai/en/stable/utilities.html#initialization-strategies.
        Default is numpyro.infer.init_to_median.
      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.
      seed: Seed to use for PRNGKey during training. For better replicability
        run all different trainings with the same seed.
    """
    if media.ndim not in (2, 3):
      raise ValueError(
          "Media data must have either 2 dims for national model or 3 for geo "
          "models.")
    if media.ndim == 3 and media_prior.ndim == 1:
      media_prior = jnp.expand_dims(media_prior, axis=-1)

    if media.shape[1] != len(media_prior):
      raise ValueError("The number of data channels provided must match the "
                       "number of cost values.")
    if media.min() < 0:
      raise ValueError("Media values must be greater or equal to zero.")

    if custom_priors:
      not_used_custom_priors = set(custom_priors.keys()).difference(
          self._prior_names)
      if not_used_custom_priors:
        raise ValueError(
            "The following passed custom priors dont have a match in the model."
            " Please double check the names have been written correctly: %s" %
            not_used_custom_priors)
      custom_priors = self._preprocess_custom_priors(
          custom_priors=custom_priors)
      geo_custom_priors = set(custom_priors.keys()).intersection(
          models.GEO_ONLY_PRIORS)
      if media.ndim == 2 and geo_custom_priors:
        raise ValueError(
            "The given data is for national models but custom_prior contains "
            "priors for the geo version of the model. Please either remove geo "
            "priors for national model or pass media data with geo dimension.")
    else:
      custom_priors = {}

    if weekday_seasonality and seasonality_frequency == 52:
      logging.warn("You have chosen daily seasonality and frequency 52 "
                   "(weekly), please check you made the right seasonality "
                   "choices.")

    if extra_features is not None:
      extra_features = jnp.array(extra_features)

    if seed is None:
      seed = utils.get_time_seed()

    train_media_size = media.shape[0]
    kernel = numpyro.infer.NUTS(
        model=self._model_function,
        target_accept_prob=target_accept_prob,
        init_strategy=init_strategy)

    mcmc = numpyro.infer.MCMC(
        sampler=kernel,
        num_warmup=number_warmup,
        num_samples=number_samples,
        num_chains=number_chains)
    mcmc.run(
        rng_key=jax.random.PRNGKey(seed),
        media_data=jnp.array(media),
        extra_features=extra_features,
        target_data=jnp.array(target),
        media_prior=jnp.array(media_prior),
        degrees_seasonality=degrees_seasonality,
        frequency=seasonality_frequency,
        transform_function=self._model_transform_function,
        weekday_seasonality=weekday_seasonality,
        custom_priors=custom_priors)

    self.custom_priors = custom_priors
    if media_names is not None:
      self.media_names = media_names
    else:
      self.media_names = [f"channel_{i}" for i in range(media.shape[1])]
    self.n_media_channels = media.shape[1]
    self.n_geos = media.shape[2] if media.ndim == 3 else 1
    self._media_prior = media_prior
    self.trace = mcmc.get_samples()
    self._number_warmup = number_warmup
    self._number_samples = number_samples
    self._number_chains = number_chains
    self._target = target
    self._train_media_size = train_media_size
    self._degrees_seasonality = degrees_seasonality
    self._seasonality_frequency = seasonality_frequency
    self._weekday_seasonality = weekday_seasonality
    self.media = media
    self._extra_features = extra_features# jax-devicearray
    self._mcmc = mcmc
    logging.info("Model has been fitted")

  def print_summary(self) -> None:
    """Calls print_summary function from numpyro to print parameters summary.
    """
    # TODO(): add name selection for print.
    self._mcmc.print_summary()

  @functools.partial(
      jax.jit,
      static_argnums=(0,),
      static_argnames=("degrees_seasonality", "weekday_seasonality",
                       "transform_function", "model"))
  def _predict(
      self,
      rng_key: jnp.ndarray,
      media_data: jnp.ndarray,
      extra_features: Optional[jnp.ndarray],
      media_prior: jnp.ndarray,
      degrees_seasonality: int, frequency: int,
      transform_function: Callable[[Any], jnp.ndarray],
      weekday_seasonality: bool,
      model: Callable[[Any], None],
      posterior_samples: Dict[str, jnp.ndarray],
      custom_priors: Dict[str, Prior]
      ) -> Dict[str, jnp.ndarray]:
    """Encapsulates the numpyro.infer.Predictive function for predict method.

    It serves as a helper jitted function for running predictions.

    Args:
      rng_key: A jax.random.PRNGKey.
      media_data: Media array for needed for the model to run predictions.
      extra_features: Extra features for needed for the model to run.
      media_prior: Cost prior used for training the model.
      degrees_seasonality: Number of degrees for the seasonality.
      frequency: Frequency of the seasonality.
      transform_function: Media transform function to use within the model.
      weekday_seasonality: Allow daily weekday estimation.
      model: Numpyro model to use for numpyro.infer.Predictive.
      posterior_samples: Mapping of the posterior samples.
      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 predictions for the given data.
    """
    return infer.Predictive(
        model=model, posterior_samples=posterior_samples)(
            rng_key=rng_key,
            media_data=media_data,
            extra_features=extra_features,
            media_prior=media_prior,
            target_data=None,
            degrees_seasonality=degrees_seasonality,
            frequency=frequency,
            transform_function=transform_function,
            custom_priors=custom_priors,
            weekday_seasonality=weekday_seasonality)

  def predict(
      self,
      media: jnp.ndarray,
      extra_features: Optional[jnp.ndarray] = None,
      media_gap: Optional[jnp.ndarray] = None,
      target_scaler: Optional[preprocessing.CustomScaler] = None,
      seed: Optional[int] = None
  ) -> jnp.ndarray:
    """Runs the model to obtain predictions for the given input data.

    Predictions returned are distributions, if point estimates are desired one
    can calculate those based on the given distribution.

    Args:
      media: Media array for needed for the model to run predictions.
      extra_features: Extra features for needed for the model to run.
      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.
      seed: Seed to use for PRNGKey during sampling. For replicability run
        this function and any other function that utilises predictions with the
        same seed.

    Returns:
      Predictions for the given media and extra features at a given date index.

    Raises:
      NotFittedModelError: When the model has not been fitted before running
        predict.
    """
    if not hasattr(self, "trace"):
      raise NotFittedModelError("Need to fit the model before running "
                                "predictions.")
    if media_gap is not None:
      if media.ndim != media_gap.ndim:
        raise ValueError("Original media data and media gap must have the same "
                         "number of dimensions.")
      if media.ndim > 1 and media.shape[1] != media_gap.shape[1]:
        raise ValueError("Media gap must have the same numer of media channels"
                         "as the original media data.")
      previous_media = jnp.concatenate(arrays=[self.media, media_gap], axis=0)
      if extra_features is not None:
        previous_extra_features = jnp.concatenate(
            arrays=[
                self._extra_features,
                jnp.zeros((media_gap.shape[0], *self._extra_features.shape[1:]))
            ],
            axis=0)
    else:
      previous_media = self.media
      previous_extra_features = self._extra_features

    full_media = jnp.concatenate(arrays=[previous_media, media], axis=0)
    if extra_features is not None:
      full_extra_features = jnp.concatenate(
          arrays=[previous_extra_features, extra_features], axis=0)
    else:
      full_extra_features = None
    if seed is None:
      seed = utils.get_time_seed()
    prediction = self._predict(
        rng_key=jax.random.PRNGKey(seed=seed),
        media_data=full_media,
        extra_features=full_extra_features,
        media_prior=jnp.array(self._media_prior),
        degrees_seasonality=self._degrees_seasonality,
        frequency=self._seasonality_frequency,
        weekday_seasonality=self._weekday_seasonality,
        transform_function=self._model_transform_function,
        model=self._model_function,
        custom_priors=self.custom_priors,
        posterior_samples=self.trace)["mu"][:, previous_media.shape[0]:]
    if target_scaler:
      prediction = target_scaler.inverse_transform(prediction)

    return prediction

  def reduce_trace(self, nsample: int = 100, seed: int = 0) -> None:
    """Reduces the samples in `trace` to speed up `predict` and optimize.

    Please note this step is not reversible. Only do this after you have
    investigated convergence of the model.

    Args:
      nsample: Target number of samples.
      seed: Random seed for down sampling.

    Raises:
      ValueError: if `nsample` is too big.
    """

    ntrace = len(self.trace["sigma"])
    if ntrace < nsample:
      raise ValueError("nsample is bigger than the actual posterior samples")
    key = jax.random.PRNGKey(seed)
    samples = jax.random.choice(key, ntrace, (nsample,), replace=False)
    for name in self.trace.keys():
      self.trace[name] = self.trace[name][samples]
    logging.info("Reduction is complete")


Prompt:
import abc
from typing import Union

import numpy as np
import torch
import tqdm

class IdentitySampler:
    def run(
        self, features: Union[torch.Tensor, np.ndarray]
    ) -> Union[torch.Tensor, np.ndarray]:
        return features

class BaseSampler(abc.ABC):
    def __init__(self, percentage: float):

        if not 0 < percentage < 1:
            raise ValueError("Percentage value not in (0, 1).")
        self.percentage = percentage


Prompt:
import abc
from typing import Union

import numpy as np
import torch
import tqdm

class IdentitySampler:
    def run(
        self, features: Union[torch.Tensor, np.ndarray]
    ) -> Union[torch.Tensor, np.ndarray]:
        return features

class BaseSampler(abc.ABC):
    def __init__(self, percentage: float):
        if not 0 < percentage < 1:
            raise ValueError("Percentage value not in (0, 1).")
        self.percentage = percentage

    @abc.abstractmethod
    def run(
        self, features: Union[torch.Tensor, np.ndarray]
    ) -> Union[torch.Tensor, np.ndarray]:
        pass

    def _store_type(self, features: Union[torch.Tensor, np.ndarray]) -> None:
        self.features_is_numpy = isinstance(features, np.ndarray)
        if not self.features_is_numpy:
            self.features_device = features.device

    def _restore_type(self, features: torch.Tensor) -> Union[torch.Tensor, np.ndarray]:
        if self.features_is_numpy:
            return features.cpu().numpy()
        return features.to(self.features_device)

class GreedyCoresetSampler(BaseSampler):
    def __init__(
        self,
        percentage: float,
        device: torch.device,
        dimension_to_project_features_to=128,
    ):
        """Greedy Coreset sampling base class."""

        super().__init__(percentage)

        self.device = device
        self.dimension_to_project_features_to = dimension_to_project_features_to


Prompt:
"""PatchCore and PatchCore detection methods."""
import logging
import os
import pickle

import numpy as np
import torch
import torch.nn.functional as F
import tqdm

import patchcore
import patchcore.backbones
import patchcore.common
import patchcore.sampler

LOGGER = logging.getLogger(__name__)

class PatchCore(torch.nn.Module):
    def __init__(self, device):
        """PatchCore anomaly detection class."""
        super(PatchCore, self).__init__()
        self.device = device

    def load(
        self,
        backbone,
        layers_to_extract_from,
        device,
        input_shape,
        pretrain_embed_dimension,
        target_embed_dimension,
        patchsize=3,
        patchstride=1,
        anomaly_score_num_nn=1,
        featuresampler=patchcore.sampler.IdentitySampler(),
        nn_method=patchcore.common.FaissNN(False, 4),
        **kwargs,
    ):
        self.backbone = backbone.to(device)
        self.layers_to_extract_from = layers_to_extract_from
        self.input_shape = input_shape

        self.device = device
        self.patch_maker = PatchMaker(patchsize, stride=patchstride)

        self.forward_modules = torch.nn.ModuleDict({})

        feature_aggregator = patchcore.common.NetworkFeatureAggregator(
            self.backbone, self.layers_to_extract_from, self.device
        )
        feature_dimensions = feature_aggregator.feature_dimensions(input_shape)
        self.forward_modules["feature_aggregator"] = feature_aggregator

        preprocessing = patchcore.common.Preprocessing(
            feature_dimensions, pretrain_embed_dimension
        )
        self.forward_modules["preprocessing"] = preprocessing

        self.target_embed_dimension = target_embed_dimension
        preadapt_aggregator = patchcore.common.Aggregator(
            target_dim=target_embed_dimension
        )

        _ = preadapt_aggregator.to(self.device)

        self.forward_modules["preadapt_aggregator"] = preadapt_aggregator

        self.anomaly_scorer = patchcore.common.NearestNeighbourScorer(
            n_nearest_neighbours=anomaly_score_num_nn, nn_method=nn_method
        )

        self.anomaly_segmentor = patchcore.common.RescaleSegmentor(
            device=self.device, target_size=input_shape[-2:]
        )

        self.featuresampler = featuresampler

    def embed(self, data):
        if isinstance(data, torch.utils.data.DataLoader):
            features = []
            for image in data:
                if isinstance(image, dict):
                    image = image["image"]
                with torch.no_grad():
                    input_image = image.to(torch.float).to(self.device)
                    features.append(self._embed(input_image))
            return features
        return self._embed(data)

    def _embed(self, images, detach=True, provide_patch_shapes=False):
        """Returns feature embeddings for images."""

        def _detach(features):

            if detach:
                return [x.detach().cpu().numpy() for x in features]
            return features