Add proto and serde support for per-channel media_effects_dist.

PiperOrigin-RevId: 938055551

Author: nmahadevan-goog

meridian/model/spec.py

@@ -232,7 +232,9 @@ class ModelSpec:
   prior: prior_distribution.PriorDistribution = dataclasses.field(
       default_factory=prior_distribution.PriorDistribution,
   )
-  media_effects_dist: str = constants.MEDIA_EFFECTS_LOG_NORMAL
+  media_effects_dist: str | Mapping[str, str] = (
+      constants.MEDIA_EFFECTS_LOG_NORMAL
+  )
   hill_before_adstock: bool = False
   max_lag: int = 8
   unique_sigma_for_each_geo: bool = False
@@ -258,10 +260,26 @@ class ModelSpec:
 
   def __post_init__(self):
     # Validate media_effects_dist.
-    if self.media_effects_dist not in constants.MEDIA_EFFECTS_DISTRIBUTIONS:
+    if isinstance(self.media_effects_dist, str):
+      if self.media_effects_dist not in constants.MEDIA_EFFECTS_DISTRIBUTIONS:
+        raise ValueError(
+            "The `media_effects_dist` parameter"
+            f" '{self.media_effects_dist}' must be one of"
+            f" {sorted(constants.MEDIA_EFFECTS_DISTRIBUTIONS)}."
+        )
+    elif isinstance(self.media_effects_dist, Mapping):
+      for channel, dist in self.media_effects_dist.items():
+        if dist not in constants.MEDIA_EFFECTS_DISTRIBUTIONS:
+          raise ValueError(
+              f"The `media_effects_dist` for channel '{channel}'"
+              " must be one of"
+              f" {sorted(constants.MEDIA_EFFECTS_DISTRIBUTIONS)},"
+              f" but got '{dist}'."
+          )
+    else:
       raise ValueError(
-          f"The `media_effects_dist` parameter '{self.media_effects_dist}' must"
-          f" be one of {sorted(constants.MEDIA_EFFECTS_DISTRIBUTIONS)}."
+          "Unsupported type for `media_effects_dist` parameter:"
+          f" {type(self.media_effects_dist)}."
       )
     # Support paid_media_prior_type for backwards compatibility.
     if self.paid_media_prior_type is not None:

meridian/model/spec_test.py

@@ -74,6 +74,34 @@ def test_spec_inits_invalid_media_effects_fails(self, dist, error_message):
     with self.assertRaisesWithLiteralMatch(ValueError, error_message):
       spec.ModelSpec(media_effects_dist=dist)
 
+  def test_spec_inits_valid_media_effects_mapping_works(self):
+    dist_map = {"ch1": "normal", "ch2": "log_normal"}
+    model_spec = spec.ModelSpec(media_effects_dist=dist_map)
+    self.assertEqual(model_spec.media_effects_dist, dist_map)
+
+  @parameterized.named_parameters(
+      (
+          "invalid_mapping_value",
+          {"ch1": "bad_dist"},
+          ValueError,
+          (
+              "The `media_effects_dist` for channel 'ch1' must be one of"
+              " ['log_normal', 'normal'], but got 'bad_dist'."
+          ),
+      ),
+      (
+          "invalid_type",
+          123,
+          ValueError,
+          "Unsupported type for `media_effects_dist` parameter: <class 'int'>.",
+      ),
+  )
+  def test_spec_inits_invalid_media_effects_mapping_or_type_fails(
+      self, dist, expected_error, error_message
+  ):
+    with self.assertRaisesWithLiteralMatch(expected_error, error_message):
+      spec.ModelSpec(media_effects_dist=dist)
+
   @parameterized.named_parameters(
       ("hill", constants.HILL),
       ("none", "none"),

meridian/schema/serde/constants.py

@@ -25,6 +25,8 @@
 GLOBAL_ADSTOCK_DECAY = 'global_adstock_decay'
 ADSTOCK_DECAY_BY_CHANNEL = 'adstock_decay_by_channel'
 DEFAULT_DECAY = 'geometric'
+MEDIA_EFFECTS_DIST_BY_CHANNEL = 'media_effects_dist_by_channel'
+
 
 SATURATION_SPEC = 'saturation_spec'
 GLOBAL_SATURATION = 'global_saturation'

meridian/schema/serde/hyperparameters.py

@@ -87,9 +87,6 @@ class HyperparametersSerde(
   def serialize(self, obj: spec.ModelSpec) -> meridian_pb.Hyperparameters:
     """Serializes the given ModelSpec into a `Hyperparameters` proto."""
     hyperparameters_proto = meridian_pb.Hyperparameters(
-        media_effects_dist=media_effects_converter.to_proto(
-            obj.media_effects_dist
-        ),
         hill_before_adstock=obj.hill_before_adstock,
         unique_sigma_for_each_geo=obj.unique_sigma_for_each_geo,
         media_prior_type=paid_media_prior_type_converter.to_proto(
@@ -112,6 +109,15 @@ def serialize(self, obj: spec.ModelSpec) -> meridian_pb.Hyperparameters:
         ),
         enable_aks=obj.enable_aks,
     )
+    if isinstance(obj.media_effects_dist, str):
+      hyperparameters_proto.media_effects_dist = (
+          media_effects_converter.to_proto(obj.media_effects_dist)
+      )
+    elif isinstance(obj.media_effects_dist, Mapping):
+      for channel, dist in obj.media_effects_dist.items():
+        hyperparameters_proto.media_effects_dist_by_channel.channel_media_effects_dists[
+            channel
+        ] = media_effects_converter.to_proto(dist)
     if obj.max_lag is not None:
       hyperparameters_proto.max_lag = obj.max_lag
 
@@ -290,10 +296,24 @@ def deserialize(
     else:
       saturation_spec = sc.DEFAULT_SATURATION
 
+    if (
+        serialized.HasField(sc.MEDIA_EFFECTS_DIST_BY_CHANNEL)
+        and serialized.media_effects_dist_by_channel.channel_media_effects_dists
+    ):
+      channel_dists = (
+          serialized.media_effects_dist_by_channel.channel_media_effects_dists
+      )
+      media_effects_dist = {
+          channel: media_effects_converter.from_proto(dist)
+          for channel, dist in channel_dists.items()
+      }
+    else:
+      media_effects_dist = media_effects_converter.from_proto(
+          serialized.media_effects_dist
+      )
+
     return spec.ModelSpec(
-        media_effects_dist=media_effects_converter.from_proto(
-            serialized.media_effects_dist
-        ),
+        media_effects_dist=media_effects_dist,
         hill_before_adstock=serialized.hill_before_adstock,
         max_lag=max_lag,
         unique_sigma_for_each_geo=serialized.unique_sigma_for_each_geo,

proto/mmm/v1/model/meridian/meridian_model.proto

@@ -334,6 +334,13 @@ enum ComputationPrecision {
   FLOAT64 = 2;
 }
 
+// Specifies the media effects distribution for each channel.
+message MediaEffectsDistByChannel {
+  // A map where keys are channel names and values are the media effects
+  // distribution to use for that channel.
+  map<string, MediaEffectsDistribution> channel_media_effects_dists = 1;
+}
+
 // Specifies the adstock decay function for each channel.
 message AdstockDecayByChannel {
   // A map where keys are channel names and values are the adstock decay
@@ -539,6 +546,10 @@ message Hyperparameters {
   // If `None`, the minimum value is used as baseline for each non-media
   // treatments channel. This attribute is used as the default value for the
   // corresponding argument to `Analyzer` methods.
+  // Channel-specific media effects distributions. Defaults to 'log_normal' for
+  // channels not specified in the map.
+  MediaEffectsDistByChannel media_effects_dist_by_channel = 28;
+
   repeated NonMediaBaselineValue non_media_baseline_values = 25;
 }