google-meridian 1.3.1__py3-none-any.whl → 1.4.0__py3-none-any.whl

meridian/model/equations.py

@@ -0,0 +1,418 @@
+# Copyright 2025 The Meridian Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Core mathematical equations for the Meridian model.
+
+This module defines the `ModelEquations` class, which encapsulates the stateless
+mathematical functions used in the Meridian MMM. This includes the core model
+definitions, such as adstock, hill, and other transformations used
+during model fitting. It requires a `ModelContext` instance for data access.
+"""
+
+from collections.abc import Sequence
+import numbers
+
+from meridian import backend
+from meridian import constants
+from meridian.model import adstock_hill
+from meridian.model import context
+
+
+__all__ = [
+    "ModelEquations",
+]
+
+
+class ModelEquations:
+  """Provides core, stateless mathematical functions for Meridian MMM."""
+
+  def __init__(self, model_context: context.ModelContext):
+    self._context = model_context
+
+  def adstock_hill_media(
+      self,
+      *,
+      media: backend.Tensor,
+      alpha: backend.Tensor,
+      ec: backend.Tensor,
+      slope: backend.Tensor,
+      decay_functions: str | Sequence[str] = constants.GEOMETRIC_DECAY,
+      n_times_output: int | None = None,
+  ) -> backend.Tensor:
+    """Transforms media or using Adstock and Hill functions in the desired order.
+
+    Args:
+      media: Tensor of dimensions `(n_geos, n_media_times, n_media_channels)`
+        containing non-negative media execution values. Typically this is
+        impressions, but it can be any metric, such as `media_spend`. Clicks are
+        often used for paid search ads.
+      alpha: Uniform distribution for Adstock and Hill calculations.
+      ec: Shifted half-normal distribution for Adstock and Hill calculations.
+      slope: Deterministic distribution for Adstock and Hill calculations.
+      decay_functions: String or sequence of strings denoting the adstock decay
+        function(s) for each channel. Default: 'geometric'.
+      n_times_output: Number of time periods to output. This argument is
+        optional when the number of time periods in `media` equals
+        `n_media_times`, in which case `n_times_output` defaults to `n_times`.
+
+    Returns:
+      Tensor with dimensions `[..., n_geos, n_times, n_media_channels]`
+      representing Adstock and Hill-transformed media.
+    """
+    if n_times_output is None and (
+        media.shape[1] == self._context.n_media_times
+    ):
+      n_times_output = self._context.n_times
+    elif n_times_output is None:
+      raise ValueError(
+          "n_times_output is required. This argument is only optional when "
+          "`media` has a number of time periods equal to `n_media_times`."
+      )
+
+    adstock_transformer = adstock_hill.AdstockTransformer(
+        alpha=alpha,
+        max_lag=self._context.model_spec.max_lag,
+        n_times_output=n_times_output,
+        decay_functions=decay_functions,
+    )
+    hill_transformer = adstock_hill.HillTransformer(
+        ec=ec,
+        slope=slope,
+    )
+    transformers_list = (
+        [hill_transformer, adstock_transformer]
+        if self._context.model_spec.hill_before_adstock
+        else [adstock_transformer, hill_transformer]
+    )
+
+    media_out = media
+    for transformer in transformers_list:
+      media_out = transformer.forward(media_out)
+    return media_out
+
+  def adstock_hill_rf(
+      self,
+      *,
+      reach: backend.Tensor,
+      frequency: backend.Tensor,
+      alpha: backend.Tensor,
+      ec: backend.Tensor,
+      slope: backend.Tensor,
+      decay_functions: str | Sequence[str] = constants.GEOMETRIC_DECAY,
+      n_times_output: int | None = None,
+  ) -> backend.Tensor:
+    """Transforms reach and frequency (RF) using Hill and Adstock functions.
+
+    Args:
+      reach: Tensor of dimensions `(n_geos, n_media_times, n_rf_channels)`
+        containing non-negative media for reach.
+      frequency: Tensor of dimensions `(n_geos, n_media_times, n_rf_channels)`
+        containing non-negative media for frequency.
+      alpha: Uniform distribution for Adstock and Hill calculations.
+      ec: Shifted half-normal distribution for Adstock and Hill calculations.
+      slope: Deterministic distribution for Adstock and Hill calculations.
+      decay_functions: String or sequence of strings denoting the adstock decay
+        function(s) for each channel. Default: 'geometric'.
+      n_times_output: Number of time periods to output. This argument is
+        optional when the number of time periods in `reach` equals
+        `n_media_times`, in which case `n_times_output` defaults to `n_times`.
+
+    Returns:
+      Tensor with dimensions `[..., n_geos, n_times, n_rf_channels]`
+      representing Hill and Adstock-transformed RF.
+    """
+    if n_times_output is None and (
+        reach.shape[1] == self._context.n_media_times
+    ):
+      n_times_output = self._context.n_times
+    elif n_times_output is None:
+      raise ValueError(
+          "n_times_output is required. This argument is only optional when "
+          "`reach` has a number of time periods equal to `n_media_times`."
+      )
+
+    hill_transformer = adstock_hill.HillTransformer(
+        ec=ec,
+        slope=slope,
+    )
+    adstock_transformer = adstock_hill.AdstockTransformer(
+        alpha=alpha,
+        max_lag=self._context.model_spec.max_lag,
+        n_times_output=n_times_output,
+        decay_functions=decay_functions,
+    )
+    adj_frequency = hill_transformer.forward(frequency)
+    rf_out = adstock_transformer.forward(reach * adj_frequency)
+
+    return rf_out
+
+  def compute_non_media_treatments_baseline(
+      self,
+      non_media_baseline_values: Sequence[str | float] | None = None,
+  ) -> backend.Tensor:
+    """Computes the baseline for each non-media treatment channel.
+
+    Args:
+      non_media_baseline_values: Optional list of shape
+        `(n_non_media_channels,)`. Each element is either a float (which means
+        that the fixed value will be used as baseline for the given channel) or
+        one of the strings "min" or "max" (which mean that the global minimum or
+        maximum value will be used as baseline for the values of the given
+        non_media treatment channel). If float values are provided, it is
+        expected that they are scaled by population for the channels where
+        `model_spec.non_media_population_scaling_id` is `True`. If `None`, the
+        `model_spec.non_media_baseline_values` is used, which defaults to the
+        minimum value for each non_media treatment channel.
+
+    Returns:
+      A tensor of shape `(n_non_media_channels,)` containing the
+      baseline values for each non-media treatment channel.
+    """
+    if non_media_baseline_values is None:
+      non_media_baseline_values = (
+          self._context.model_spec.non_media_baseline_values
+      )
+
+    no_op_scaling_factor = backend.ones_like(self._context.population)[
+        :, backend.newaxis, backend.newaxis
+    ]
+    if self._context.model_spec.non_media_population_scaling_id is not None:
+      scaling_factors = backend.where(
+          self._context.model_spec.non_media_population_scaling_id,
+          self._context.population[:, backend.newaxis, backend.newaxis],
+          no_op_scaling_factor,
+      )
+    else:
+      scaling_factors = no_op_scaling_factor
+
+    non_media_treatments_population_scaled = backend.divide_no_nan(
+        self._context.non_media_treatments, scaling_factors
+    )
+
+    if non_media_baseline_values is None:
+      # If non_media_baseline_values is not provided, use the minimum
+      # value for each non_media treatment channel as the baseline.
+      non_media_baseline_values_filled = [
+          constants.NON_MEDIA_BASELINE_MIN
+      ] * non_media_treatments_population_scaled.shape[-1]
+    else:
+      non_media_baseline_values_filled = non_media_baseline_values
+
+    if non_media_treatments_population_scaled.shape[-1] != len(
+        non_media_baseline_values_filled
+    ):
+      raise ValueError(
+          "The number of non-media channels"
+          f" ({non_media_treatments_population_scaled.shape[-1]}) does not"
+          " match the number of baseline values"
+          f" ({len(non_media_baseline_values_filled)})."
+      )
+
+    baseline_list = []
+    for channel in range(non_media_treatments_population_scaled.shape[-1]):
+      baseline_value = non_media_baseline_values_filled[channel]
+
+      if baseline_value == constants.NON_MEDIA_BASELINE_MIN:
+        baseline_for_channel = backend.reduce_min(
+            non_media_treatments_population_scaled[..., channel], axis=[0, 1]
+        )
+      elif baseline_value == constants.NON_MEDIA_BASELINE_MAX:
+        baseline_for_channel = backend.reduce_max(
+            non_media_treatments_population_scaled[..., channel], axis=[0, 1]
+        )
+      elif isinstance(baseline_value, numbers.Number):
+        baseline_for_channel = backend.to_tensor(
+            baseline_value, dtype=backend.float32
+        )
+      else:
+        raise ValueError(
+            f"Invalid non_media_baseline_values value: '{baseline_value}'. Only"
+            " float numbers and strings 'min' and 'max' are supported."
+        )
+
+      baseline_list.append(baseline_for_channel)
+
+    return backend.stack(baseline_list, axis=-1)
+
+  def linear_predictor_counterfactual_difference_media(
+      self,
+      *,
+      media_transformed: backend.Tensor,
+      alpha_m: backend.Tensor,
+      ec_m: backend.Tensor,
+      slope_m: backend.Tensor,
+  ) -> backend.Tensor:
+    """Calculates linear predictor counterfactual difference for non-RF media.
+
+    For non-RF media variables (paid or organic), this function calculates the
+    linear predictor difference between the treatment variable and its
+    counterfactual. "Linear predictor" refers to the output of the hill/adstock
+    function, which is multiplied by the geo-level coefficient.
+
+    This function does the calculation efficiently by only calculating calling
+    the hill/adstock function if the prior counterfactual is not all zeros.
+
+    Args:
+      media_transformed: The output of the hill/adstock function for actual
+        historical media data.
+      alpha_m: The adstock alpha parameter values.
+      ec_m: The adstock ec parameter values.
+      slope_m: The adstock hill slope parameter values.
+
+    Returns:
+      The linear predictor difference between the treatment variable and its
+      counterfactual.
+    """
+    if self._context.media_tensors.prior_media_scaled_counterfactual is None:
+      return media_transformed
+    media_transformed_counterfactual = self.adstock_hill_media(
+        media=self._context.media_tensors.prior_media_scaled_counterfactual,
+        alpha=alpha_m,
+        ec=ec_m,
+        slope=slope_m,
+        decay_functions=self._context.adstock_decay_spec.media,
+    )
+    # Absolute values is needed because the difference is negative for mROI
+    # priors and positive for ROI and contribution priors.
+    return backend.absolute(
+        media_transformed - media_transformed_counterfactual
+    )
+
+  def linear_predictor_counterfactual_difference_rf(
+      self,
+      *,
+      rf_transformed: backend.Tensor,
+      alpha_rf: backend.Tensor,
+      ec_rf: backend.Tensor,
+      slope_rf: backend.Tensor,
+  ) -> backend.Tensor:
+    """Calculates linear predictor counterfactual difference for RF media.
+
+    For RF media variables (paid or organic), this function calculates the
+    linear predictor difference between the treatment variable and its
+    counterfactual. "Linear predictor" refers to the output of the hill/adstock
+    function, which is multiplied by the geo-level coefficient.
+
+    This function does the calculation efficiently by only calculating calling
+    the hill/adstock function if the prior counterfactual is not all zeros.
+
+    Args:
+      rf_transformed: The output of the hill/adstock function for actual
+        historical media data.
+      alpha_rf: The adstock alpha parameter values.
+      ec_rf: The adstock ec parameter values.
+      slope_rf: The adstock hill slope parameter values.
+
+    Returns:
+      The linear predictor difference between the treatment variable and its
+      counterfactual.
+    """
+    if self._context.rf_tensors.prior_reach_scaled_counterfactual is None:
+      return rf_transformed
+    rf_transformed_counterfactual = self.adstock_hill_rf(
+        reach=self._context.rf_tensors.prior_reach_scaled_counterfactual,
+        frequency=self._context.rf_tensors.frequency,
+        alpha=alpha_rf,
+        ec=ec_rf,
+        slope=slope_rf,
+        decay_functions=self._context.adstock_decay_spec.rf,
+    )
+    # Absolute values is needed because the difference is negative for mROI
+    # priors and positive for ROI and contribution priors.
+    return backend.absolute(rf_transformed - rf_transformed_counterfactual)
+
+  def calculate_beta_x(
+      self,
+      *,
+      is_non_media: bool,
+      incremental_outcome_x: backend.Tensor,
+      linear_predictor_counterfactual_difference: backend.Tensor,
+      eta_x: backend.Tensor,
+      beta_gx_dev: backend.Tensor,
+  ) -> backend.Tensor:
+    """Calculates coefficient mean parameter for any treatment variable type.
+
+    The "beta_x" in the function name refers to the coefficient mean parameter
+    of any treatment variable. The "x" can represent "m", "rf", "om", or "orf".
+    This function can also be used to calculate "gamma_n" for any non-media
+    treatments.
+
+    Args:
+      is_non_media: Boolean indicating whether the treatment variable is a
+        non-media treatment. This argument is used to determine whether the
+        coefficient random effects are normal or log-normal. If `True`, then
+        random effects are assumed to be normal. Otherwise, the distribution is
+        inferred from `self._context.media_effects_dist`.
+      incremental_outcome_x: The incremental outcome of the treatment variable,
+        which depends on the parameter values of a particular prior or posterior
+        draw. The "_x" indicates that this is a tensor with length equal to the
+        dimension of the treatment variable.
+      linear_predictor_counterfactual_difference: The difference between the
+        treatment variable and its counterfactual on the linear predictor scale.
+        "Linear predictor" refers to the quantity that is multiplied by the
+        geo-level coefficient. For media variables, this is the output of the
+        hill/adstock transformation function. For non-media treatments, this is
+        simply the treatment variable after centering/scaling transformations.
+        This tensor has dimensions for geo, time, and channel.
+      eta_x: The random effect standard deviation parameter values. For media
+        variables, the "x" represents "m", "rf", "om", or "orf". For non-media
+        treatments, this argument should be set to `xi_n`, which is analogous to
+        "eta".
+      beta_gx_dev: The latent standard normal parameter values of the geo-level
+        coefficients. For media variables, the "x" represents "m", "rf", "om",
+        or "orf". For non-media treatments, this argument should be set to
+        `gamma_gn_dev`, which is analogous to "beta_gx_dev".
+
+    Returns:
+      The coefficient mean parameter of the treatment variable, which has
+      dimension equal to the number of treatment channels..
+    """
+    if is_non_media:
+      random_effects_normal = True
+    else:
+      random_effects_normal = (
+          self._context.media_effects_dist == constants.MEDIA_EFFECTS_NORMAL
+      )
+    if self._context.revenue_per_kpi is None:
+      revenue_per_kpi = backend.ones(
+          [self._context.n_geos, self._context.n_times], dtype=backend.float32
+      )
+    else:
+      revenue_per_kpi = self._context.revenue_per_kpi
+    incremental_outcome_gx_over_beta_gx = backend.einsum(
+        "...gtx,gt,g,->...gx",
+        linear_predictor_counterfactual_difference,
+        revenue_per_kpi,
+        self._context.population,
+        self._context.kpi_transformer.population_scaled_stdev,
+    )
+    if random_effects_normal:
+      numerator_term_x = backend.einsum(
+          "...gx,...gx,...x->...x",
+          incremental_outcome_gx_over_beta_gx,
+          beta_gx_dev,
+          eta_x,
+      )
+      denominator_term_x = backend.einsum(
+          "...gx->...x", incremental_outcome_gx_over_beta_gx
+      )
+      return (incremental_outcome_x - numerator_term_x) / denominator_term_x
+    # For log-normal random effects, beta_x and eta_x are not mean & std.
+    # The parameterization is beta_gx ~ exp(beta_x + eta_x * N(0, 1)).
+    denominator_term_x = backend.einsum(
+        "...gx,...gx->...x",
+        incremental_outcome_gx_over_beta_gx,
+        backend.exp(beta_gx_dev * eta_x[..., backend.newaxis, :]),
+    )
+    return backend.log(incremental_outcome_x) - backend.log(denominator_term_x)

meridian/model/knots.py

@@ -14,17 +14,17 @@
 
 """Auxiliary functions for knots calculations."""
 
-import bisect
 from collections.abc import Collection, Sequence
 import copy
 import dataclasses
 import math
+import pprint
 from typing import Any
+
 from meridian import constants
 from meridian.data import input_data
 import numpy as np
-# TODO: b/437393442 - migrate patsy
-from patsy import highlevel
+from scipy import interpolate
 from statsmodels.regression import linear_model
 
 
@@ -35,40 +35,35 @@ __all__ = [
 ]
 
 
-# TODO: Reimplement with a more readable method.
-def _find_neighboring_knots_indices(
+def _find_left_knot_indices(
+    *,
     times: np.ndarray,
     knot_locations: np.ndarray,
-) -> Sequence[Sequence[int] | None]:
-  """Return indices of neighboring knot locations.
-
-  Returns indices in `knot_locations` that correspond to the neighboring knot
-  locations for each time period. If a time point is at or before the first
-  knot, the first knot is the only neighboring knot. If a time point is after
-  the last knot, the last knot is the only neighboring knot.
+) -> Sequence[int]:
+  """Return the index of the left neighboring knot for each time point.
 
   Args:
     times: Times `0, 1, 2,..., (n_times-1)`.
     knot_locations: The location of knots within `0, 1, 2,..., (n_times-1)`.
 
   Returns:
-    List of length `n_times`. Each element is the indices of the neighboring
-    knot locations for the respective time period. If a time point is at or
-    before the first knot, the first knot is the only neighboring knot. If a
-    time point is after the last knot, the last knot is the only neighboring
-    knot.
+    A list of indices of the left neighboring knot for each time point. The
+    length of the list is equal to the length of `times`.
+    - If a time point is at or before the first knot, the index is 0.
+    - If a time point is at or after the last knot, the index is `n_knots - 1`.
+    - Otherwise, it's the index of the knot just to the left.
   """
-  neighboring_knots_indices = [None] * len(times)
-  for t in times:
-    # knot_locations assumed to be sorted.
-    if t <= knot_locations[0]:
-      neighboring_knots_indices[t] = [0]
-    elif t >= knot_locations[-1]:
-      neighboring_knots_indices[t] = [len(knot_locations) - 1]
-    else:
-      bisect_index = bisect.bisect_left(knot_locations, t)
-      neighboring_knots_indices[t] = [bisect_index - 1, bisect_index]
-  return neighboring_knots_indices
+  n_knots = len(knot_locations)
+  # Find indices such that knot_locations[i-1] < times <= knot_locations[i]
+  insert_indices = np.searchsorted(knot_locations, times, side='right')
+  left_knot_indices = insert_indices - 1
+
+  # Handle edge cases for times before the first knot
+  left_knot_indices[times < knot_locations[0]] = 0
+  # Handle edge cases for times at or after the last knot
+  left_knot_indices[times >= knot_locations[-1]] = n_knots - 1
+
+  return left_knot_indices
 
 
 def l1_distance_weights(
@@ -114,16 +109,31 @@ def l1_distance_weights(
   time_minus_knot = abs(knot_locations[:, np.newaxis] - times[np.newaxis, :])
 
   w = np.zeros(time_minus_knot.shape, dtype=np.float32)
-  neighboring_knots_indices = _find_neighboring_knots_indices(
-      times, knot_locations
+  left_knot_indices = _find_left_knot_indices(
+      times=times, knot_locations=knot_locations
   )
+
   for t in times:
-    idx = neighboring_knots_indices[t]
-    if len(idx) == 1:
-      w[idx, t] = 1
+    left_idx = left_knot_indices[t]
+    current_time = times[t]
+
+    if current_time in knot_locations:
+      # If time is exactly at a knot, give all weight to that knot.
+      knot_idx = np.where(knot_locations == current_time)[0][0]
+      w[knot_idx, t] = 1.0
+    elif current_time < knot_locations[0] or current_time > knot_locations[-1]:
+      # Outside the knot range, assign full weight to the closest endpoint knot.
+      w[left_idx, t] = 1.0
     else:
-      # Weight is in proportion to how close the two neighboring knots are.
-      w[idx, t] = 1 - (time_minus_knot[idx, t] / time_minus_knot[idx, t].sum())
+      # Time is between left_idx and left_idx + 1.
+      left_dist = time_minus_knot[left_idx, t]
+      right_dist = time_minus_knot[left_idx + 1, t]
+      total_dist = left_dist + right_dist
+
+      # Assign weight inversely proportional to distance.
+      # The closer knot gets more weight.
+      w[left_idx, t] = right_dist / total_dist
+      w[left_idx + 1, t] = left_dist / total_dist
 
   return w
 
@@ -289,6 +299,22 @@ class AKS:
     penalty = geo_scaling_factor * base_penalty
 
     aspline = self.aspline(x=x, y=y, knots=knots, penalty=penalty)
+    # Ensure defined knot range covers at least one of the available knot sets.
+    available_knots_lengths = np.unique(
+        np.fromiter(
+            (len(x) for x in aspline[constants.KNOTS_SELECTED]), dtype=int
+        )
+    ).tolist()
+    if not any(
+        min_internal_knots <= k <= max_internal_knots
+        for k in available_knots_lengths
+    ):
+      raise ValueError(
+          f'The range [{min_internal_knots}, {max_internal_knots}] does not'
+          ' contain any of the available knot lengths:'
+          f' {pprint.pformat(available_knots_lengths)}'
+      )
+
     n_knots = np.array([len(x) for x in aspline[constants.KNOTS_SELECTED]])
     feasible_idx = np.where(
         (n_knots >= min_internal_knots) & (n_knots <= max_internal_knots)
@@ -302,6 +328,17 @@ class AKS:
 
     return AKSResult(knots_sel[opt_idx], model[opt_idx])
 
+  def _get_bspline_matrix(self, x, knots):
+    """Replaces patsy.highlevel.dmatrix('bs(...)', ...)"""
+    # Pad knots at boundaries to match patsy's 'bs()' behavior
+    t = np.concatenate((
+        [x.min()] * (self._DEGREE + 1),
+        np.sort(knots),
+        [x.max()] * (self._DEGREE + 1),
+    ))
+    # Create design matrix (standard numpy array)
+    return interpolate.BSpline.design_matrix(x, t, self._DEGREE).toarray()
+
   def _calculate_initial_knots(
       self,
       x: np.ndarray,
@@ -416,14 +453,7 @@ class AKS:
           'Provided x and y args for aspline must both be 1 dimensional!'
       )
 
-    bs_cmd = (
-        'bs(x,knots=['
-        + ','.join(map(str, knots))
-        + '],degree='
-        + str(self._DEGREE)
-        + ',include_intercept=True)-1'
-    )
-    xmat = highlevel.dmatrix(bs_cmd, {'x': x})
+    xmat = self._get_bspline_matrix(x, knots)
     nrow = xmat.shape[0]
     ncol = xmat.shape[1]
 
@@ -455,10 +485,8 @@ class AKS:
       if converge:
         sel_ls[index_penalty] = sel
         knots_sel[index_penalty] = knots[sel > 0.99]
-        bs_cmd_iter = (
-            f"bs(x,knots=[{','.join(map(str, knots_sel[index_penalty]))}],degree={self._DEGREE},include_intercept=True)-1"
-        )
-        design_mat = highlevel.dmatrix(bs_cmd_iter, {'x': x})
+
+        design_mat = self._get_bspline_matrix(x, knots_sel[index_penalty])
         x_sel[index_penalty] = design_mat
         bs_model = linear_model.OLS(y, x_sel[index_penalty]).fit()
         model[index_penalty] = bs_model