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

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

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import preprocessing

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

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

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

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

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

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

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

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

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

  Returns:
    A list of tuples with the lower and upper bound for each media channel.
  """
  if media.ndim == 3:
    lower_pct = jnp.expand_dims(lower_pct, axis=-1)
    upper_pct = jnp.expand_dims(upper_pct, axis=-1)

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

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

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

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

def _generate_starting_values(
    n_time_periods: int, media: jnp.ndarray,
    media_scaler: preprocessing.CustomScaler,
    budget: Union[float, int],
    prices: jnp.ndarray,
) -> jnp.ndarray:
  """Generates starting values based on historic allocation and budget.

  In order to make a comparison we can take the allocation of the last
  `n_time_periods` and scale it based on the given budget. Given this, one can
  compare how this initial values (based on average historic allocation) compare
  to the output of the optimisation in terms of sales/KPI.

  Args:
    n_time_periods: Number of time periods the optimization will be done with.
    media: Historic media data the model was trained with.
    media_scaler: Scaler that was used to scale the media data before training.
    budget: Total budget to allocate during the optimization time.
    prices: An array with shape (n_media_channels,) for the cost of each media
      channel unit.

  Returns:
    An array with the starting value for each media channel for the
      optimization.
  """

  previous_allocation = media.mean(axis=0) * n_time_periods
  if media_scaler:  # Scale before sum as geo scaler has shape (c, g).
    previous_allocation = media_scaler.inverse_transform(previous_allocation)

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

  avg_spend_per_channel = previous_allocation * prices
  pct_spend_per_channel = avg_spend_per_channel / avg_spend_per_channel.sum()
  budget_per_channel = budget * pct_spend_per_channel
  media_unit_per_channel = budget_per_channel / prices
  return media_unit_per_channel


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

"""Utilities for preprocessing dataset for training LightweightMMM."""

import copy
from typing import Callable, List, Optional, Sequence, Tuple, Union

import jax.numpy as jnp
import pandas as pd
from sklearn import base

from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant
from lightweight_mmm.core import core_utils

class NotFittedScalerError(Exception):
  pass

class CustomScaler(base.TransformerMixin):
  """Class to scale your data based on multiplications and divisions.

  This scaler can be used in two fashions for both the multiplication and
  division operation.
  - By specifying a value to use for the scaling operation.
  - By specifying an operation used at column level to calculate the value
    for the actual scaling operation.

  Eg. if one wants to scale the dataset by multiply by 100 you can directly
  pass multiply_by=100. Value can also be an array with as many values
  as column has the data being scaled. But if you want to multiply by the mean
  value of each column, then you can pass multiply_operation=jnp.mean (or any
  other operation desired).

  Operation parameters have the upper hand in the cases where both values and
  operations are passed, values will be ignored in this case.

  Scaler must be fit first in order to call the transform method.

  Attributes.
    divide_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for division during scaling.
    divide_by: Numbers(s) by which to divide data in the scaling process. Since
      the scaler is applied to axis 0 of the data, the shape of divide_by must
      be consistent with division into the data. For example, if data.shape =
      (100, 3, 5) then divide_by.shape can be (3, 5) or (5,) or a number. If
      divide_operation is given, this divide_by value will be ignored.
    multiply_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for multiplication during scaling.
    multiply_by: Numbers(s) by which to multiply data in the scaling process.
      Since the scaler is applied to axis 0 of the data, the shape of
      multiply_by must be consistent with multiplication into the data. For
      example, if data.shape = (100, 3, 5) then multiply_by.shape can be (3, 5)
      or (5,) or a number. If multiply_operation is given, this multiply_by
      value will be ignored.
  """

  def __init__(
      self,
      divide_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      divide_by: Optional[Union[float, int, jnp.ndarray]] = 1,
      multiply_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      multiply_by: Optional[Union[float, int, jnp.ndarray]] = 1.) -> None:
    """Constructor for the CustomScaler class."""
    if all([
        divide_by is None, divide_operation is None, multiply_by is None,
        multiply_operation is None
    ]):
      raise ValueError("No values for transformations were provided and this "
                       "scaler will fail. Please instantiate a valid one")

    if divide_operation is None and divide_by is None:
      raise ValueError(
          "Either a division operation or value needs to be passed. If "
          "you dont want to use a division to scale your data just "
          "pass divide_by=1.")
    elif divide_operation is not None:
      self.divide_operation = divide_operation
    else:
      self.divide_by = divide_by

    if multiply_operation is None and multiply_by is None:
      raise ValueError(
          "Either a multiplication operation or value needs to be passed. If "
          "you dont want to use a multiplication to scale your data just "
          "pass multiply_by=1.")
    elif multiply_operation is not None:
      self.multiply_operation = multiply_operation
    else:
      self.multiply_by = multiply_by

  def fit(self, data: jnp.ndarray) -> None:
    """Figures out values for transformations based on the specified operations.

    Args:
      data: Input dataset to use for fitting.
    """

    if hasattr(self, "divide_operation"):
      self.divide_by = jnp.apply_along_axis(
          func1d=self.divide_operation, axis=0, arr=data)
    elif isinstance(self.divide_by, int) or isinstance(self.divide_by, float):
      self.divide_by = self.divide_by * jnp.ones(data.shape[1:])
    if hasattr(self, "multiply_operation"):
      self.multiply_by = jnp.apply_along_axis(
          func1d=self.multiply_operation, axis=0, arr=data)
    elif isinstance(self.multiply_by, int) or isinstance(
        self.multiply_by, float):
      self.multiply_by = self.multiply_by * jnp.ones(data.shape[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.

"""Utilities for preprocessing dataset for training LightweightMMM."""

import copy
from typing import Callable, List, Optional, Sequence, Tuple, Union

import jax.numpy as jnp
import pandas as pd
from sklearn import base

from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant
from lightweight_mmm.core import core_utils

class NotFittedScalerError(Exception):
  pass

class CustomScaler(base.TransformerMixin):
  """Class to scale your data based on multiplications and divisions.

  This scaler can be used in two fashions for both the multiplication and
  division operation.
  - By specifying a value to use for the scaling operation.
  - By specifying an operation used at column level to calculate the value
    for the actual scaling operation.

  Eg. if one wants to scale the dataset by multiply by 100 you can directly
  pass multiply_by=100. Value can also be an array with as many values
  as column has the data being scaled. But if you want to multiply by the mean
  value of each column, then you can pass multiply_operation=jnp.mean (or any
  other operation desired).

  Operation parameters have the upper hand in the cases where both values and
  operations are passed, values will be ignored in this case.

  Scaler must be fit first in order to call the transform method.

  Attributes.
    divide_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for division during scaling.
    divide_by: Numbers(s) by which to divide data in the scaling process. Since
      the scaler is applied to axis 0 of the data, the shape of divide_by must
      be consistent with division into the data. For example, if data.shape =
      (100, 3, 5) then divide_by.shape can be (3, 5) or (5,) or a number. If
      divide_operation is given, this divide_by value will be ignored.
    multiply_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for multiplication during scaling.
    multiply_by: Numbers(s) by which to multiply data in the scaling process.
      Since the scaler is applied to axis 0 of the data, the shape of
      multiply_by must be consistent with multiplication into the data. For
      example, if data.shape = (100, 3, 5) then multiply_by.shape can be (3, 5)
      or (5,) or a number. If multiply_operation is given, this multiply_by
      value will be ignored.
  """

  def __init__(
      self,
      divide_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      divide_by: Optional[Union[float, int, jnp.ndarray]] = 1,
      multiply_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      multiply_by: Optional[Union[float, int, jnp.ndarray]] = 1.) -> None:
    """Constructor for the CustomScaler class."""
    if all([
        divide_by is None, divide_operation is None, multiply_by is None,
        multiply_operation is None
    ]):
      raise ValueError("No values for transformations were provided and this "
                       "scaler will fail. Please instantiate a valid one")

    if divide_operation is None and divide_by is None:
      raise ValueError(
          "Either a division operation or value needs to be passed. If "
          "you dont want to use a division to scale your data just "
          "pass divide_by=1.")
    elif divide_operation is not None:
      self.divide_operation = divide_operation
    else:
      self.divide_by = divide_by

    if multiply_operation is None and multiply_by is None:
      raise ValueError(
          "Either a multiplication operation or value needs to be passed. If "
          "you dont want to use a multiplication to scale your data just "
          "pass multiply_by=1.")
    elif multiply_operation is not None:
      self.multiply_operation = multiply_operation
    else:
      self.multiply_by = multiply_by

  def fit(self, data: jnp.ndarray) -> None:
    """Figures out values for transformations based on the specified operations.

    Args:
      data: Input dataset to use for fitting.
    """
    if hasattr(self, "divide_operation"):
      self.divide_by = jnp.apply_along_axis(
          func1d=self.divide_operation, axis=0, arr=data)
    elif isinstance(self.divide_by, int) or isinstance(self.divide_by, float):
      self.divide_by = self.divide_by * jnp.ones(data.shape[1:])
    if hasattr(self, "multiply_operation"):
      self.multiply_by = jnp.apply_along_axis(
          func1d=self.multiply_operation, axis=0, arr=data)
    elif isinstance(self.multiply_by, int) or isinstance(
        self.multiply_by, float):
      self.multiply_by = self.multiply_by * jnp.ones(data.shape[1:])

  def transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Applies transformation based on fitted values.

    It can only be called if scaler was fit first.

    Args:
      data: Input dataset to transform.

    Returns:
      Transformed array.
    """

    if not hasattr(self, "divide_by") or not hasattr(self, "multiply_by"):
      raise NotFittedScalerError(
          "transform is called without fit being called previously. Please "
          "fit scaler first.")
    return self.multiply_by * data / self.divide_by


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

"""Utilities for preprocessing dataset for training LightweightMMM."""

import copy
from typing import Callable, List, Optional, Sequence, Tuple, Union

import jax.numpy as jnp
import pandas as pd
from sklearn import base

from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant
from lightweight_mmm.core import core_utils

class NotFittedScalerError(Exception):
  pass

class CustomScaler(base.TransformerMixin):
  """Class to scale your data based on multiplications and divisions.

  This scaler can be used in two fashions for both the multiplication and
  division operation.
  - By specifying a value to use for the scaling operation.
  - By specifying an operation used at column level to calculate the value
    for the actual scaling operation.

  Eg. if one wants to scale the dataset by multiply by 100 you can directly
  pass multiply_by=100. Value can also be an array with as many values
  as column has the data being scaled. But if you want to multiply by the mean
  value of each column, then you can pass multiply_operation=jnp.mean (or any
  other operation desired).

  Operation parameters have the upper hand in the cases where both values and
  operations are passed, values will be ignored in this case.

  Scaler must be fit first in order to call the transform method.

  Attributes.
    divide_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for division during scaling.
    divide_by: Numbers(s) by which to divide data in the scaling process. Since
      the scaler is applied to axis 0 of the data, the shape of divide_by must
      be consistent with division into the data. For example, if data.shape =
      (100, 3, 5) then divide_by.shape can be (3, 5) or (5,) or a number. If
      divide_operation is given, this divide_by value will be ignored.
    multiply_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for multiplication during scaling.
    multiply_by: Numbers(s) by which to multiply data in the scaling process.
      Since the scaler is applied to axis 0 of the data, the shape of
      multiply_by must be consistent with multiplication into the data. For
      example, if data.shape = (100, 3, 5) then multiply_by.shape can be (3, 5)
      or (5,) or a number. If multiply_operation is given, this multiply_by
      value will be ignored.
  """

  def __init__(
      self,
      divide_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      divide_by: Optional[Union[float, int, jnp.ndarray]] = 1,
      multiply_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      multiply_by: Optional[Union[float, int, jnp.ndarray]] = 1.) -> None:
    """Constructor for the CustomScaler class."""
    if all([
        divide_by is None, divide_operation is None, multiply_by is None,
        multiply_operation is None
    ]):
      raise ValueError("No values for transformations were provided and this "
                       "scaler will fail. Please instantiate a valid one")

    if divide_operation is None and divide_by is None:
      raise ValueError(
          "Either a division operation or value needs to be passed. If "
          "you dont want to use a division to scale your data just "
          "pass divide_by=1.")
    elif divide_operation is not None:
      self.divide_operation = divide_operation
    else:
      self.divide_by = divide_by

    if multiply_operation is None and multiply_by is None:
      raise ValueError(
          "Either a multiplication operation or value needs to be passed. If "
          "you dont want to use a multiplication to scale your data just "
          "pass multiply_by=1.")
    elif multiply_operation is not None:
      self.multiply_operation = multiply_operation
    else:
      self.multiply_by = multiply_by

  def fit(self, data: jnp.ndarray) -> None:
    """Figures out values for transformations based on the specified operations.

    Args:
      data: Input dataset to use for fitting.
    """
    if hasattr(self, "divide_operation"):
      self.divide_by = jnp.apply_along_axis(
          func1d=self.divide_operation, axis=0, arr=data)
    elif isinstance(self.divide_by, int) or isinstance(self.divide_by, float):
      self.divide_by = self.divide_by * jnp.ones(data.shape[1:])
    if hasattr(self, "multiply_operation"):
      self.multiply_by = jnp.apply_along_axis(
          func1d=self.multiply_operation, axis=0, arr=data)
    elif isinstance(self.multiply_by, int) or isinstance(
        self.multiply_by, float):
      self.multiply_by = self.multiply_by * jnp.ones(data.shape[1:])

  def transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Applies transformation based on fitted values.

    It can only be called if scaler was fit first.

    Args:
      data: Input dataset to transform.

    Returns:
      Transformed array.
    """
    if not hasattr(self, "divide_by") or not hasattr(self, "multiply_by"):
      raise NotFittedScalerError(
          "transform is called without fit being called previously. Please "
          "fit scaler first.")
    return self.multiply_by * data / self.divide_by

  def fit_transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Fits the values and applies transformation to the input data.

    Args:
      data: Input dataset.

    Returns:
      Transformed array.
    """
    self.fit(data)
    return self.transform(data)

  def inverse_transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Runs inverse transformation to get original values.

    Args:
      data: Input dataset.

    Returns:
      Dataset with the inverse transformation applied.
    """
    return self.divide_by * data / self.multiply_by

def _compute_correlations(
    features: jnp.ndarray,
    target: jnp.ndarray,
    feature_names: List[str],
    ) -> List[pd.DataFrame]:
  """Computes feature-feature and feature-target correlations.

  Helper function for DataQualityCheck.

  Args:
    features: Features for media mix model (media and non-media variables).
    target: Target variable for media mix model.
    feature_names: Names of media channels to be added to the output dataframes.

  Returns:
    List of dataframes containing Pearson correlation coefficients between each
      feature, as well as between features and the target variable. For
      national-level data the list contains just one dataframe, and for
      geo-level data the list contains one dataframe for each geo.

  Raises:
    ValueError: If features and target have incompatible shapes (e.g. one is
      geo-level and the other national-level).
  """
  if not ((features.ndim == 2 and target.ndim == 1) or
          (features.ndim == 3 and target.ndim == 2)):
    raise ValueError(f"Incompatible shapes between features {features.shape}"
                     f" and target {target.shape}.")

  number_of_geos = core_utils.get_number_geos(features)
  correlation_matrix_output = []
  for i_geo in range(number_of_geos):

    if number_of_geos == 1:
      features_and_target = jnp.concatenate(
          [features, jnp.expand_dims(target, axis=1)], axis=1)
    else:
      features_and_target = jnp.concatenate(
          [features[:, :, i_geo],
           jnp.expand_dims(target[:, i_geo], axis=1)],
          axis=1)

    covariance_matrix = jnp.cov(features_and_target, rowvar=False)
    standard_deviations = jnp.std(features_and_target, axis=0, ddof=1)
    correlation_matrix = covariance_matrix / jnp.outer(standard_deviations,
                                                       standard_deviations)
    correlation_matrix = pd.DataFrame(
        correlation_matrix,
        columns=feature_names + ["target"],
        index=feature_names + ["target"],
        dtype=float)
    correlation_matrix_output.append(correlation_matrix)

  return correlation_matrix_output

def _compute_variances(
    features: jnp.ndarray,
    feature_names: Sequence[str],
    geo_names: Sequence[str],
) -> pd.DataFrame:
  """Computes variances over time for each feature.

  In general, higher variance is better since it creates more signal for the
  regression analysis. However, if the features have not been scaled (divided by
  the mean), then the variance can take any value and this analysis is not
  meaningful.

  Args:
    features: Features for media mix model (media and non-media variables).
    feature_names: Names of media channels to be added to the output dataframe.
    geo_names: Names of geos to be added to the output dataframes.

  Returns:
    Dataframe containing the variance over time for each feature. This dataframe
      contains one row per geo, and just a single row for national data.

  Raises:
    ValueError: If the number of geos in features does not match the number of
    supplied geo_names.
  """

  number_of_geos = core_utils.get_number_geos(features)

  if len(geo_names) != number_of_geos:
    raise ValueError("The number of geos in features does not match the length "
                     "of geo_names")

  variances_as_series = []
  for i_geo in range(number_of_geos):
    features_for_this_geo = features[...,
                                     i_geo] if number_of_geos > 1 else features
    variances_as_series.append(
        pd.DataFrame(data=features_for_this_geo).var(axis=0, ddof=0))

  variances = pd.concat(variances_as_series, axis=1)
  variances.columns = geo_names
  variances.index = copy.copy(feature_names)

  return variances


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

"""Utilities for preprocessing dataset for training LightweightMMM."""

import copy
from typing import Callable, List, Optional, Sequence, Tuple, Union

import jax.numpy as jnp
import pandas as pd
from sklearn import base

from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant
from lightweight_mmm.core import core_utils

class NotFittedScalerError(Exception):
  pass

class CustomScaler(base.TransformerMixin):
  """Class to scale your data based on multiplications and divisions.

  This scaler can be used in two fashions for both the multiplication and
  division operation.
  - By specifying a value to use for the scaling operation.
  - By specifying an operation used at column level to calculate the value
    for the actual scaling operation.

  Eg. if one wants to scale the dataset by multiply by 100 you can directly
  pass multiply_by=100. Value can also be an array with as many values
  as column has the data being scaled. But if you want to multiply by the mean
  value of each column, then you can pass multiply_operation=jnp.mean (or any
  other operation desired).

  Operation parameters have the upper hand in the cases where both values and
  operations are passed, values will be ignored in this case.

  Scaler must be fit first in order to call the transform method.

  Attributes.
    divide_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for division during scaling.
    divide_by: Numbers(s) by which to divide data in the scaling process. Since
      the scaler is applied to axis 0 of the data, the shape of divide_by must
      be consistent with division into the data. For example, if data.shape =
      (100, 3, 5) then divide_by.shape can be (3, 5) or (5,) or a number. If
      divide_operation is given, this divide_by value will be ignored.
    multiply_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for multiplication during scaling.
    multiply_by: Numbers(s) by which to multiply data in the scaling process.
      Since the scaler is applied to axis 0 of the data, the shape of
      multiply_by must be consistent with multiplication into the data. For
      example, if data.shape = (100, 3, 5) then multiply_by.shape can be (3, 5)
      or (5,) or a number. If multiply_operation is given, this multiply_by
      value will be ignored.
  """

  def __init__(
      self,
      divide_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      divide_by: Optional[Union[float, int, jnp.ndarray]] = 1,
      multiply_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      multiply_by: Optional[Union[float, int, jnp.ndarray]] = 1.) -> None:
    """Constructor for the CustomScaler class."""
    if all([
        divide_by is None, divide_operation is None, multiply_by is None,
        multiply_operation is None
    ]):
      raise ValueError("No values for transformations were provided and this "
                       "scaler will fail. Please instantiate a valid one")

    if divide_operation is None and divide_by is None:
      raise ValueError(
          "Either a division operation or value needs to be passed. If "
          "you dont want to use a division to scale your data just "
          "pass divide_by=1.")
    elif divide_operation is not None:
      self.divide_operation = divide_operation
    else:
      self.divide_by = divide_by

    if multiply_operation is None and multiply_by is None:
      raise ValueError(
          "Either a multiplication operation or value needs to be passed. If "
          "you dont want to use a multiplication to scale your data just "
          "pass multiply_by=1.")
    elif multiply_operation is not None:
      self.multiply_operation = multiply_operation
    else:
      self.multiply_by = multiply_by

  def fit(self, data: jnp.ndarray) -> None:
    """Figures out values for transformations based on the specified operations.

    Args:
      data: Input dataset to use for fitting.
    """
    if hasattr(self, "divide_operation"):
      self.divide_by = jnp.apply_along_axis(
          func1d=self.divide_operation, axis=0, arr=data)
    elif isinstance(self.divide_by, int) or isinstance(self.divide_by, float):
      self.divide_by = self.divide_by * jnp.ones(data.shape[1:])
    if hasattr(self, "multiply_operation"):
      self.multiply_by = jnp.apply_along_axis(
          func1d=self.multiply_operation, axis=0, arr=data)
    elif isinstance(self.multiply_by, int) or isinstance(
        self.multiply_by, float):
      self.multiply_by = self.multiply_by * jnp.ones(data.shape[1:])

  def transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Applies transformation based on fitted values.

    It can only be called if scaler was fit first.

    Args:
      data: Input dataset to transform.

    Returns:
      Transformed array.
    """
    if not hasattr(self, "divide_by") or not hasattr(self, "multiply_by"):
      raise NotFittedScalerError(
          "transform is called without fit being called previously. Please "
          "fit scaler first.")
    return self.multiply_by * data / self.divide_by

  def fit_transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Fits the values and applies transformation to the input data.

    Args:
      data: Input dataset.

    Returns:
      Transformed array.
    """
    self.fit(data)
    return self.transform(data)

  def inverse_transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Runs inverse transformation to get original values.

    Args:
      data: Input dataset.

    Returns:
      Dataset with the inverse transformation applied.
    """
    return self.divide_by * data / self.multiply_by

def _compute_correlations(
    features: jnp.ndarray,
    target: jnp.ndarray,
    feature_names: List[str],
    ) -> List[pd.DataFrame]:
  """Computes feature-feature and feature-target correlations.

  Helper function for DataQualityCheck.

  Args:
    features: Features for media mix model (media and non-media variables).
    target: Target variable for media mix model.
    feature_names: Names of media channels to be added to the output dataframes.

  Returns:
    List of dataframes containing Pearson correlation coefficients between each
      feature, as well as between features and the target variable. For
      national-level data the list contains just one dataframe, and for
      geo-level data the list contains one dataframe for each geo.

  Raises:
    ValueError: If features and target have incompatible shapes (e.g. one is
      geo-level and the other national-level).
  """
  if not ((features.ndim == 2 and target.ndim == 1) or
          (features.ndim == 3 and target.ndim == 2)):
    raise ValueError(f"Incompatible shapes between features {features.shape}"
                     f" and target {target.shape}.")

  number_of_geos = core_utils.get_number_geos(features)
  correlation_matrix_output = []
  for i_geo in range(number_of_geos):

    if number_of_geos == 1:
      features_and_target = jnp.concatenate(
          [features, jnp.expand_dims(target, axis=1)], axis=1)
    else:
      features_and_target = jnp.concatenate(
          [features[:, :, i_geo],
           jnp.expand_dims(target[:, i_geo], axis=1)],
          axis=1)

    covariance_matrix = jnp.cov(features_and_target, rowvar=False)
    standard_deviations = jnp.std(features_and_target, axis=0, ddof=1)
    correlation_matrix = covariance_matrix / jnp.outer(standard_deviations,
                                                       standard_deviations)
    correlation_matrix = pd.DataFrame(
        correlation_matrix,
        columns=feature_names + ["target"],
        index=feature_names + ["target"],
        dtype=float)
    correlation_matrix_output.append(correlation_matrix)

  return correlation_matrix_output

def _compute_variances(
    features: jnp.ndarray,
    feature_names: Sequence[str],
    geo_names: Sequence[str],
) -> pd.DataFrame:
  """Computes variances over time for each feature.

  In general, higher variance is better since it creates more signal for the
  regression analysis. However, if the features have not been scaled (divided by
  the mean), then the variance can take any value and this analysis is not
  meaningful.

  Args:
    features: Features for media mix model (media and non-media variables).
    feature_names: Names of media channels to be added to the output dataframe.
    geo_names: Names of geos to be added to the output dataframes.

  Returns:
    Dataframe containing the variance over time for each feature. This dataframe
      contains one row per geo, and just a single row for national data.

  Raises:
    ValueError: If the number of geos in features does not match the number of
    supplied geo_names.
  """
  number_of_geos = core_utils.get_number_geos(features)

  if len(geo_names) != number_of_geos:
    raise ValueError("The number of geos in features does not match the length "
                     "of geo_names")

  variances_as_series = []
  for i_geo in range(number_of_geos):
    features_for_this_geo = features[...,
                                     i_geo] if number_of_geos > 1 else features
    variances_as_series.append(
        pd.DataFrame(data=features_for_this_geo).var(axis=0, ddof=0))

  variances = pd.concat(variances_as_series, axis=1)
  variances.columns = geo_names
  variances.index = copy.copy(feature_names)

  return variances

def _compute_spend_fractions(
    cost_data: jnp.ndarray,
    channel_names: Optional[Sequence[str]] = None,
    output_column_name: str = "fraction of spend") -> pd.DataFrame:
  """Computes fraction of total spend for each media channel.

  Args:
    cost_data: Spend (can be normalized or not) per channel.
    channel_names: Names of media channels to be added to the output dataframe.
    output_column_name: Name of the column in the output dataframe, denoting the
      fraction of the total spend in each media channel.

  Returns:
    Dataframe containing fraction of the total spend in each channel.

  Raises:
    ValueError if any of the costs are zero or negative.
  """

  cost_df = pd.DataFrame(
      cost_data, index=channel_names, columns=[output_column_name])

  if (cost_df[output_column_name] <= 0).any():
    raise ValueError("Values in cost_data must all be positive.")

  normalized_cost_df = cost_df.div(cost_df.sum(axis=0), axis=1).round(4)
  return normalized_cost_df


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

"""Utilities for preprocessing dataset for training LightweightMMM."""

import copy
from typing import Callable, List, Optional, Sequence, Tuple, Union

import jax.numpy as jnp
import pandas as pd
from sklearn import base

from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant
from lightweight_mmm.core import core_utils

class NotFittedScalerError(Exception):
  pass

class CustomScaler(base.TransformerMixin):
  """Class to scale your data based on multiplications and divisions.

  This scaler can be used in two fashions for both the multiplication and
  division operation.
  - By specifying a value to use for the scaling operation.
  - By specifying an operation used at column level to calculate the value
    for the actual scaling operation.

  Eg. if one wants to scale the dataset by multiply by 100 you can directly
  pass multiply_by=100. Value can also be an array with as many values
  as column has the data being scaled. But if you want to multiply by the mean
  value of each column, then you can pass multiply_operation=jnp.mean (or any
  other operation desired).

  Operation parameters have the upper hand in the cases where both values and
  operations are passed, values will be ignored in this case.

  Scaler must be fit first in order to call the transform method.

  Attributes.
    divide_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for division during scaling.
    divide_by: Numbers(s) by which to divide data in the scaling process. Since
      the scaler is applied to axis 0 of the data, the shape of divide_by must
      be consistent with division into the data. For example, if data.shape =
      (100, 3, 5) then divide_by.shape can be (3, 5) or (5,) or a number. If
      divide_operation is given, this divide_by value will be ignored.
    multiply_operation: Operation to apply over axis 0 of the fitting data to
      obtain the value that will be used for multiplication during scaling.
    multiply_by: Numbers(s) by which to multiply data in the scaling process.
      Since the scaler is applied to axis 0 of the data, the shape of
      multiply_by must be consistent with multiplication into the data. For
      example, if data.shape = (100, 3, 5) then multiply_by.shape can be (3, 5)
      or (5,) or a number. If multiply_operation is given, this multiply_by
      value will be ignored.
  """

  def __init__(
      self,
      divide_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      divide_by: Optional[Union[float, int, jnp.ndarray]] = 1,
      multiply_operation: Optional[Callable[[jnp.ndarray], jnp.float32]] = None,
      multiply_by: Optional[Union[float, int, jnp.ndarray]] = 1.) -> None:
    """Constructor for the CustomScaler class."""
    if all([
        divide_by is None, divide_operation is None, multiply_by is None,
        multiply_operation is None
    ]):
      raise ValueError("No values for transformations were provided and this "
                       "scaler will fail. Please instantiate a valid one")

    if divide_operation is None and divide_by is None:
      raise ValueError(
          "Either a division operation or value needs to be passed. If "
          "you dont want to use a division to scale your data just "
          "pass divide_by=1.")
    elif divide_operation is not None:
      self.divide_operation = divide_operation
    else:
      self.divide_by = divide_by

    if multiply_operation is None and multiply_by is None:
      raise ValueError(
          "Either a multiplication operation or value needs to be passed. If "
          "you dont want to use a multiplication to scale your data just "
          "pass multiply_by=1.")
    elif multiply_operation is not None:
      self.multiply_operation = multiply_operation
    else:
      self.multiply_by = multiply_by

  def fit(self, data: jnp.ndarray) -> None:
    """Figures out values for transformations based on the specified operations.

    Args:
      data: Input dataset to use for fitting.
    """
    if hasattr(self, "divide_operation"):
      self.divide_by = jnp.apply_along_axis(
          func1d=self.divide_operation, axis=0, arr=data)
    elif isinstance(self.divide_by, int) or isinstance(self.divide_by, float):
      self.divide_by = self.divide_by * jnp.ones(data.shape[1:])
    if hasattr(self, "multiply_operation"):
      self.multiply_by = jnp.apply_along_axis(
          func1d=self.multiply_operation, axis=0, arr=data)
    elif isinstance(self.multiply_by, int) or isinstance(
        self.multiply_by, float):
      self.multiply_by = self.multiply_by * jnp.ones(data.shape[1:])

  def transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Applies transformation based on fitted values.

    It can only be called if scaler was fit first.

    Args:
      data: Input dataset to transform.

    Returns:
      Transformed array.
    """
    if not hasattr(self, "divide_by") or not hasattr(self, "multiply_by"):
      raise NotFittedScalerError(
          "transform is called without fit being called previously. Please "
          "fit scaler first.")
    return self.multiply_by * data / self.divide_by

  def fit_transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Fits the values and applies transformation to the input data.

    Args:
      data: Input dataset.

    Returns:
      Transformed array.
    """
    self.fit(data)
    return self.transform(data)

  def inverse_transform(self, data: jnp.ndarray) -> jnp.ndarray:
    """Runs inverse transformation to get original values.

    Args:
      data: Input dataset.

    Returns:
      Dataset with the inverse transformation applied.
    """
    return self.divide_by * data / self.multiply_by

def _compute_correlations(
    features: jnp.ndarray,
    target: jnp.ndarray,
    feature_names: List[str],
    ) -> List[pd.DataFrame]:
  """Computes feature-feature and feature-target correlations.

  Helper function for DataQualityCheck.

  Args:
    features: Features for media mix model (media and non-media variables).
    target: Target variable for media mix model.
    feature_names: Names of media channels to be added to the output dataframes.

  Returns:
    List of dataframes containing Pearson correlation coefficients between each
      feature, as well as between features and the target variable. For
      national-level data the list contains just one dataframe, and for
      geo-level data the list contains one dataframe for each geo.

  Raises:
    ValueError: If features and target have incompatible shapes (e.g. one is
      geo-level and the other national-level).
  """
  if not ((features.ndim == 2 and target.ndim == 1) or
          (features.ndim == 3 and target.ndim == 2)):
    raise ValueError(f"Incompatible shapes between features {features.shape}"
                     f" and target {target.shape}.")

  number_of_geos = core_utils.get_number_geos(features)
  correlation_matrix_output = []
  for i_geo in range(number_of_geos):

    if number_of_geos == 1:
      features_and_target = jnp.concatenate(
          [features, jnp.expand_dims(target, axis=1)], axis=1)
    else:
      features_and_target = jnp.concatenate(
          [features[:, :, i_geo],
           jnp.expand_dims(target[:, i_geo], axis=1)],
          axis=1)

    covariance_matrix = jnp.cov(features_and_target, rowvar=False)
    standard_deviations = jnp.std(features_and_target, axis=0, ddof=1)
    correlation_matrix = covariance_matrix / jnp.outer(standard_deviations,
                                                       standard_deviations)
    correlation_matrix = pd.DataFrame(
        correlation_matrix,
        columns=feature_names + ["target"],
        index=feature_names + ["target"],
        dtype=float)
    correlation_matrix_output.append(correlation_matrix)

  return correlation_matrix_output

def _compute_variances(
    features: jnp.ndarray,
    feature_names: Sequence[str],
    geo_names: Sequence[str],
) -> pd.DataFrame:
  """Computes variances over time for each feature.

  In general, higher variance is better since it creates more signal for the
  regression analysis. However, if the features have not been scaled (divided by
  the mean), then the variance can take any value and this analysis is not
  meaningful.

  Args:
    features: Features for media mix model (media and non-media variables).
    feature_names: Names of media channels to be added to the output dataframe.
    geo_names: Names of geos to be added to the output dataframes.

  Returns:
    Dataframe containing the variance over time for each feature. This dataframe
      contains one row per geo, and just a single row for national data.

  Raises:
    ValueError: If the number of geos in features does not match the number of
    supplied geo_names.
  """
  number_of_geos = core_utils.get_number_geos(features)

  if len(geo_names) != number_of_geos:
    raise ValueError("The number of geos in features does not match the length "
                     "of geo_names")

  variances_as_series = []
  for i_geo in range(number_of_geos):
    features_for_this_geo = features[...,
                                     i_geo] if number_of_geos > 1 else features
    variances_as_series.append(
        pd.DataFrame(data=features_for_this_geo).var(axis=0, ddof=0))

  variances = pd.concat(variances_as_series, axis=1)
  variances.columns = geo_names
  variances.index = copy.copy(feature_names)

  return variances

def _compute_spend_fractions(
    cost_data: jnp.ndarray,
    channel_names: Optional[Sequence[str]] = None,
    output_column_name: str = "fraction of spend") -> pd.DataFrame:
  """Computes fraction of total spend for each media channel.

  Args:
    cost_data: Spend (can be normalized or not) per channel.
    channel_names: Names of media channels to be added to the output dataframe.
    output_column_name: Name of the column in the output dataframe, denoting the
      fraction of the total spend in each media channel.

  Returns:
    Dataframe containing fraction of the total spend in each channel.

  Raises:
    ValueError if any of the costs are zero or negative.
  """
  cost_df = pd.DataFrame(
      cost_data, index=channel_names, columns=[output_column_name])

  if (cost_df[output_column_name] <= 0).any():
    raise ValueError("Values in cost_data must all be positive.")

  normalized_cost_df = cost_df.div(cost_df.sum(axis=0), axis=1).round(4)
  return normalized_cost_df

def _compute_variance_inflation_factors(
    features: jnp.ndarray, feature_names: Sequence[str],
    geo_names: Sequence[str]) -> pd.DataFrame:
  """Computes variance inflation factors for all features.

  Helper function for DataQualityCheck.

  Args:
    features: Features for media mix model (media and non-media variables).
    feature_names: Names of media channels to be added to the output dataframe.
    geo_names: Names of geos to be added to the output dataframes.

  Returns:
    Dataframe containing variance inflation factors for each feature. For
      national-level data the dataframe contains just one column, and for
      geo-level data the list contains one column for each geo.

  Raises:
    ValueError: If the number of geos in features does not match the number of
    supplied geo_names.
  """

  number_of_geos = core_utils.get_number_geos(features)

  if len(geo_names) != number_of_geos:
    raise ValueError("The number of geos in features does not match the length "
                     "of geo_names")

  vifs_for_each_geo = []
  for i_geo in range(number_of_geos):
    features_for_this_geo = features[...,
                                     i_geo] if number_of_geos > 1 else features
    features_for_this_geo = add_constant(
        pd.DataFrame(features_for_this_geo, dtype=float), has_constant="skip")

    vifs_for_this_geo = []
    for i, feature in enumerate(features_for_this_geo.columns):
      if feature != "const":
        vifs_for_this_geo.append(
            variance_inflation_factor(features_for_this_geo.values, i))

    vifs_for_each_geo.append(vifs_for_this_geo)

  vif_df = pd.DataFrame(data=zip(*vifs_for_each_geo), dtype=float)
  vif_df.columns = geo_names
  vif_df.index = copy.copy(feature_names)

  return vif_df


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

"""Set of utilities for LightweighMMM package."""
import pickle
import time
from typing import Any, List, Optional, Tuple

from absl import logging
from jax import random
import jax.numpy as jnp
import numpy as np
import pandas as pd
from scipy import interpolate
from scipy import optimize
from scipy import spatial
from scipy import stats
from tensorflow.io import gfile

from lightweight_mmm import media_transforms

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


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

"""Set of utilities for LightweighMMM package."""
import pickle
import time
from typing import Any, List, Optional, Tuple

from absl import logging
from jax import random
import jax.numpy as jnp
import numpy as np
import pandas as pd
from scipy import interpolate
from scipy import optimize
from scipy import spatial
from scipy import stats
from tensorflow.io import gfile

from lightweight_mmm import media_transforms

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


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

"""Set of utilities for LightweighMMM package."""
import pickle
import time
from typing import Any, List, Optional, Tuple

from absl import logging
from jax import random
import jax.numpy as jnp
import numpy as np
import pandas as pd
from scipy import interpolate
from scipy import optimize
from scipy import spatial
from scipy import stats
from tensorflow.io import gfile

from lightweight_mmm import media_transforms

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


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

"""Set of utilities for LightweighMMM package."""
import pickle
import time
from typing import Any, List, Optional, Tuple

from absl import logging
from jax import random
import jax.numpy as jnp
import numpy as np
import pandas as pd
from scipy import interpolate
from scipy import optimize
from scipy import spatial
from scipy import stats
from tensorflow.io import gfile

from lightweight_mmm import media_transforms

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()


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

"""Set of utilities for LightweighMMM package."""
import pickle
import time
from typing import Any, List, Optional, Tuple

from absl import logging
from jax import random
import jax.numpy as jnp
import numpy as np
import pandas as pd
from scipy import interpolate
from scipy import optimize
from scipy import spatial
from scipy import stats
from tensorflow.io import gfile

from lightweight_mmm import media_transforms

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()