google-meridian 1.1.3__py3-none-any.whl → 1.1.5__py3-none-any.whl

meridian/model/model.py

@@ -295,10 +295,6 @@ class Meridian:
   def is_national(self) -> bool:
     return self.n_geos == 1
 
-  @property
-  def _sigma_shape(self) -> int:
-    return len(self.input_data.geo) if self.unique_sigma_for_each_geo else 1
-
   @functools.cached_property
   def knot_info(self) -> knots.KnotInfo:
     return knots.get_knot_info(
@@ -389,6 +385,7 @@ class Meridian:
   @functools.cached_property
   def unique_sigma_for_each_geo(self) -> bool:
     if self.is_national:
+      # Should evaluate to False.
       return constants.NATIONAL_MODEL_SPEC_ARGS[
           constants.UNIQUE_SIGMA_FOR_EACH_GEO
       ]
@@ -449,7 +446,7 @@ class Meridian:
         n_organic_rf_channels=self.n_organic_rf_channels,
         n_controls=self.n_controls,
         n_non_media_channels=self.n_non_media_channels,
-        sigma_shape=self._sigma_shape,
+        unique_sigma_for_each_geo=self.unique_sigma_for_each_geo,
         n_knots=self.knot_info.n_knots,
         is_national=self.is_national,
         set_total_media_contribution_prior=self._set_total_media_contribution_prior,
@@ -663,10 +660,6 @@ class Meridian:
     self._validate_injected_inference_data_group_coord(
         inference_data, group, constants.TIME, self.n_times
     )
-    if not self.model_spec.unique_sigma_for_each_geo:
-      self._validate_injected_inference_data_group_coord(
-          inference_data, group, constants.SIGMA_DIM, self._sigma_shape
-      )
     self._validate_injected_inference_data_group_coord(
         inference_data,
         group,
@@ -1429,7 +1422,7 @@ class Meridian:
     if self.unique_sigma_for_each_geo:
       inference_dims[constants.SIGMA] = [constants.GEO]
     else:
-      inference_dims[constants.SIGMA] = [constants.SIGMA_DIM]
+      inference_dims[constants.SIGMA] = []
 
     return {
         param: [constants.CHAIN, constants.DRAW] + list(dims)

meridian/model/prior_distribution.py

@@ -509,7 +509,7 @@ class PriorDistribution:
       n_organic_rf_channels: int,
       n_controls: int,
       n_non_media_channels: int,
-      sigma_shape: int,
+      unique_sigma_for_each_geo: bool,
       n_knots: int,
       is_national: bool,
       set_total_media_contribution_prior: bool,
@@ -527,9 +527,9 @@ class PriorDistribution:
         used.
       n_controls: Number of controls used.
       n_non_media_channels: Number of non-media channels used.
-      sigma_shape: A number describing the shape of the sigma parameter. It's
-        either `1` (if `sigma_for_each_geo=False`) or `n_geos` (if
-        `sigma_for_each_geo=True`). For more information, see `ModelSpec`.
+      unique_sigma_for_each_geo: A boolean indicator whether to use the same
+        sigma parameter for all geos. Only used if `n_geos > 1`. For more
+        information, see `ModelSpec`.
       n_knots: Number of knots used.
       is_national: A boolean indicator whether the prior distribution will be
         adapted for a national model.
@@ -801,6 +801,9 @@ class PriorDistribution:
     slope_orf = tfp.distributions.BatchBroadcast(
         self.slope_orf, n_organic_rf_channels, name=constants.SLOPE_ORF
     )
+
+    # If `unique_sigma_for_each_geo == False`, then make a scalar batch.
+    sigma_shape = n_geos if (n_geos > 1 and unique_sigma_for_each_geo) else []
     sigma = tfp.distributions.BatchBroadcast(
         self.sigma, sigma_shape, name=constants.SIGMA
     )

meridian/model/prior_sampler.py

@@ -510,6 +510,8 @@ class PriorDistributionSampler:
       tf.keras.utils.set_random_seed(1)
 
     prior = mmm.prior_broadcast
+    # `sample_shape` is prepended to the shape of each BatchBroadcast in `prior`
+    # when it is sampled.
     sample_shape = [1, n_draws]
     sample_kwargs = {constants.SAMPLE_SHAPE: sample_shape, constants.SEED: seed}
 

meridian/version.py

@@ -14,4 +14,4 @@
 
 """Module for Meridian version."""
 
-__version__ = "1.1.3"
+__version__ = "1.1.5"