google-meridian 1.2.0__py3-none-any.whl → 1.2.1__py3-none-any.whl

meridian/model/eda/eda_engine.py

@@ -14,6 +14,7 @@
 
 """Meridian EDA Engine."""
 
+import dataclasses
 import functools
 from typing import Callable, Dict, Optional, TypeAlias
 from meridian import constants
@@ -28,11 +29,64 @@ _DEFAULT_DA_VAR_AGG_FUNCTION = np.sum
 AggregationMap: TypeAlias = Dict[str, Callable[[xr.DataArray], np.ndarray]]
 
 
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class ReachFrequencyData:
+  """Holds reach and frequency data arrays.
+
+  Attributes:
+    reach_raw_da: Raw reach data.
+    reach_scaled_da: Scaled reach data.
+    reach_raw_da_national: National raw reach data.
+    reach_scaled_da_national: National scaled reach data.
+    frequency_da: Frequency data.
+    frequency_da_national: National frequency data.
+    rf_impressions_scaled_da: Scaled reach * frequency impressions data.
+    rf_impressions_scaled_da_national: National scaled reach * frequency
+      impressions data.
+    rf_impressions_raw_da: Raw reach * frequency impressions data.
+    rf_impressions_raw_da_national: National raw reach * frequency impressions
+      data.
+  """
+
+  reach_raw_da: xr.DataArray
+  reach_scaled_da: xr.DataArray
+  reach_raw_da_national: xr.DataArray
+  reach_scaled_da_national: xr.DataArray
+  frequency_da: xr.DataArray
+  frequency_da_national: xr.DataArray
+  rf_impressions_scaled_da: xr.DataArray
+  rf_impressions_scaled_da_national: xr.DataArray
+  rf_impressions_raw_da: xr.DataArray
+  rf_impressions_raw_da_national: xr.DataArray
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class AggregationConfig:
+  """Configuration for custom aggregation functions.
+
+  Attributes:
+    control_variables: A dictionary mapping control variable names to
+      aggregation functions. Defaults to `np.sum` if a variable is not
+      specified.
+    non_media_treatments: A dictionary mapping non-media variable names to
+      aggregation functions. Defaults to `np.sum` if a variable is not
+      specified.
+  """
+
+  control_variables: AggregationMap = dataclasses.field(default_factory=dict)
+  non_media_treatments: AggregationMap = dataclasses.field(default_factory=dict)
+
+
 class EDAEngine:
   """Meridian EDA Engine."""
 
-  def __init__(self, meridian: model.Meridian):
+  def __init__(
+      self,
+      meridian: model.Meridian,
+      agg_config: AggregationConfig = AggregationConfig(),
+  ):
     self._meridian = meridian
+    self._agg_config = agg_config
 
   @functools.cached_property
   def controls_scaled_da(self) -> xr.DataArray | None:
@@ -44,6 +98,26 @@ class EDAEngine:
     )
     return controls_scaled_da
 
+  @functools.cached_property
+  def controls_scaled_da_national(self) -> xr.DataArray | None:
+    """Returns the national controls data array."""
+    if self._meridian.input_data.controls is None:
+      return None
+    if self._meridian.is_national:
+      if self.controls_scaled_da is None:
+        # This case should be impossible given the check above.
+        raise RuntimeError(
+            'controls_scaled_da is None when controls is not None.'
+        )
+      return self.controls_scaled_da.squeeze(constants.GEO)
+    else:
+      return self._aggregate_and_scale_geo_da(
+          self._meridian.input_data.controls,
+          transformers.CenteringAndScalingTransformer,
+          constants.CONTROL_VARIABLE,
+          self._agg_config.control_variables,
+      )
+
   @functools.cached_property
   def media_raw_da(self) -> xr.DataArray | None:
     if self._meridian.input_data.media is None:
@@ -71,13 +145,32 @@ class EDAEngine:
     # No need to truncate the media time for media spend.
     return media_spend_da
 
+  @functools.cached_property
+  def media_spend_da_national(self) -> xr.DataArray | None:
+    """Returns the national media spend data array."""
+    if self._meridian.input_data.media_spend is None:
+      return None
+    if self._meridian.is_national:
+      if self.media_spend_da is None:
+        # This case should be impossible given the check above.
+        raise RuntimeError(
+            'media_spend_da is None when media_spend is not None.'
+        )
+      return self.media_spend_da.squeeze(constants.GEO)
+    else:
+      return self._aggregate_and_scale_geo_da(
+          self._meridian.input_data.media_spend,
+          None,
+      )
+
   @functools.cached_property
   def media_raw_da_national(self) -> xr.DataArray | None:
     if self.media_raw_da is None:
       return None
     if self._meridian.is_national:
-      return self.media_raw_da
+      return self.media_raw_da.squeeze(constants.GEO)
     else:
+      # Note that media is summable by assumption.
       return self._aggregate_and_scale_geo_da(
           self.media_raw_da,
           None,
@@ -88,8 +181,9 @@ class EDAEngine:
     if self.media_scaled_da is None:
       return None
     if self._meridian.is_national:
-      return self.media_scaled_da
+      return self.media_scaled_da.squeeze(constants.GEO)
     else:
+      # Note that media is summable by assumption.
       return self._aggregate_and_scale_geo_da(
           self.media_raw_da,
           transformers.MediaTransformer,
@@ -116,8 +210,9 @@ class EDAEngine:
     if self.organic_media_raw_da is None:
       return None
     if self._meridian.is_national:
-      return self.organic_media_raw_da
+      return self.organic_media_raw_da.squeeze(constants.GEO)
     else:
+      # Note that organic media is summable by assumption.
       return self._aggregate_and_scale_geo_da(self.organic_media_raw_da, None)
 
   @functools.cached_property
@@ -125,8 +220,9 @@ class EDAEngine:
     if self.organic_media_scaled_da is None:
       return None
     if self._meridian.is_national:
-      return self.organic_media_scaled_da
+      return self.organic_media_scaled_da.squeeze(constants.GEO)
     else:
+      # Note that organic media is summable by assumption.
       return self._aggregate_and_scale_geo_da(
           self.organic_media_raw_da,
           transformers.MediaTransformer,
@@ -142,6 +238,26 @@ class EDAEngine:
     )
     return non_media_scaled_da
 
+  @functools.cached_property
+  def non_media_scaled_da_national(self) -> xr.DataArray | None:
+    """Returns the national non-media treatment data array."""
+    if self._meridian.input_data.non_media_treatments is None:
+      return None
+    if self._meridian.is_national:
+      if self.non_media_scaled_da is None:
+        # This case should be impossible given the check above.
+        raise RuntimeError(
+            'non_media_scaled_da is None when non_media_treatments is not None.'
+        )
+      return self.non_media_scaled_da.squeeze(constants.GEO)
+    else:
+      return self._aggregate_and_scale_geo_da(
+          self._meridian.input_data.non_media_treatments,
+          transformers.CenteringAndScalingTransformer,
+          constants.NON_MEDIA_CHANNEL,
+          self._agg_config.non_media_treatments,
+      )
+
   @functools.cached_property
   def rf_spend_da(self) -> xr.DataArray | None:
     if self._meridian.input_data.rf_spend is None:
@@ -157,12 +273,226 @@ class EDAEngine:
     if self._meridian.input_data.rf_spend is None:
       return None
     if self._meridian.is_national:
-      return self.rf_spend_da
+      if self.rf_spend_da is None:
+        # This case should be impossible given the check above.
+        raise RuntimeError('rf_spend_da is None when rf_spend is not None.')
+      return self.rf_spend_da.squeeze(constants.GEO)
     else:
       return self._aggregate_and_scale_geo_da(
           self._meridian.input_data.rf_spend, None
       )
 
+  @functools.cached_property
+  def _rf_data(self) -> ReachFrequencyData | None:
+    if self._meridian.input_data.reach is None:
+      return None
+    return self._get_rf_data(
+        self._meridian.input_data.reach,
+        self._meridian.input_data.frequency,
+        is_organic=False,
+    )
+
+  @property
+  def reach_raw_da(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.reach_raw_da
+
+  @property
+  def reach_scaled_da(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.reach_scaled_da
+
+  @property
+  def reach_raw_da_national(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.reach_raw_da_national
+
+  @property
+  def reach_scaled_da_national(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.reach_scaled_da_national
+
+  @property
+  def frequency_da(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.frequency_da
+
+  @property
+  def frequency_da_national(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.frequency_da_national
+
+  @property
+  def rf_impressions_raw_da(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.rf_impressions_raw_da
+
+  @property
+  def rf_impressions_raw_da_national(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.rf_impressions_raw_da_national
+
+  @property
+  def rf_impressions_scaled_da(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.rf_impressions_scaled_da
+
+  @property
+  def rf_impressions_scaled_da_national(self) -> xr.DataArray | None:
+    if self._rf_data is None:
+      return None
+    return self._rf_data.rf_impressions_scaled_da_national
+
+  @functools.cached_property
+  def _organic_rf_data(self) -> ReachFrequencyData | None:
+    if self._meridian.input_data.organic_reach is None:
+      return None
+    return self._get_rf_data(
+        self._meridian.input_data.organic_reach,
+        self._meridian.input_data.organic_frequency,
+        is_organic=True,
+    )
+
+  @property
+  def organic_reach_raw_da(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.reach_raw_da
+
+  @property
+  def organic_reach_scaled_da(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.reach_scaled_da
+
+  @property
+  def organic_reach_raw_da_national(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.reach_raw_da_national
+
+  @property
+  def organic_reach_scaled_da_national(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.reach_scaled_da_national
+
+  @property
+  def organic_rf_impressions_scaled_da(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.rf_impressions_scaled_da
+
+  @property
+  def organic_rf_impressions_scaled_da_national(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.rf_impressions_scaled_da_national
+
+  @property
+  def organic_frequency_da(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.frequency_da
+
+  @property
+  def organic_frequency_da_national(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.frequency_da_national
+
+  @property
+  def organic_rf_impressions_raw_da(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.rf_impressions_raw_da
+
+  @property
+  def organic_rf_impressions_raw_da_national(self) -> xr.DataArray | None:
+    if self._organic_rf_data is None:
+      return None
+    return self._organic_rf_data.rf_impressions_raw_da_national
+
+  @functools.cached_property
+  def geo_population_da(self) -> xr.DataArray | None:
+    if self._meridian.is_national:
+      return None
+    return xr.DataArray(
+        self._meridian.population,
+        coords={constants.GEO: self._meridian.input_data.geo.values},
+        dims=[constants.GEO],
+        name=constants.POPULATION,
+    )
+
+  @functools.cached_property
+  def kpi_scaled_da(self) -> xr.DataArray:
+    return _data_array_like(
+        da=self._meridian.input_data.kpi,
+        values=self._meridian.kpi_scaled,
+    )
+
+  @functools.cached_property
+  def kpi_scaled_da_national(self) -> xr.DataArray:
+    if self._meridian.is_national:
+      return self.kpi_scaled_da.squeeze(constants.GEO)
+    else:
+      # Note that kpi is summable by assumption.
+      return self._aggregate_and_scale_geo_da(
+          self._meridian.input_data.kpi,
+          transformers.CenteringAndScalingTransformer,
+      )
+
+  @functools.cached_property
+  def treatment_control_scaled_ds(self) -> xr.Dataset:
+    """Returns a Dataset containing all scaled treatments and controls.
+
+    This includes media, RF impressions, organic media, organic RF impressions,
+    non-media treatments, and control variables, all at the geo level.
+    """
+    to_merge = [
+        da
+        for da in [
+            self.media_scaled_da,
+            self.rf_impressions_scaled_da,
+            self.organic_media_scaled_da,
+            self.organic_rf_impressions_scaled_da,
+            self.controls_scaled_da,
+            self.non_media_scaled_da,
+        ]
+        if da is not None
+    ]
+    return xr.merge(to_merge, join='inner')
+
+  @functools.cached_property
+  def treatment_control_scaled_ds_national(self) -> xr.Dataset:
+    """Returns a Dataset containing all scaled treatments and controls.
+
+    This includes media, RF impressions, organic media, organic RF impressions,
+    non-media treatments, and control variables, all at the national level.
+    """
+    to_merge_national = [
+        da
+        for da in [
+            self.media_scaled_da_national,
+            self.rf_impressions_scaled_da_national,
+            self.organic_media_scaled_da_national,
+            self.organic_rf_impressions_scaled_da_national,
+            self.controls_scaled_da_national,
+            self.non_media_scaled_da_national,
+        ]
+        if da is not None
+    ]
+    return xr.merge(to_merge_national, join='inner')
+
   def _truncate_media_time(self, da: xr.DataArray) -> xr.DataArray:
     """Truncates the first `start` elements of the media time of a variable."""
     # This should not happen. If it does, it means this function is mis-used.
@@ -183,15 +513,17 @@ class EDAEngine:
       population: tf.Tensor = tf.constant([1.0], dtype=tf.float32),
   ):
     """Scales xarray values with a TensorTransformer."""
+    da = xarray.copy()
+
     if transformer_class is None:
-      return xarray
+      return da
     elif transformer_class is transformers.CenteringAndScalingTransformer:
       xarray_transformer = transformers.CenteringAndScalingTransformer(
-          tensor=xarray.values, population=population
+          tensor=da.values, population=population
       )
     elif transformer_class is transformers.MediaTransformer:
       xarray_transformer = transformers.MediaTransformer(
-          media=xarray.values, population=population
+          media=da.values, population=population
       )
     else:
       raise ValueError(
@@ -200,8 +532,8 @@ class EDAEngine:
           + '.\nMust be one of: CenteringAndScalingTransformer or'
           ' MediaTransformer.'
       )
-    xarray.values = xarray_transformer.forward(xarray.values)
-    return xarray
+    da.values = xarray_transformer.forward(da.values)
+    return da
 
   def _aggregate_variables(
       self,
@@ -282,6 +614,103 @@ class EDAEngine:
 
     return da_national.sel({constants.GEO: temp_geo_dim}, drop=True)
 
+  def _get_rf_data(
+      self,
+      reach_raw_da: xr.DataArray,
+      freq_raw_da: xr.DataArray,
+      is_organic: bool,
+  ) -> ReachFrequencyData:
+    """Get impressions and frequencies data arrays for RF channels."""
+    if is_organic:
+      scaled_reach_values = (
+          self._meridian.organic_rf_tensors.organic_reach_scaled
+      )
+    else:
+      scaled_reach_values = self._meridian.rf_tensors.reach_scaled
+    reach_scaled_da = _data_array_like(
+        da=reach_raw_da, values=scaled_reach_values
+    )
+    # Truncate the media time for reach and scaled reach.
+    reach_raw_da = self._truncate_media_time(reach_raw_da)
+    reach_scaled_da = self._truncate_media_time(reach_scaled_da)
+
+    # The geo level frequency
+    frequency_da = self._truncate_media_time(freq_raw_da)
+
+    # The raw geo level impression
+    # It's equal to reach * frequency.
+    impressions_raw_da = reach_raw_da * frequency_da
+    impressions_raw_da.name = (
+        constants.ORGANIC_RF_IMPRESSIONS
+        if is_organic
+        else constants.RF_IMPRESSIONS
+    )
+    impressions_raw_da.values = tf.cast(impressions_raw_da.values, tf.float32)
+
+    if self._meridian.is_national:
+      reach_raw_da_national = reach_raw_da.squeeze(constants.GEO)
+      reach_scaled_da_national = reach_scaled_da.squeeze(constants.GEO)
+      impressions_raw_da_national = impressions_raw_da.squeeze(constants.GEO)
+      frequency_da_national = frequency_da.squeeze(constants.GEO)
+
+      # Scaled impressions
+      impressions_scaled_da = self._scale_xarray(
+          impressions_raw_da, transformers.MediaTransformer
+      )
+      impressions_scaled_da_national = impressions_scaled_da.squeeze(
+          constants.GEO
+      )
+    else:
+      reach_raw_da_national = self._aggregate_and_scale_geo_da(
+          reach_raw_da, None
+      )
+      reach_scaled_da_national = self._aggregate_and_scale_geo_da(
+          reach_raw_da, transformers.MediaTransformer
+      )
+      impressions_raw_da_national = self._aggregate_and_scale_geo_da(
+          impressions_raw_da, None
+      )
+
+      # National frequency is a weighted average of geo frequencies,
+      # weighted by reach.
+      frequency_da_national = xr.where(
+          reach_raw_da_national == 0.0,
+          0.0,
+          impressions_raw_da_national / reach_raw_da_national,
+      )
+      frequency_da_national.name = (
+          constants.ORGANIC_PREFIX if is_organic else ''
+      ) + constants.FREQUENCY
+      frequency_da_national.values = tf.cast(
+          frequency_da_national.values, tf.float32
+      )
+
+      # Scale the impressions by population
+      impressions_scaled_da = self._scale_xarray(
+          impressions_raw_da,
+          transformers.MediaTransformer,
+          population=self._meridian.population,
+      )
+
+      # Scale the national impressions
+      impressions_scaled_da_national = self._aggregate_and_scale_geo_da(
+          impressions_raw_da,
+          transformers.MediaTransformer,
+      )
+
+    return ReachFrequencyData(
+        reach_raw_da=reach_raw_da,
+        reach_scaled_da=reach_scaled_da,
+        reach_raw_da_national=reach_raw_da_national,
+        reach_scaled_da_national=reach_scaled_da_national,
+        frequency_da=frequency_da,
+        frequency_da_national=frequency_da_national,
+        rf_impressions_scaled_da=impressions_scaled_da,
+        rf_impressions_scaled_da_national=impressions_scaled_da_national,
+        rf_impressions_raw_da=impressions_raw_da,
+        rf_impressions_raw_da_national=impressions_raw_da_national,
+    )
+
 
 def _data_array_like(
     *, da: xr.DataArray, values: np.ndarray | tf.Tensor

meridian/model/knots.py

@@ -276,7 +276,7 @@ class AKS:
     feasible_idx = np.where(
         (n_knots >= min_internal_knots) & (n_knots <= max_internal_knots)
     )[0]
-    information_criterion = aspline[constants.EBIC][feasible_idx]
+    information_criterion = aspline[constants.AIC][feasible_idx]
     knots_sel = [aspline[constants.KNOTS_SELECTED][i] for i in feasible_idx]
     model = [aspline[constants.MODEL][i] for i in feasible_idx]
     opt_idx = max(

meridian/model/model_test_data.py

@@ -52,7 +52,9 @@ def _convert_with_swap(array: xr.DataArray, n_burnin: int) -> backend.Tensor:
   else:
     pad_value = 0.0 if array.dtype.kind == "f" else 0
 
-  burnin = backend.fill([n_burnin] + transposed_tensor.shape[1:], pad_value)
+  burnin = backend.fill(
+      [n_burnin] + list(transposed_tensor.shape[1:]), pad_value
+  )
   return backend.concatenate(
       [burnin, transposed_tensor],
       axis=0,
@@ -122,18 +124,13 @@ class WithInputDataSamples:
   _N_MEDIA_CHANNELS = 3
   _N_RF_CHANNELS = 2
   _N_CONTROLS = 2
-  _ROI_CALIBRATION_PERIOD = backend.cast(
-      backend.ones((_N_MEDIA_TIMES_SHORT, _N_MEDIA_CHANNELS)),
-      dtype=backend.bool_,
-  )
-  _RF_ROI_CALIBRATION_PERIOD = backend.cast(
-      backend.ones((_N_MEDIA_TIMES_SHORT, _N_RF_CHANNELS)),
-      dtype=backend.bool_,
-  )
   _N_ORGANIC_MEDIA_CHANNELS = 4
   _N_ORGANIC_RF_CHANNELS = 1
   _N_NON_MEDIA_CHANNELS = 2
 
+  _ROI_CALIBRATION_PERIOD: backend.Tensor
+  _RF_ROI_CALIBRATION_PERIOD: backend.Tensor
+
   # Private class variables to hold the base test data.
   _input_data_non_revenue_no_revenue_per_kpi: input_data.InputData
   _input_data_media_and_rf_non_revenue_no_revenue_per_kpi: input_data.InputData
@@ -170,6 +167,15 @@ class WithInputDataSamples:
   @classmethod
   def setup(cls):
     """Sets up input data samples."""
+    cls._ROI_CALIBRATION_PERIOD = backend.cast(
+        backend.ones((cls._N_MEDIA_TIMES_SHORT, cls._N_MEDIA_CHANNELS)),
+        dtype=backend.bool_,
+    )
+    cls._RF_ROI_CALIBRATION_PERIOD = backend.cast(
+        backend.ones((cls._N_MEDIA_TIMES_SHORT, cls._N_RF_CHANNELS)),
+        dtype=backend.bool_,
+    )
+
     cls._input_data_non_revenue_no_revenue_per_kpi = (
         test_utils.sample_input_data_non_revenue_no_revenue_per_kpi(
             n_geos=cls._N_GEOS,