google-meridian 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

meridian/analysis/visualizer.py

@@ -1233,10 +1233,10 @@ class MediaEffects:
       bands for the Hill curves.
     """
     if c.MEDIA in list(df_channel_type[c.CHANNEL_TYPE]):
-      x_axis_title = 'Media Units per Capita'
+      x_axis_title = summary_text.HILL_X_AXIS_MEDIA_LABEL
       shaded_area_title = summary_text.HILL_SHADED_REGION_MEDIA_LABEL
     else:
-      x_axis_title = 'Average Frequency'
+      x_axis_title = summary_text.HILL_X_AXIS_RF_LABEL
       shaded_area_title = summary_text.HILL_SHADED_REGION_RF_LABEL
     domain_list = [
         c.POSTERIOR,
@@ -1380,6 +1380,7 @@ class MediaSummary:
       confidence_level: float = c.DEFAULT_CONFIDENCE_LEVEL,
       selected_times: Sequence[str] | None = None,
       marginal_roi_by_reach: bool = True,
+      non_media_baseline_values: Sequence[str | float] | None = None,
   ):
     """Initializes the media summary metrics based on the model data and params.
 
@@ -1394,12 +1395,20 @@ class MediaSummary:
         next dollar spent only impacts reach, holding frequency constant.  If
         this argument is False, we assume the next dollar spent only impacts
         frequency, holding reach constant.
+      non_media_baseline_values: Optional list of shape (n_non_media_channels,).
+        Each element is either a float (which means that the fixed value will be
+        used as baseline for the given channel) or one of the strings "min" or
+        "max" (which mean that the global minimum or maximum value will be used
+        as baseline for the values of the given non_media treatment channel). If
+        None, the minimum value is used as baseline for each non_media treatment
+        channel.
     """
     self._meridian = meridian
     self._analyzer = analyzer.Analyzer(meridian)
     self._confidence_level = confidence_level
     self._selected_times = selected_times
     self._marginal_roi_by_reach = marginal_roi_by_reach
+    self._non_media_baseline_values = non_media_baseline_values
 
   @functools.cached_property
   def paid_summary_metrics(self) -> xr.Dataset:
@@ -1438,6 +1447,7 @@ class MediaSummary:
         use_kpi=self._meridian.input_data.revenue_per_kpi is None,
         confidence_level=self._confidence_level,
         include_non_paid_channels=True,
+        non_media_baseline_values=self._non_media_baseline_values,
     )
 
   def summary_table(
@@ -1560,6 +1570,7 @@ class MediaSummary:
       confidence_level: float | None = None,
       selected_times: Sequence[str] | None = None,
       marginal_roi_by_reach: bool = True,
+      non_media_baseline_values: Sequence[str | float] | None = None,
   ):
     """Runs the computation for the media summary metrics with new parameters.
 
@@ -1574,10 +1585,18 @@ class MediaSummary:
         dollar spent only impacts reach, holding frequency constant. If `False`,
         the assumption is the next dollar spent only impacts frequency, holding
         reach constant.
+      non_media_baseline_values: Optional list of shape (n_non_media_channels,).
+        Each element is either a float (which means that the fixed value will be
+        used as baseline for the given channel) or one of the strings "min" or
+        "max" (which mean that the global minimum or maximum value will be used
+        as baseline for the values of the given non_media treatment channel). If
+        None, the minimum value is used as baseline for each non_media treatment
+        channel.
     """
     self._confidence_level = confidence_level or self._confidence_level
     self._selected_times = selected_times
     self._marginal_roi_by_reach = marginal_roi_by_reach
+    self._non_media_baseline_values = non_media_baseline_values
 
   def plot_contribution_waterfall_chart(self) -> alt.Chart:
     """Plots a waterfall chart of the contribution share per channel.

meridian/data/__init__.py

@@ -14,6 +14,7 @@
 
 """Data handling API for Meridian."""
 
+from meridian.data import arg_builder
 from meridian.data import input_data
 from meridian.data import load
 from meridian.data import time_coordinates

meridian/data/arg_builder.py

@@ -0,0 +1,107 @@
+# Copyright 2024 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Data argument builders for various API surfaces in Meridian."""
+
+from collections.abc import Sequence
+from typing import Generic, TypeVar
+
+
+__all__ = [
+    'OrderedListArgumentBuilder',
+]
+
+
+T = TypeVar('T')
+
+
+class OrderedListArgumentBuilder(Generic[T]):
+  """A simple builder for an argument that expects an ordered list of values.
+
+  For example, some Meridian function requires a list/array of some values that:
+
+  - Must have the same length as some channel coordinates in the input data.
+  - Must be in the same order as some channel coordinates in the input data.
+
+  This argument builder can be bound to one such channel coordinates list, and
+  provides a user-friendly and human-readable way to seed it with user given
+  values.
+
+  For example, to create an array (list) of values that must be indexed on the
+  *paid media* channels in the input data:
+
+  ```python
+  paid_media_channels_arg_builder = (
+      OrderedListArgumentBuilder[float](input_data.get_all_paid_channels())
+  )
+  # Note: rather than creating this builder directly, use methods like
+  # `InputData.get_all_paid_channels_argument_builder()` where the container
+  # determines which coordinates to bind the builder with.
+
+  # Use `.with_default_value()` to set a default value for coordinates that are
+  # not given in `__call__`.
+  paid_media_channels_arg_builder = (
+      paid_media_channels_arg_builder.with_default_value(0.3)
+  )
+
+  # Assuming we have paid media channels ['display', 'search', social'].
+  some_arg = paid_media_channels_arg_builder(
+      display=0.1,
+      social=0.25,
+  )
+  # some_arg == [0.1, 0.3, 0.25]
+  ```
+
+  See: `InputData.get_paid_channels_argument_builder()`.
+  """
+
+  def __init__(self, ordered_coords: Sequence[str]):
+    self._ordered_coords = list(ordered_coords)
+    self._default_value = None  # Applied when a coordinate value is not given.
+
+  def with_default_value(
+      self, default_value: T
+  ) -> 'OrderedListArgumentBuilder':
+    """Sets the default value for coordinates that are not given in `__call__`.
+
+    Args:
+      default_value: The default value to use for coordinates that are not given
+        in `__call__`. If unset (or set to `None`), then `__call__` will raise
+        an error if any bound coordinate's value is not given.
+
+    Returns:
+      This builder itself for fluent chaining.
+    """
+    self._default_value = default_value
+    return self
+
+  def __call__(self, **kwargs) -> list[T]:
+    """Builds an ordered argument values list, given the bound coordinates list.
+
+    Args:
+      **kwargs: The keys in `kwargs` are channel names. All channel names must
+        be present in the `ordered_coords` bound to this builder.
+
+    Returns:
+      A list of values, in the same order as the `ordered_coords` bound to this
+      builder.
+    """
+    if self._default_value is None and (
+        set(kwargs.keys()) != set(self._ordered_coords)
+    ):
+      raise ValueError(
+          'All coordinates must be present in the given keyword arguments: '
+          f'Given: {kwargs.keys()} vs Expected: {self._ordered_coords}'
+      )
+    return [kwargs.get(c, self._default_value) for c in self._ordered_coords]

meridian/data/input_data.py

@@ -24,6 +24,7 @@ import functools
 import warnings
 
 from meridian import constants
+from meridian.data import arg_builder
 from meridian.data import time_coordinates as tc
 import numpy as np
 import xarray as xr
@@ -678,6 +679,28 @@ class InputData:
       raise ValueError("Both RF and media channel values are missing.")
     # pytype: enable=attribute-error
 
+  def get_paid_channels_argument_builder(
+      self,
+  ) -> arg_builder.OrderedListArgumentBuilder:
+    """Returns an argument builder for all *paid* channels."""
+    return arg_builder.OrderedListArgumentBuilder(self.get_all_paid_channels())
+
+  def get_paid_media_channels_argument_builder(
+      self,
+  ) -> arg_builder.OrderedListArgumentBuilder:
+    """Returns an argument builder for *paid* media channels *only*."""
+    if self.media_channel is None:
+      raise ValueError("There are no media channels in the input data.")
+    return arg_builder.OrderedListArgumentBuilder(self.media_channel.values)
+
+  def get_paid_rf_channels_argument_builder(
+      self,
+  ) -> arg_builder.OrderedListArgumentBuilder:
+    """Returns an argument builder for *paid* RF channels *only*."""
+    if self.rf_channel is None:
+      raise ValueError("There are no RF channels in the input data.")
+    return arg_builder.OrderedListArgumentBuilder(self.rf_channel.values)
+
   def get_all_channels(self) -> np.ndarray:
     """Returns all the channel dimensions.
 

meridian/data/test_utils.py

@@ -1480,7 +1480,9 @@ def sample_coord_to_columns(
   )
 
 
-def sample_input_data_from_dataset(dataset: xr.Dataset, kpi_type: str):
+def sample_input_data_from_dataset(
+    dataset: xr.Dataset, kpi_type: str
+) -> input_data.InputData:
   """Generates a sample `InputData` from a full xarray Dataset."""
   return input_data.InputData(
       kpi=dataset.kpi,
@@ -1507,7 +1509,7 @@ def sample_input_data_revenue(
     n_organic_media_channels: int | None = None,
     n_organic_rf_channels: int | None = None,
     seed: int = 0,
-):
+) -> input_data.InputData:
   """Generates sample InputData for `kpi_type='revenue'`."""
   dataset = random_dataset(
       n_geos=n_geos,
@@ -1555,7 +1557,7 @@ def sample_input_data_non_revenue_revenue_per_kpi(
     n_organic_media_channels: int | None = None,
     n_organic_rf_channels: int | None = None,
     seed: int = 0,
-):
+) -> input_data.InputData:
   """Generates sample InputData for `non_revenue` KPI w/ revenue_per_kpi."""
   dataset = random_dataset(
       n_geos=n_geos,
@@ -1602,7 +1604,7 @@ def sample_input_data_non_revenue_no_revenue_per_kpi(
     n_organic_media_channels: int | None = None,
     n_organic_rf_channels: int | None = None,
     seed: int = 0,
-):
+) -> input_data.InputData:
   """Generates sample InputData for `non_revenue` KPI w/o revenue_per_kpi."""
   dataset = random_dataset(
       n_geos=n_geos,

meridian/model/__init__.py

@@ -18,6 +18,8 @@ from meridian.model import adstock_hill
 from meridian.model import knots
 from meridian.model import media
 from meridian.model import model
+from meridian.model import posterior_sampler
 from meridian.model import prior_distribution
+from meridian.model import prior_sampler
 from meridian.model import spec
 from meridian.model import transformers