google-meridian 1.0.9__py3-none-any.whl → 1.1.0__py3-none-any.whl

meridian/model/media.py

@@ -34,9 +34,29 @@ __all__ = [
 ]
 
 
+def _roi_calibration_scaled_counterfactual(
+    metric_scaled: tf.Tensor,
+    calibration_period: tf.Tensor,
+) -> tf.Tensor:
+  """Calculate ROI calibration scaled counterfactual media or reach.
+
+  Args:
+    metric_scaled: A tensor of scaled metric values with shape `(n_geos,
+      n_times, n_channels)`.
+    calibration_period: A boolean tensor indicating which time periods are used
+      for calculation.
+
+  Returns:
+    A tensor of scaled metric values with shape `(n_geos, n_times, n_channels)`
+    where media values are set to zero during the calibration period.
+  """
+  factors = tf.where(calibration_period, 0.0, 1.0)
+  return tf.einsum("gtm,tm->gtm", metric_scaled, factors)
+
+
 @dataclasses.dataclass(frozen=True)
 class MediaTensors:
-  """Container for media tensors.
+  """Container for (paid) media tensors.
 
   Attributes:
     media: A tensor constructed from `InputData.media`.
@@ -45,25 +65,31 @@ class MediaTensors:
       model's media data.
     media_scaled: The media tensor normalized by population and by the median
       value.
-    media_counterfactual: A tensor containing the media counterfactual values.
-      If ROI priors are used, then the ROI of media channels is based on the
-      difference in expected sales between the `media` tensor and this
-      `media_counterfactual` tensor.
-    media_counterfactual_scaled: A tensor containing the media counterfactual
-      scaled values.
-    media_spend_counterfactual: A tensor containing the media spend
-      counterfactual values. If ROI priors are used, then the ROI of media
-      channels is based on the spend difference between `media_spend` tensor and
-      this `media_spend_counterfactual` tensor.
+    prior_media_scaled_counterfactual: A tensor containing `media_scaled` values
+      corresponding to the counterfactual scenario required for the prior
+      calculation. For ROI priors, the counterfactual scenario is where media is
+      set to zero during the calibration period. For mROI priors, the
+      counterfactual scenario is where media is increased by a small factor for
+      all `n_media_times`. For contribution priors, the counterfactual scenario
+      is where media is set to zero for all `n_media_times`. This attribute is
+      set to `None` when it would otherwise be a tensor of zeros, i.e., when
+      contribution contribution priors are used, or when ROI priors are used and
+      `roi_calibration_period` is `None`.
+    prior_denominator: If ROI, mROI, or contribution priors are used, this
+      represents the denominator. It is a tensor with dimension equal to
+      `n_media_channels`. For ROI priors, it is the spend during the overlapping
+      time periods between the calibration period and the modeling time window.
+      For mROI priors, it is the ROI prior denominator multiplied by a small
+      factor. For contribution priors, it is the total observed outcome
+      (repeated for each channel.)
   """
 
   media: tf.Tensor | None = None
   media_spend: tf.Tensor | None = None
   media_transformer: transformers.MediaTransformer | None = None
   media_scaled: tf.Tensor | None = None
-  media_counterfactual: tf.Tensor | None = None
-  media_counterfactual_scaled: tf.Tensor | None = None
-  media_spend_counterfactual: tf.Tensor | None = None
+  prior_media_scaled_counterfactual: tf.Tensor | None = None
+  prior_denominator: tf.Tensor | None = None
 
 
 def build_media_tensors(
@@ -81,43 +107,55 @@ def build_media_tensors(
       media, tf.convert_to_tensor(input_data.population, dtype=tf.float32)
   )
   media_scaled = media_transformer.forward(media)
-
-  # Derive counterfactual media tensors depending on whether mroi or roi priors
-  # are used and whether roi_calibration_period is specified.
-  if (
-      model_spec.paid_media_prior_type
-      == constants.PAID_MEDIA_PRIOR_TYPE_MROI
-  ):
-    factor = constants.MROI_FACTOR
+  prior_type = model_spec.effective_media_prior_type
+  calibration_period = model_spec.roi_calibration_period
+  if calibration_period is not None:
+    calibration_period_tf = tf.convert_to_tensor(
+        calibration_period, dtype=tf.bool
+    )
   else:
-    factor = 0
+    calibration_period_tf = None
 
-  if model_spec.roi_calibration_period is None:
-    media_counterfactual = factor * media
-    media_counterfactual_scaled = factor * media_scaled
-    media_spend_counterfactual = factor * media_spend
+  aggregated_media_spend = tf.convert_to_tensor(
+      input_data.aggregate_media_spend(
+          calibration_period=calibration_period
+      ),
+      dtype=tf.float32,
+  )
+  # Set `prior_media_scaled_counterfactual` and `prior_denominator` depending on
+  # the prior type.
+  if prior_type == constants.TREATMENT_PRIOR_TYPE_ROI:
+    prior_denominator = aggregated_media_spend
+
+    if calibration_period is None:
+      prior_media_scaled_counterfactual = None
+    else:
+      prior_media_scaled_counterfactual = (
+          _roi_calibration_scaled_counterfactual(
+              media_scaled,
+              calibration_period=calibration_period_tf,
+          )
+      )
+  elif prior_type == constants.TREATMENT_PRIOR_TYPE_MROI:
+    prior_media_scaled_counterfactual = media_scaled * constants.MROI_FACTOR
+    prior_denominator = aggregated_media_spend * (constants.MROI_FACTOR - 1.0)
+  elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
+    prior_media_scaled_counterfactual = None
+    total_outcome = tf.cast(input_data.get_total_outcome(), tf.float32)
+    prior_denominator = tf.repeat(total_outcome, len(input_data.media_channel))
+  elif prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
+    prior_media_scaled_counterfactual = None
+    prior_denominator = None
   else:
-    media_counterfactual = tf.where(
-        model_spec.roi_calibration_period, factor * media, media
-    )
-    media_counterfactual_scaled = tf.where(
-        model_spec.roi_calibration_period, factor * media_scaled, media_scaled
-    )
-    n_times = len(input_data.time)
-    media_spend_counterfactual = tf.where(
-        model_spec.roi_calibration_period[..., -n_times:, :],
-        factor * media_spend,
-        media_spend,
-    )
+    raise ValueError(f"Unsupported prior type: {prior_type}")
 
   return MediaTensors(
       media=media,
       media_spend=media_spend,
       media_transformer=media_transformer,
       media_scaled=media_scaled,
-      media_counterfactual=media_counterfactual,
-      media_counterfactual_scaled=media_counterfactual_scaled,
-      media_spend_counterfactual=media_spend_counterfactual,
+      prior_media_scaled_counterfactual=prior_media_scaled_counterfactual,
+      prior_denominator=prior_denominator,
   )
 
 
@@ -131,17 +169,11 @@ class OrganicMediaTensors:
       the model's organic media data.
     organic_media_scaled: The organic media tensor normalized by population and
       by the median value.
-    organic_media_counterfactual: A tensor containing the organic media
-      counterfactual values.
-    organic_media_counterfactual_scaled: A tensor containing the organic media
-      counterfactual scaled values.
   """
 
   organic_media: tf.Tensor | None = None
   organic_media_transformer: transformers.MediaTransformer | None = None
   organic_media_scaled: tf.Tensor | None = None
-  organic_media_counterfactual: tf.Tensor | None = None
-  organic_media_counterfactual_scaled: tf.Tensor | None = None
 
 
 def build_organic_media_tensors(
@@ -160,15 +192,11 @@ def build_organic_media_tensors(
       tf.convert_to_tensor(input_data.population, dtype=tf.float32),
   )
   organic_media_scaled = organic_media_transformer.forward(organic_media)
-  organic_media_counterfactual = tf.zeros_like(organic_media)
-  organic_media_counterfactual_scaled = tf.zeros_like(organic_media_scaled)
 
   return OrganicMediaTensors(
       organic_media=organic_media,
       organic_media_transformer=organic_media_transformer,
       organic_media_scaled=organic_media_scaled,
-      organic_media_counterfactual=organic_media_counterfactual,
-      organic_media_counterfactual_scaled=organic_media_counterfactual_scaled,
   )
 
 
@@ -184,16 +212,23 @@ class RfTensors:
       model's RF data.
     reach_scaled: A reach tensor normalized by population and by the median
       value.
-    reach_counterfactual: A reach tensor with media counterfactual values. If
-      ROI priors are used, then the ROI of R&F channels is based on the
-      difference in expected sales between the `reach` tensor and this
-      `reach_counterfactual` tensor.
-    reach_counterfactual_scaled: A reach tensor with media counterfactual scaled
-      values.
-    rf_spend_counterfactual: A reach tensor with media spend counterfactual
-      values. If ROI priors are used, then the ROI of R&F channels is based on
-      the spend difference between `rf_spend` tensor and this
-      `rf_spend_counterfactual` tensor.
+    prior_reach_scaled_counterfactual: A tensor containing `reach_scaled` values
+      corresponding to the counterfactual scenario required for the prior
+      calculation. For ROI priors, the counterfactual scenario is where reach is
+      set to zero during the calibration period. For mROI priors, the
+      counterfactual scenario is where reach is increased by a small factor for
+      all `n_rf_times`. For contribution priors, the counterfactual scenario is
+      where reach is set to zero for all `n_rf_times`. This attribute is set to
+      `None` when it would otherwise be a tensor of zeros, i.e., when
+      contribution contribution priors are used, or when ROI priors are used and
+      `rf_roi_calibration_period` is `None`.
+    prior_denominator: If ROI, mROI, or contribution priors are used, this
+      represents the denominator. It is a tensor with dimension equal to
+      `n_rf_channels`. For ROI priors, it is the spend during the overlapping
+      time periods between the calibration period and the modeling time window.
+      For mROI priors, it is the ROI prior denominator multiplied by a small
+      factor. For contribution priors, it is the total observed outcome
+      (repeated for each channel).
   """
 
   reach: tf.Tensor | None = None
@@ -201,9 +236,8 @@ class RfTensors:
   rf_spend: tf.Tensor | None = None
   reach_transformer: transformers.MediaTransformer | None = None
   reach_scaled: tf.Tensor | None = None
-  reach_counterfactual: tf.Tensor | None = None
-  reach_counterfactual_scaled: tf.Tensor | None = None
-  rf_spend_counterfactual: tf.Tensor | None = None
+  prior_reach_scaled_counterfactual: tf.Tensor | None = None
+  prior_denominator: tf.Tensor | None = None
 
 
 def build_rf_tensors(
@@ -221,27 +255,39 @@ def build_rf_tensors(
       reach, tf.convert_to_tensor(input_data.population, dtype=tf.float32)
   )
   reach_scaled = reach_transformer.forward(reach)
-
-  # marginal ROI by reach equals the ROI. The conversion between `beta_rf` and
-  # `roi_rf` is the same, regardless of whether `roi_rf` represents ROI or
-  # marginal ROI by reach.
-  if model_spec.rf_roi_calibration_period is None:
-    reach_counterfactual = tf.zeros_like(reach)
-    reach_counterfactual_scaled = tf.zeros_like(reach_scaled)
-    rf_spend_counterfactual = tf.zeros_like(rf_spend)
+  prior_type = model_spec.effective_rf_prior_type
+  calibration_period = model_spec.rf_roi_calibration_period
+  if calibration_period is not None:
+    calibration_period = tf.convert_to_tensor(calibration_period, dtype=tf.bool)
+  aggregated_rf_spend = tf.convert_to_tensor(
+      input_data.aggregate_rf_spend(calibration_period=calibration_period),
+      dtype=tf.float32,
+  )
+  # Set `prior_reach_scaled_counterfactual` and `prior_denominator` depending on
+  # the prior type.
+  if prior_type == constants.TREATMENT_PRIOR_TYPE_ROI:
+    prior_denominator = aggregated_rf_spend
+    if calibration_period is None:
+      prior_reach_scaled_counterfactual = None
+    else:
+      prior_reach_scaled_counterfactual = (
+          _roi_calibration_scaled_counterfactual(
+              reach_scaled,
+              calibration_period=calibration_period,
+          )
+      )
+  elif prior_type == constants.TREATMENT_PRIOR_TYPE_MROI:
+    prior_reach_scaled_counterfactual = reach_scaled * constants.MROI_FACTOR
+    prior_denominator = aggregated_rf_spend * (constants.MROI_FACTOR - 1.0)
+  elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
+    prior_reach_scaled_counterfactual = None
+    total_outcome = tf.cast(input_data.get_total_outcome(), tf.float32)
+    prior_denominator = tf.repeat(total_outcome, len(input_data.rf_channel))
+  elif prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
+    prior_reach_scaled_counterfactual = None
+    prior_denominator = None
   else:
-    reach_counterfactual = tf.where(
-        model_spec.rf_roi_calibration_period, 0, reach
-    )
-    reach_counterfactual_scaled = tf.where(
-        model_spec.rf_roi_calibration_period, 0, reach_scaled
-    )
-    n_times = len(input_data.time)
-    rf_spend_counterfactual = tf.where(
-        model_spec.rf_roi_calibration_period[..., -n_times:, :],
-        0,
-        rf_spend,
-    )
+    raise ValueError(f"Unsupported prior type: {prior_type}")
 
   return RfTensors(
       reach=reach,
@@ -249,9 +295,8 @@ def build_rf_tensors(
       rf_spend=rf_spend,
       reach_transformer=reach_transformer,
       reach_scaled=reach_scaled,
-      reach_counterfactual=reach_counterfactual,
-      reach_counterfactual_scaled=reach_counterfactual_scaled,
-      rf_spend_counterfactual=rf_spend_counterfactual,
+      prior_reach_scaled_counterfactual=prior_reach_scaled_counterfactual,
+      prior_denominator=prior_denominator,
   )
 
 
@@ -266,18 +311,12 @@ class OrganicRfTensors:
       using the model's organic RF data.
     organic_reach_scaled: An organic reach tensor normalized by population and
       by the median value.
-    organic_reach_counterfactual: An organic reach tensor with media
-      counterfactual values.
-    organic_reach_counterfactual_scaled: An organic reach tensor with media
-      counterfactual scaled values.
   """
 
   organic_reach: tf.Tensor | None = None
   organic_frequency: tf.Tensor | None = None
   organic_reach_transformer: transformers.MediaTransformer | None = None
   organic_reach_scaled: tf.Tensor | None = None
-  organic_reach_counterfactual: tf.Tensor | None = None
-  organic_reach_counterfactual_scaled: tf.Tensor | None = None
 
 
 def build_organic_rf_tensors(
@@ -298,14 +337,10 @@ def build_organic_rf_tensors(
       tf.convert_to_tensor(input_data.population, dtype=tf.float32),
   )
   organic_reach_scaled = organic_reach_transformer.forward(organic_reach)
-  organic_reach_counterfactual = tf.zeros_like(organic_reach)
-  organic_reach_counterfactual_scaled = tf.zeros_like(organic_reach_scaled)
 
   return OrganicRfTensors(
       organic_reach=organic_reach,
       organic_frequency=organic_frequency,
       organic_reach_transformer=organic_reach_transformer,
       organic_reach_scaled=organic_reach_scaled,
-      organic_reach_counterfactual=organic_reach_counterfactual,
-      organic_reach_counterfactual_scaled=organic_reach_counterfactual_scaled,
   )