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

meridian/model/model_test_data.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.
@@ -70,6 +70,10 @@ class WithInputDataSamples:
       _TEST_DIR,
       "sample_prior_media_only.nc",
   )
+  _TEST_SAMPLE_PRIOR_MEDIA_ONLY_NO_CONTROLS_PATH = os.path.join(
+      _TEST_DIR,
+      "sample_prior_media_only_no_controls.nc",
+  )
   _TEST_SAMPLE_PRIOR_RF_ONLY_PATH = os.path.join(
       _TEST_DIR,
       "sample_prior_rf_only.nc",
@@ -82,6 +86,10 @@ class WithInputDataSamples:
       _TEST_DIR,
       "sample_posterior_media_only.nc",
   )
+  _TEST_SAMPLE_POSTERIOR_MEDIA_ONLY_NO_CONTROLS_PATH = os.path.join(
+      _TEST_DIR,
+      "sample_posterior_media_only_no_controls.nc",
+  )
   _TEST_SAMPLE_POSTERIOR_RF_ONLY_PATH = os.path.join(
       _TEST_DIR,
       "sample_posterior_rf_only.nc",
@@ -130,6 +138,17 @@ class WithInputDataSamples:
             seed=0,
         )
     )
+    self.input_data_media_and_rf_non_revenue_no_revenue_per_kpi = (
+        test_utils.sample_input_data_non_revenue_no_revenue_per_kpi(
+            n_geos=self._N_GEOS,
+            n_times=self._N_TIMES,
+            n_media_times=self._N_MEDIA_TIMES,
+            n_controls=self._N_CONTROLS,
+            n_media_channels=self._N_MEDIA_CHANNELS,
+            n_rf_channels=self._N_RF_CHANNELS,
+            seed=0,
+        )
+    )
     self.input_data_with_media_only = (
         test_utils.sample_input_data_non_revenue_revenue_per_kpi(
             n_geos=self._N_GEOS,
@@ -161,6 +180,17 @@ class WithInputDataSamples:
             seed=0,
         )
     )
+    self.input_data_with_media_and_rf_no_controls = (
+        test_utils.sample_input_data_non_revenue_revenue_per_kpi(
+            n_geos=self._N_GEOS,
+            n_times=self._N_TIMES,
+            n_media_times=self._N_MEDIA_TIMES,
+            n_controls=None,
+            n_media_channels=self._N_MEDIA_CHANNELS,
+            n_rf_channels=self._N_RF_CHANNELS,
+            seed=0,
+        )
+    )
     self.short_input_data_with_media_only = (
         test_utils.sample_input_data_non_revenue_revenue_per_kpi(
             n_geos=self._N_GEOS,
@@ -171,6 +201,16 @@ class WithInputDataSamples:
             seed=0,
         )
     )
+    self.short_input_data_with_media_only_no_controls = (
+        test_utils.sample_input_data_non_revenue_revenue_per_kpi(
+            n_geos=self._N_GEOS,
+            n_times=self._N_TIMES_SHORT,
+            n_media_times=self._N_MEDIA_TIMES_SHORT,
+            n_controls=0,
+            n_media_channels=self._N_MEDIA_CHANNELS,
+            seed=0,
+        )
+    )
     self.short_input_data_with_rf_only = (
         test_utils.sample_input_data_non_revenue_revenue_per_kpi(
             n_geos=self._N_GEOS,
@@ -220,6 +260,9 @@ class WithInputDataSamples:
     test_prior_media_only = xr.open_dataset(
         self._TEST_SAMPLE_PRIOR_MEDIA_ONLY_PATH
     )
+    test_prior_media_only_no_controls = xr.open_dataset(
+        self._TEST_SAMPLE_PRIOR_MEDIA_ONLY_NO_CONTROLS_PATH
+    )
     test_prior_rf_only = xr.open_dataset(self._TEST_SAMPLE_PRIOR_RF_ONLY_PATH)
     self.test_dist_media_and_rf = collections.OrderedDict({
         param: tf.convert_to_tensor(test_prior_media_and_rf[param])
@@ -232,6 +275,18 @@ class WithInputDataSamples:
         for param in constants.COMMON_PARAMETER_NAMES
         + constants.MEDIA_PARAMETER_NAMES
     })
+    self.test_dist_media_only_no_controls = collections.OrderedDict({
+        param: tf.convert_to_tensor(test_prior_media_only_no_controls[param])
+        for param in (
+            set(
+                constants.COMMON_PARAMETER_NAMES
+                + constants.MEDIA_PARAMETER_NAMES
+            )
+            - set(
+                constants.CONTROL_PARAMETERS + constants.GEO_CONTROL_PARAMETERS
+            )
+        )
+    })
     self.test_dist_rf_only = collections.OrderedDict({
         param: tf.convert_to_tensor(test_prior_rf_only[param])
         for param in constants.COMMON_PARAMETER_NAMES
@@ -244,6 +299,9 @@ class WithInputDataSamples:
     test_posterior_media_only = xr.open_dataset(
         self._TEST_SAMPLE_POSTERIOR_MEDIA_ONLY_PATH
     )
+    test_posterior_media_only_no_controls = xr.open_dataset(
+        self._TEST_SAMPLE_POSTERIOR_MEDIA_ONLY_NO_CONTROLS_PATH
+    )
     test_posterior_rf_only = xr.open_dataset(
         self._TEST_SAMPLE_POSTERIOR_RF_ONLY_PATH
     )
@@ -262,6 +320,21 @@ class WithInputDataSamples:
         for param in constants.COMMON_PARAMETER_NAMES
         + constants.MEDIA_PARAMETER_NAMES
     }
+    posterior_params_to_tensors_media_only_no_controls = {
+        param: _convert_with_swap(
+            test_posterior_media_only_no_controls[param],
+            n_burnin=self._N_BURNIN,
+        )
+        for param in (
+            set(
+                constants.COMMON_PARAMETER_NAMES
+                + constants.MEDIA_PARAMETER_NAMES
+            )
+            - set(
+                constants.CONTROL_PARAMETERS + constants.GEO_CONTROL_PARAMETERS
+            )
+        )
+    }
     posterior_params_to_tensors_rf_only = {
         param: _convert_with_swap(
             test_posterior_rf_only[param], n_burnin=self._N_BURNIN
@@ -279,6 +352,18 @@ class WithInputDataSamples:
         "StructTuple",
         constants.COMMON_PARAMETER_NAMES + constants.MEDIA_PARAMETER_NAMES,
     )(**posterior_params_to_tensors_media_only)
+    self.test_posterior_states_media_only_no_controls = collections.namedtuple(
+        "StructTuple",
+        (
+            set(
+                constants.COMMON_PARAMETER_NAMES
+                + constants.MEDIA_PARAMETER_NAMES
+            )
+            - set(
+                constants.CONTROL_PARAMETERS + constants.GEO_CONTROL_PARAMETERS
+            )
+        ),
+    )(**posterior_params_to_tensors_media_only_no_controls)
     self.test_posterior_states_rf_only = collections.namedtuple(
         "StructTuple",
         constants.COMMON_PARAMETER_NAMES + constants.RF_PARAMETER_NAMES,

meridian/model/posterior_sampler.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.
@@ -110,23 +110,16 @@ class PosteriorMCMCSampler:
     organic_media_tensors = mmm.organic_media_tensors
     organic_rf_tensors = mmm.organic_rf_tensors
     controls_scaled = mmm.controls_scaled
-    non_media_treatments_scaled = mmm.non_media_treatments_scaled
+    non_media_treatments_normalized = mmm.non_media_treatments_normalized
     media_effects_dist = mmm.media_effects_dist
     adstock_hill_media_fn = mmm.adstock_hill_media
     adstock_hill_rf_fn = mmm.adstock_hill_rf
-    get_roi_prior_beta_m_value_fn = (
-        mmm.prior_sampler_callable.get_roi_prior_beta_m_value
-    )
-    get_roi_prior_beta_rf_value_fn = (
-        mmm.prior_sampler_callable.get_roi_prior_beta_rf_value
-    )
+    total_outcome = mmm.total_outcome
 
     @tfp.distributions.JointDistributionCoroutineAutoBatched
     def joint_dist_unpinned():
       # Sample directly from prior.
       knot_values = yield prior_broadcast.knot_values
-      gamma_c = yield prior_broadcast.gamma_c
-      xi_c = yield prior_broadcast.xi_c
       sigma = yield prior_broadcast.sigma
 
       tau_g_excl_baseline = yield tfp.distributions.Sample(
@@ -167,26 +160,39 @@ class PosteriorMCMCSampler:
             ec=ec_m,
             slope=slope_m,
         )
-        prior_type = mmm.model_spec.paid_media_prior_type
-        if prior_type in constants.PAID_MEDIA_ROI_PRIOR_TYPES:
-          if prior_type == constants.PAID_MEDIA_PRIOR_TYPE_ROI:
-            roi_or_mroi_m = yield prior_broadcast.roi_m
+        prior_type = mmm.model_spec.effective_media_prior_type
+        if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
+          beta_m = yield prior_broadcast.beta_m
+        else:
+          if prior_type == constants.TREATMENT_PRIOR_TYPE_ROI:
+            treatment_parameter_m = yield prior_broadcast.roi_m
+          elif prior_type == constants.TREATMENT_PRIOR_TYPE_MROI:
+            treatment_parameter_m = yield prior_broadcast.mroi_m
+          elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
+            treatment_parameter_m = yield prior_broadcast.contribution_m
           else:
-            roi_or_mroi_m = yield prior_broadcast.mroi_m
-          beta_m_value = get_roi_prior_beta_m_value_fn(
-              alpha_m,
-              beta_gm_dev,
-              ec_m,
-              eta_m,
-              roi_or_mroi_m,
-              slope_m,
-              media_transformed,
+            raise ValueError(f"Unsupported prior type: {prior_type}")
+          incremental_outcome_m = (
+              treatment_parameter_m * media_tensors.prior_denominator
+          )
+          linear_predictor_counterfactual_difference = (
+              mmm.linear_predictor_counterfactual_difference_media(
+                  media_transformed=media_transformed,
+                  alpha_m=alpha_m,
+                  ec_m=ec_m,
+                  slope_m=slope_m,
+              )
+          )
+          beta_m_value = mmm.calculate_beta_x(
+              is_non_media=False,
+              incremental_outcome_x=incremental_outcome_m,
+              linear_predictor_counterfactual_difference=linear_predictor_counterfactual_difference,
+              eta_x=eta_m,
+              beta_gx_dev=beta_gm_dev,
           )
           beta_m = yield tfp.distributions.Deterministic(
               beta_m_value, name=constants.BETA_M
           )
-        else:
-          beta_m = yield prior_broadcast.beta_m
 
         beta_eta_combined = beta_m + eta_m * beta_gm_dev
         beta_gm_value = (
@@ -220,27 +226,39 @@ class PosteriorMCMCSampler:
             slope=slope_rf,
         )
 
-        prior_type = mmm.model_spec.paid_media_prior_type
-        if prior_type in constants.PAID_MEDIA_ROI_PRIOR_TYPES:
-          if prior_type == constants.PAID_MEDIA_PRIOR_TYPE_ROI:
-            roi_or_mroi_rf = yield prior_broadcast.roi_rf
+        prior_type = mmm.model_spec.effective_rf_prior_type
+        if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
+          beta_rf = yield prior_broadcast.beta_rf
+        else:
+          if prior_type == constants.TREATMENT_PRIOR_TYPE_ROI:
+            treatment_parameter_rf = yield prior_broadcast.roi_rf
+          elif prior_type == constants.TREATMENT_PRIOR_TYPE_MROI:
+            treatment_parameter_rf = yield prior_broadcast.mroi_rf
+          elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
+            treatment_parameter_rf = yield prior_broadcast.contribution_rf
           else:
-            roi_or_mroi_rf = yield prior_broadcast.mroi_rf
-          beta_rf_value = get_roi_prior_beta_rf_value_fn(
-              alpha_rf,
-              beta_grf_dev,
-              ec_rf,
-              eta_rf,
-              roi_or_mroi_rf,
-              slope_rf,
-              rf_transformed,
+            raise ValueError(f"Unsupported prior type: {prior_type}")
+          incremental_outcome_rf = (
+              treatment_parameter_rf * rf_tensors.prior_denominator
+          )
+          linear_predictor_counterfactual_difference = (
+              mmm.linear_predictor_counterfactual_difference_rf(
+                  rf_transformed=rf_transformed,
+                  alpha_rf=alpha_rf,
+                  ec_rf=ec_rf,
+                  slope_rf=slope_rf,
+              )
+          )
+          beta_rf_value = mmm.calculate_beta_x(
+              is_non_media=False,
+              incremental_outcome_x=incremental_outcome_rf,
+              linear_predictor_counterfactual_difference=linear_predictor_counterfactual_difference,
+              eta_x=eta_rf,
+              beta_gx_dev=beta_grf_dev,
           )
           beta_rf = yield tfp.distributions.Deterministic(
-              beta_rf_value,
-              name=constants.BETA_RF,
+              beta_rf_value, name=constants.BETA_RF
           )
-        else:
-          beta_rf = yield prior_broadcast.beta_rf
 
         beta_eta_combined = beta_rf + eta_rf * beta_grf_dev
         beta_grf_value = (
@@ -272,7 +290,24 @@ class PosteriorMCMCSampler:
             ec=ec_om,
             slope=slope_om,
         )
-        beta_om = yield prior_broadcast.beta_om
+        prior_type = mmm.model_spec.organic_media_prior_type
+        if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
+          beta_om = yield prior_broadcast.beta_om
+        elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
+          contribution_om = yield prior_broadcast.contribution_om
+          incremental_outcome_om = contribution_om * total_outcome
+          beta_om_value = mmm.calculate_beta_x(
+              is_non_media=False,
+              incremental_outcome_x=incremental_outcome_om,
+              linear_predictor_counterfactual_difference=organic_media_transformed,
+              eta_x=eta_om,
+              beta_gx_dev=beta_gom_dev,
+          )
+          beta_om = yield tfp.distributions.Deterministic(
+              beta_om_value, name=constants.BETA_OM
+          )
+        else:
+          raise ValueError(f"Unsupported prior type: {prior_type}")
 
         beta_eta_combined = beta_om + eta_om * beta_gom_dev
         beta_gom_value = (
@@ -306,7 +341,24 @@ class PosteriorMCMCSampler:
             slope=slope_orf,
         )
 
-        beta_orf = yield prior_broadcast.beta_orf
+        prior_type = mmm.model_spec.organic_rf_prior_type
+        if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
+          beta_orf = yield prior_broadcast.beta_orf
+        elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
+          contribution_orf = yield prior_broadcast.contribution_orf
+          incremental_outcome_orf = contribution_orf * total_outcome
+          beta_orf_value = mmm.calculate_beta_x(
+              is_non_media=False,
+              incremental_outcome_x=incremental_outcome_orf,
+              linear_predictor_counterfactual_difference=organic_rf_transformed,
+              eta_x=eta_orf,
+              beta_gx_dev=beta_gorf_dev,
+          )
+          beta_orf = yield tfp.distributions.Deterministic(
+              beta_orf_value, name=constants.BETA_ORF
+          )
+        else:
+          raise ValueError(f"Unsupported prior type: {prior_type}")
 
         beta_eta_combined = beta_orf + eta_orf * beta_gorf_dev
         beta_gorf_value = (
@@ -323,33 +375,62 @@ class PosteriorMCMCSampler:
         combined_beta = tf.concat([combined_beta, beta_gorf], axis=-1)
 
       sigma_gt = tf.transpose(tf.broadcast_to(sigma, [n_times, n_geos]))
-      gamma_gc_dev = yield tfp.distributions.Sample(
-          tfp.distributions.Normal(0, 1),
-          [n_geos, n_controls],
-          name=constants.GAMMA_GC_DEV,
-      )
-      gamma_gc = yield tfp.distributions.Deterministic(
-          gamma_c + xi_c * gamma_gc_dev, name=constants.GAMMA_GC
-      )
-      y_pred_combined_media = (
-          tau_gt
-          + tf.einsum("gtm,gm->gt", combined_media_transformed, combined_beta)
-          + tf.einsum("gtc,gc->gt", controls_scaled, gamma_gc)
+      y_pred_combined_media = tau_gt + tf.einsum(
+          "gtm,gm->gt", combined_media_transformed, combined_beta
       )
+      # Omit gamma_c, xi_c, and gamma_gc from joint distribution output if
+      # there are no control variables in the model.
+      if n_controls:
+        gamma_c = yield prior_broadcast.gamma_c
+        xi_c = yield prior_broadcast.xi_c
+        gamma_gc_dev = yield tfp.distributions.Sample(
+            tfp.distributions.Normal(0, 1),
+            [n_geos, n_controls],
+            name=constants.GAMMA_GC_DEV,
+        )
+        gamma_gc = yield tfp.distributions.Deterministic(
+            gamma_c + xi_c * gamma_gc_dev, name=constants.GAMMA_GC
+        )
+        y_pred_combined_media += tf.einsum(
+            "gtc,gc->gt", controls_scaled, gamma_gc
+        )
 
       if mmm.non_media_treatments is not None:
-        gamma_n = yield prior_broadcast.gamma_n
         xi_n = yield prior_broadcast.xi_n
         gamma_gn_dev = yield tfp.distributions.Sample(
             tfp.distributions.Normal(0, 1),
             [n_geos, n_non_media_channels],
             name=constants.GAMMA_GN_DEV,
         )
+        prior_type = mmm.model_spec.non_media_treatments_prior_type
+        if prior_type == constants.TREATMENT_PRIOR_TYPE_COEFFICIENT:
+          gamma_n = yield prior_broadcast.gamma_n
+        elif prior_type == constants.TREATMENT_PRIOR_TYPE_CONTRIBUTION:
+          contribution_n = yield prior_broadcast.contribution_n
+          incremental_outcome_n = contribution_n * total_outcome
+          baseline_scaled = mmm.non_media_transformer.forward(  # pytype: disable=attribute-error
+              mmm.compute_non_media_treatments_baseline()
+          )
+          linear_predictor_counterfactual_difference = (
+              non_media_treatments_normalized - baseline_scaled
+          )
+          gamma_n_value = mmm.calculate_beta_x(
+              is_non_media=True,
+              incremental_outcome_x=incremental_outcome_n,
+              linear_predictor_counterfactual_difference=linear_predictor_counterfactual_difference,
+              eta_x=xi_n,
+              beta_gx_dev=gamma_gn_dev,
+          )
+          gamma_n = yield tfp.distributions.Deterministic(
+              gamma_n_value, name=constants.GAMMA_N
+          )
+        else:
+          raise ValueError(f"Unsupported prior type: {prior_type}")
         gamma_gn = yield tfp.distributions.Deterministic(
             gamma_n + xi_n * gamma_gn_dev, name=constants.GAMMA_GN
         )
         y_pred = y_pred_combined_media + tf.einsum(
-            "gtn,gn->gt", non_media_treatments_scaled, gamma_gn
+            "gtn,gn->gt", non_media_treatments_normalized, gamma_gn
         )
       else:
         y_pred = y_pred_combined_media

meridian/model/prior_distribution.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.
@@ -84,6 +84,11 @@ class PriorDistribution:
   | `roi_rf`              | `n_rf_channels`            |
   | `mroi_m`              | `n_media_channels`         |
   | `mroi_rf`             | `n_rf_channels`            |
+  | `contribution_m`      | `n_media_channels`         |
+  | `contribution_rf`     | `n_rf_channels`            |
+  | `contribution_om`     | `n_organic_media_channels` |
+  | `contribution_orf`    | `n_organic_f_channels`     |
+  | `contribution_n`      | `n_non_media_channels`     |
 
   (σ) `n_geos` if `unique_sigma_for_each_geo`, otherwise this is `1`
 
@@ -233,6 +238,36 @@ class PriorDistribution:
       is interpreted as the marginal incremental KPI units per monetary unit
       spent. In this case, a default distribution is not provided, so the user
       must specify it.
+    contribution_m: Prior distribution on the contribution of each media channel
+      as a percentage of total outcome. This parameter is only used when
+      `paid_media_prior_type` is `'contribution'`, in which case `beta_m` is
+      calculated as a deterministic function of `contribution_m`, `alpha_m`,
+      `ec_m`, `slope_m`, and the total outcome. Default distribution is
+      `Beta(1.0, 99.0)`.
+    contribution_rf: Prior distribution on the contribution of each Reach &
+      Frequency channel as a percentage of total outcome. This parameter is only
+      used when `paid_media_prior_type` is `'contribution'`, in which case
+      `beta_rf` is calculated as a deterministic function of `contribution_rf`,
+      `alpha_rf`, `ec_rf`, `slope_rf`, and the total outcome. Default
+      distribution is `Beta(1.0, 99.0)`.
+    contribution_om: Prior distribution on the contribution of each organic
+      media channel as a percentage of total outcome. This parameter is only
+      used when `organic_media_prior_type` is `'contribution'`, in which case
+      `beta_om` is calculated as a deterministic function of `contribution_om`,
+      `alpha_om`, `ec_om`, `slope_om`, and the total outcome. Default
+      distribution is `Beta(1.0, 99.0)`.
+    contribution_orf: Prior distribution on the contribution of each organic
+      Reach & Frequency channel as a percentage of total outcome. This parameter
+      is only used when `organic_media_prior_type` is `'contribution'`, in which
+      case `beta_orf` is calculated as a deterministic function of
+      `contribution_orf`, `alpha_orf`, `ec_orf`, `slope_orf`, and the total
+      outcome. Default distribution is `Beta(1.0, 99.0)`.
+    contribution_n: Prior distribution on the contribution of each non-media
+      treatment channel as a percentage of total outcome. This parameter is only
+      used when `non_media_treatment_prior_type` is `'contribution'`, in which
+      case `gamma_n` is calculated as a deterministic function of
+      `contribution_n` and the total outcome. Default distribution is
+      `TruncatedNormal(0.0, 0.1, -1.0, 1.0)`.
   """
 
   knot_values: tfp.distributions.Distribution = dataclasses.field(
@@ -394,6 +429,31 @@ class PriorDistribution:
           0.0, 0.5, name=constants.MROI_RF
       ),
   )
+  contribution_m: tfp.distributions.Distribution = dataclasses.field(
+      default_factory=lambda: tfp.distributions.Beta(
+          1.0, 99.0, name=constants.CONTRIBUTION_M
+      ),
+  )
+  contribution_rf: tfp.distributions.Distribution = dataclasses.field(
+      default_factory=lambda: tfp.distributions.Beta(
+          1.0, 99.0, name=constants.CONTRIBUTION_RF
+      ),
+  )
+  contribution_om: tfp.distributions.Distribution = dataclasses.field(
+      default_factory=lambda: tfp.distributions.Beta(
+          1.0, 99.0, name=constants.CONTRIBUTION_OM
+      ),
+  )
+  contribution_orf: tfp.distributions.Distribution = dataclasses.field(
+      default_factory=lambda: tfp.distributions.Beta(
+          1.0, 99.0, name=constants.CONTRIBUTION_ORF
+      ),
+  )
+  contribution_n: tfp.distributions.Distribution = dataclasses.field(
+      default_factory=lambda: tfp.distributions.TruncatedNormal(
+          loc=0.0, scale=0.1, low=-1.0, high=1.0, name=constants.CONTRIBUTION_N
+      ),
+  )
 
   def __setstate__(self, state):
     # Override to support pickling.
@@ -455,7 +515,6 @@ class PriorDistribution:
       set_total_media_contribution_prior: bool,
       kpi: float,
       total_spend: np.ndarray,
-      media_effects_dist: str,
   ) -> PriorDistribution:
     """Returns a new `PriorDistribution` with broadcast distribution attributes.
 
@@ -481,8 +540,6 @@ class PriorDistribution:
         `set_total_media_contribution_prior=True`.
       total_spend: Spend per media channel summed across geos and time. Required
         if `set_total_media_contribution_prior=True`.
-      media_effects_dist: A string to specify the distribution of media random
-        effects across geos.
 
     Returns:
       A new `PriorDistribution` broadcast from this prior distribution,
@@ -508,6 +565,7 @@ class PriorDistribution:
 
     _validate_media_custom_priors(self.roi_m)
     _validate_media_custom_priors(self.mroi_m)
+    _validate_media_custom_priors(self.contribution_m)
     _validate_media_custom_priors(self.alpha_m)
     _validate_media_custom_priors(self.ec_m)
     _validate_media_custom_priors(self.slope_m)
@@ -529,6 +587,7 @@ class PriorDistribution:
             'that channel.'
         )
 
+    _validate_organic_media_custom_priors(self.contribution_om)
     _validate_organic_media_custom_priors(self.alpha_om)
     _validate_organic_media_custom_priors(self.ec_om)
     _validate_organic_media_custom_priors(self.slope_om)
@@ -550,6 +609,7 @@ class PriorDistribution:
             'for that channel.'
         )
 
+    _validate_organic_rf_custom_priors(self.contribution_orf)
     _validate_organic_rf_custom_priors(self.alpha_orf)
     _validate_organic_rf_custom_priors(self.ec_orf)
     _validate_organic_rf_custom_priors(self.slope_orf)
@@ -569,6 +629,7 @@ class PriorDistribution:
 
     _validate_rf_custom_priors(self.roi_rf)
     _validate_rf_custom_priors(self.mroi_rf)
+    _validate_rf_custom_priors(self.contribution_rf)
     _validate_rf_custom_priors(self.alpha_rf)
     _validate_rf_custom_priors(self.ec_rf)
     _validate_rf_custom_priors(self.slope_rf)
@@ -604,6 +665,7 @@ class PriorDistribution:
             'that channel.'
         )
 
+    _validate_non_media_custom_priors(self.contribution_n)
     _validate_non_media_custom_priors(self.gamma_n)
     _validate_non_media_custom_priors(self.xi_n)
 
@@ -743,57 +805,50 @@ class PriorDistribution:
         self.sigma, sigma_shape, name=constants.SIGMA
     )
 
-    default_distribution = PriorDistribution()
-    if set_total_media_contribution_prior and distributions_are_equal(
-        self.roi_m, default_distribution.roi_m
-    ):
-      warnings.warn(
-          'Consider setting custom ROI priors, as kpi_type was specified as'
-          ' `non_revenue` with no `revenue_per_kpi` being set. Otherwise, the'
-          ' total media contribution prior will be used with'
-          f' `p_mean={constants.P_MEAN}` and `p_sd={constants.P_SD}`. Further'
-          ' documentation available at '
-          ' https://developers.google.com/meridian/docs/advanced-modeling/unknown-revenue-kpi-custom#set-total-paid-media-contribution-prior',
-      )
+    if set_total_media_contribution_prior:
       roi_m_converted = _get_total_media_contribution_prior(
           kpi, total_spend, constants.ROI_M
       )
-    else:
-      roi_m_converted = self.roi_m
-    _check_for_negative_effect(roi_m_converted, media_effects_dist)
-    roi_m = tfp.distributions.BatchBroadcast(
-        roi_m_converted, n_media_channels, name=constants.ROI_M
-    )
-
-    if set_total_media_contribution_prior and distributions_are_equal(
-        self.roi_rf, default_distribution.roi_rf
-    ):
-      warnings.warn(
-          'Consider setting custom ROI priors, as kpi_type was specified as'
-          ' `non_revenue` with no `revenue_per_kpi` being set. Otherwise, the'
-          ' total media contribution prior will be used with'
-          f' `p_mean={constants.P_MEAN}` and `p_sd={constants.P_SD}`. Further'
-          ' documentation available at '
-          ' https://developers.google.com/meridian/docs/advanced-modeling/unknown-revenue-kpi-custom#set-total-paid-media-contribution-prior',
-      )
       roi_rf_converted = _get_total_media_contribution_prior(
           kpi, total_spend, constants.ROI_RF
       )
     else:
+      roi_m_converted = self.roi_m
       roi_rf_converted = self.roi_rf
-    _check_for_negative_effect(roi_rf_converted, media_effects_dist)
+    roi_m = tfp.distributions.BatchBroadcast(
+        roi_m_converted, n_media_channels, name=constants.ROI_M
+    )
     roi_rf = tfp.distributions.BatchBroadcast(
         roi_rf_converted, n_rf_channels, name=constants.ROI_RF
     )
-    _check_for_negative_effect(self.mroi_m, media_effects_dist)
+
     mroi_m = tfp.distributions.BatchBroadcast(
         self.mroi_m, n_media_channels, name=constants.MROI_M
     )
-    _check_for_negative_effect(self.mroi_rf, media_effects_dist)
     mroi_rf = tfp.distributions.BatchBroadcast(
         self.mroi_rf, n_rf_channels, name=constants.MROI_RF
     )
 
+    contribution_m = tfp.distributions.BatchBroadcast(
+        self.contribution_m, n_media_channels, name=constants.CONTRIBUTION_M
+    )
+    contribution_rf = tfp.distributions.BatchBroadcast(
+        self.contribution_rf, n_rf_channels, name=constants.CONTRIBUTION_RF
+    )
+    contribution_om = tfp.distributions.BatchBroadcast(
+        self.contribution_om,
+        n_organic_media_channels,
+        name=constants.CONTRIBUTION_OM,
+    )
+    contribution_orf = tfp.distributions.BatchBroadcast(
+        self.contribution_orf,
+        n_organic_rf_channels,
+        name=constants.CONTRIBUTION_ORF,
+    )
+    contribution_n = tfp.distributions.BatchBroadcast(
+        self.contribution_n, n_non_media_channels, name=constants.CONTRIBUTION_N
+    )
+
     return PriorDistribution(
         knot_values=knot_values,
         tau_g_excl_baseline=tau_g_excl_baseline,
@@ -826,6 +881,11 @@ class PriorDistribution:
         roi_rf=roi_rf,
         mroi_m=mroi_m,
         mroi_rf=mroi_rf,
+        contribution_m=contribution_m,
+        contribution_rf=contribution_rf,
+        contribution_om=contribution_om,
+        contribution_orf=contribution_orf,
+        contribution_n=contribution_n,
     )
 
 
@@ -891,21 +951,6 @@ def _get_total_media_contribution_prior(
   return tfp.distributions.LogNormal(lognormal_mu, lognormal_sigma, name=name)
 
 
-def _check_for_negative_effect(
-    dist: tfp.distributions.Distribution, media_effects_dist: str
-):
-  """Checks for negative effect in the model."""
-  if (
-      media_effects_dist == constants.MEDIA_EFFECTS_LOG_NORMAL
-      and np.any(dist.cdf(0)) > 0
-  ):
-    raise ValueError(
-        'Media priors must have non-negative support when'
-        f' `media_effects_dist`="{media_effects_dist}". Found negative effect'
-        f' in {dist.name}.'
-    )
-
-
 def distributions_are_equal(
     a: tfp.distributions.Distribution, b: tfp.distributions.Distribution
 ) -> bool: