google-meridian 1.3.2__py3-none-any.whl → 1.5.0__py3-none-any.whl

schema/processors/reach_frequency_optimization_processor.py

@@ -0,0 +1,584 @@
+# 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.
+# 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.
+
+"""Defines a processor for reach and frequency optimization inference on a Meridian model.
+
+This module provides the `ReachFrequencyOptimizationProcessor`, which optimizes
+the average frequency for reach and frequency (R&F) media channels in a trained
+Meridian model to maximize ROI.
+
+The processor takes a trained model and a `ReachFrequencyOptimizationSpec`
+object. The spec defines the constraints for the optimization, such as the
+minimum and maximum average frequency to consider for each channel.
+
+Key Features:
+
+-   Optimizes average frequency for all R&F channels simultaneously.
+-   Allows setting minimum and maximum frequency constraints.
+-   Generates detailed results, including the optimal average frequency for
+    each channel, the expected outcomes at this optimal frequency, and
+    response curves showing KPI/Revenue as a function of spend.
+-   Outputs results in a structured protobuf format
+    (`ReachFrequencyOptimization`).
+
+Key Classes:
+
+-   `ReachFrequencyOptimizationSpec`: Dataclass to specify optimization
+    parameters and constraints.
+-   `ReachFrequencyOptimizationProcessor`: The main processor class to execute
+    the R&F optimization.
+
+Example Usage:
+
+```python
+from schema.processors import reach_frequency_optimization_processor
+from schema.processors import common
+from schema.processors import model_processor
+import datetime
+
+# Assuming 'mmm' is a trained Meridian model object with R&F channels
+trained_model = model_processor.TrainedModel(mmm)
+
+spec = reach_frequency_optimization_processor.ReachFrequencyOptimizationSpec(
+    optimization_name="rf_optimize_q1",
+    start_date=datetime.date(2023, 1, 1),
+    end_date=datetime.date(2023, 4, 1),
+    min_frequency=1.0,
+    max_frequency=10.0,  # Optional, defaults to model's max frequency
+    kpi_type=common.KpiType.REVENUE,
+)
+
+processor = (
+    reach_frequency_optimization_processor.ReachFrequencyOptimizationProcessor(
+        trained_model
+    )
+)
+# result is a rf_pb.ReachFrequencyOptimization proto
+result = processor.execute([spec])
+
+print(f"R&F Optimization results for {spec.optimization_name}:")
+# Access results from the proto, e.g.:
+# result.results[0].optimized_channel_frequencies
+# result.results[0].optimized_marketing_analysis
+# result.results[0].frequency_outcome_grid
+```
+
+Note: You can provide the processor with multiple specs. This would result in
+a `ReachFrequencyOptimization` output with multiple results therein.
+"""
+
+from collections.abc import Sequence
+import dataclasses
+
+from meridian import backend
+from meridian import constants
+from mmm.v1 import mmm_pb2 as pb
+from mmm.v1.common import kpi_type_pb2 as kpi_type_pb
+from mmm.v1.marketing.analysis import marketing_analysis_pb2 as analysis_pb
+from mmm.v1.marketing.analysis import media_analysis_pb2 as media_analysis_pb
+from mmm.v1.marketing.analysis import outcome_pb2 as outcome_pb
+from mmm.v1.marketing.analysis import response_curve_pb2
+from mmm.v1.marketing.optimization import constraints_pb2 as constraints_pb
+from mmm.v1.marketing.optimization import reach_frequency_optimization_pb2 as rf_pb
+from schema.processors import common
+from schema.processors import model_processor
+from schema.utils import time_record
+import numpy as np
+import xarray as xr
+
+
+__all__ = [
+    "ReachFrequencyOptimizationSpec",
+    "ReachFrequencyOptimizationProcessor",
+]
+
+
+_STEP_SIZE_DECIMAL_PRECISION = 1
+_STEP_SIZE = _STEP_SIZE_DECIMAL_PRECISION / 10
+_TOL = 1e-6
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class ReachFrequencyOptimizationSpec(model_processor.OptimizationSpec):
+  """Spec dataclass for marketing reach and frequency optimization processor.
+
+  A frequency grid is generated using the range `[rounded_min_frequency,
+  rounded_max_frequency]` and a step size of `STEP_SIZE=0.1`.
+  `rounded_min_frequency` and `rounded_max_frequency` are rounded to the
+  nearest multiple of `STEP_SIZE`.
+
+  This spec is used both as user input to inform the R&F optimization processor
+  of its constraints and parameters, as well as an output structure that is
+  serializable to a `ReachFrequencyOptimizationSpec` proto. The latter serves as
+  a metadata embedded in a `ReachFrequencyOptimizationResult`. The output spec
+  in the proto reflects the actual numbers used to generate the reach and
+  frequency optimization result.
+
+  Attributes:
+    min_frequency: The minimum frequency constraint for each channel. Must be
+      greater than or equal to `1.0`. Defaults to `1.0`.
+    max_frequency: The maximum frequency constraint for each channel. Must be
+      greater than min_frequency. Defaults to None. If this value is set to
+      None, the model's max frequency will be used.
+    rf_channels: The R&F media channels in the model. When resolved with a
+      model, the model's R&F channels will be present here. Ignored when used as
+      input.
+    kpi_type: A `common.KpiType` enum denoting whether the optimized KPI is of a
+      `'revenue'` or `'non-revenue'` type.
+  """
+
+  min_frequency: float = 1.0
+  max_frequency: float | None = None
+  rf_channels: Sequence[str] = dataclasses.field(default_factory=list)
+  kpi_type: common.KpiType = common.KpiType.REVENUE
+
+  @property
+  def selected_times(self) -> tuple[str | None, str | None] | None:
+    """The start and end dates, as a tuple of date strings."""
+    start, end = (None, None)
+    if self.start_date is not None:
+      start = self.start_date.strftime(constants.DATE_FORMAT)
+    if self.end_date is not None:
+      end = self.end_date.strftime(constants.DATE_FORMAT)
+
+    if start or end:
+      return (start, end)
+    return None
+
+  @property
+  def objective(self) -> common.TargetMetric:
+    """A Meridian budget optimization objective is always ROI."""
+    return common.TargetMetric.ROI
+
+  def validate(self):
+    super().validate()
+    if self.min_frequency < 0:
+      raise ValueError("Min frequency must be non-negative.")
+    if (
+        self.max_frequency is not None
+        and self.max_frequency < self.min_frequency
+    ):
+      raise ValueError("Max frequency must be greater than min frequency.")
+
+  def to_proto(self) -> rf_pb.ReachFrequencyOptimizationSpec:
+    # When invoked as an output proto, the spec should have been fully resolved.
+    if self.start_date is None or self.end_date is None:
+      raise ValueError(
+          "Start and end dates must be resolved before this spec can be"
+          " serialized."
+      )
+
+    return rf_pb.ReachFrequencyOptimizationSpec(
+        date_interval=time_record.create_date_interval_pb(
+            self.start_date, self.end_date, tag=self.date_interval_tag
+        ),
+        rf_channel_constraints=[
+            rf_pb.RfChannelConstraint(
+                channel_name=channel,
+                frequency_constraint=constraints_pb.FrequencyConstraint(
+                    min_frequency=self.min_frequency,
+                    max_frequency=self.max_frequency,
+                ),
+            )
+            for channel in self.rf_channels
+        ],
+        objective=self.objective.value,
+        kpi_type=(
+            kpi_type_pb.KpiType.REVENUE
+            if self.kpi_type == common.KpiType.REVENUE
+            else kpi_type_pb.KpiType.NON_REVENUE
+        ),
+    )
+
+
+class ReachFrequencyOptimizationProcessor(
+    model_processor.ModelProcessor[
+        ReachFrequencyOptimizationSpec, rf_pb.ReachFrequencyOptimization
+    ],
+):
+  """A Processor for marketing reach and frequency optimization."""
+
+  def __init__(
+      self,
+      trained_model: model_processor.ModelType,
+  ):
+    trained_model = model_processor.ensure_trained_model(trained_model)
+    self._internal_analyzer = trained_model.internal_analyzer
+    self._meridian = trained_model.mmm
+
+    if trained_model.mmm.input_data.rf_channel is None:
+      raise ValueError("RF channels must be set in the model.")
+
+    self._all_rf_channels = trained_model.mmm.input_data.rf_channel.data
+
+  @classmethod
+  def spec_type(cls) -> type[ReachFrequencyOptimizationSpec]:
+    return ReachFrequencyOptimizationSpec
+
+  @classmethod
+  def output_type(cls) -> type[rf_pb.ReachFrequencyOptimization]:
+    return rf_pb.ReachFrequencyOptimization
+
+  def _to_target_precision(self, value: float) -> float:
+    return round(value, _STEP_SIZE_DECIMAL_PRECISION)
+
+  def _set_output(
+      self, output: pb.Mmm, result: rf_pb.ReachFrequencyOptimization
+  ):
+    output.marketing_optimization.reach_frequency_optimization.CopyFrom(result)
+
+  def execute(
+      self, specs: Sequence[ReachFrequencyOptimizationSpec]
+  ) -> rf_pb.ReachFrequencyOptimization:
+    output = rf_pb.ReachFrequencyOptimization()
+
+    group_ids = [spec.group_id for spec in specs if spec.group_id]
+    if len(set(group_ids)) != len(group_ids):
+      raise ValueError(
+          "Specified group_id must be unique among the given group of specs."
+      )
+
+    for spec in specs:
+      selected_times = spec.resolver(
+          self._meridian.input_data.time_coordinates
+      ).resolve_to_enumerated_selected_times()
+
+      grid_min_freq = self._to_target_precision(spec.min_frequency)
+      # If the max frequency is not set, use the model's max frequency.
+      grid_max_freq = self._to_target_precision(
+          spec.max_frequency or np.max(self._meridian.rf_tensors.frequency)
+      )
+      grid = [
+          self._to_target_precision(f)
+          for f in np.arange(grid_min_freq, grid_max_freq + _TOL, _STEP_SIZE)
+      ]
+
+      # Note that the internal analyzer, like the budget optimizer, maximizes
+      # non-revenue KPI if input data is of non-revenue and the user selects
+      # `use_kpi=True`. Otherwise, it maximizes revenue KPI.
+      optimal_frequency = self._internal_analyzer.optimal_freq(
+          selected_times=selected_times,
+          confidence_level=spec.confidence_level,
+          freq_grid=grid,
+          use_kpi=(spec.kpi_type == common.KpiType.NON_REVENUE),
+          max_frequency=spec.max_frequency,
+      )
+      response_curve = self._internal_analyzer.response_curves(
+          confidence_level=spec.confidence_level,
+          selected_times=selected_times,
+          by_reach=False,
+          use_kpi=(spec.kpi_type == common.KpiType.NON_REVENUE),
+          use_optimal_frequency=True,
+      )
+
+      spend_data = self._compute_spend_data(selected_times=selected_times)
+
+      # Obtain the output spec.
+      start, end = spec.resolver(
+          self._meridian.input_data.time_coordinates
+      ).resolve_to_date_interval_open_end()
+
+      # Copy the current spec, and resolve its date interval as well as model-
+      # dependent parameters.
+      output_spec = dataclasses.replace(
+          spec,
+          rf_channels=self._all_rf_channels,
+          min_frequency=grid_min_freq,
+          max_frequency=grid_max_freq,
+          start_date=start,
+          end_date=end,
+      )
+
+      output.results.append(
+          self._to_reach_frequency_optimization_result(
+              output_spec,
+              optimal_frequency,
+              response_curve,
+              spend_data,
+          )
+      )
+    return output
+
+  def _compute_spend_data(
+      self, selected_times: list[str] | None = None
+  ) -> xr.Dataset:
+    aggregated_spends = self._internal_analyzer.get_historical_spend(
+        selected_times
+    )
+    aggregated_rf_spend = aggregated_spends.sel(
+        {constants.CHANNEL: self._all_rf_channels}
+    ).data
+    total_spend = np.sum(aggregated_spends.data)
+    pct_of_spend = 100.0 * aggregated_rf_spend / total_spend
+
+    xr_dims = (constants.RF_CHANNEL,)
+    xr_coords = {
+        constants.RF_CHANNEL: (
+            [constants.RF_CHANNEL],
+            list(self._all_rf_channels),
+        ),
+    }
+    xr_data_vars = {
+        constants.SPEND: (xr_dims, aggregated_rf_spend),
+        constants.PCT_OF_SPEND: (xr_dims, pct_of_spend),
+    }
+
+    return xr.Dataset(
+        data_vars=xr_data_vars,
+        coords=xr_coords,
+    )
+
+  def _to_reach_frequency_optimization_result(
+      self,
+      spec: ReachFrequencyOptimizationSpec,
+      optimal_frequency: xr.Dataset,
+      response_curve: xr.Dataset,
+      spend_data: xr.Dataset,
+  ) -> rf_pb.ReachFrequencyOptimizationResult:
+    """Converts given optimal frequency dataset to protobuf form."""
+    result = rf_pb.ReachFrequencyOptimizationResult(
+        name=spec.optimization_name,
+        spec=spec.to_proto(),
+        optimized_channel_frequencies=_create_optimized_channel_frequencies(
+            optimal_frequency
+        ),
+        optimized_marketing_analysis=self._to_marketing_analysis(
+            spec,
+            optimal_frequency,
+            response_curve,
+            spend_data,
+        ),
+        frequency_outcome_grid=self._create_frequency_outcome_grid(
+            optimal_frequency,
+            spec,
+        ),
+    )
+    if spec.group_id:
+      result.group_id = spec.group_id
+    return result
+
+  def _to_marketing_analysis(
+      self,
+      spec: ReachFrequencyOptimizationSpec,
+      optimal_frequency: xr.Dataset,
+      response_curve: xr.Dataset,
+      spend_data: xr.Dataset,
+  ) -> analysis_pb.MarketingAnalysis:
+    """Converts an optimal frequency dataset to a `MarketingAnalysis` proto."""
+    # `spec` should have been resolved with concrete date interval parameters.
+    assert spec.start_date is not None and spec.end_date is not None
+
+    optimized_marketing_analysis = analysis_pb.MarketingAnalysis(
+        date_interval=time_record.create_date_interval_pb(
+            start_date=spec.start_date,
+            end_date=spec.end_date,
+        ),
+    )
+
+    # Create a per-channel MediaAnalysis.
+    channels = optimal_frequency.coords[constants.RF_CHANNEL].data
+    for channel in channels:
+      channel_optimal_frequency = optimal_frequency.sel(rf_channel=channel)
+      channel_spend_data = spend_data.sel(rf_channel=channel)
+
+      # TODO Add non-media analyses.
+      channel_media_analysis = media_analysis_pb.MediaAnalysis(
+          channel_name=channel,
+          response_curve=_compute_response_curve(
+              response_curve,
+              channel,
+          ),
+          spend_info=media_analysis_pb.SpendInfo(
+              spend=channel_spend_data[constants.SPEND].data.item(),
+              spend_share=(
+                  channel_spend_data[constants.PCT_OF_SPEND].data.item()
+              ),
+          ),
+      )
+
+      # Output one outcome per channel: either revenue or non-revenue.
+      channel_media_analysis.media_outcomes.append(
+          _to_outcome(
+              channel_optimal_frequency,
+              is_revenue_kpi=optimal_frequency.attrs[constants.IS_REVENUE_KPI],
+          )
+      )
+
+      optimized_marketing_analysis.media_analyses.append(channel_media_analysis)
+
+    return optimized_marketing_analysis
+
+  def _create_frequency_outcome_grid(
+      self,
+      optimal_frequency_dataset: xr.Dataset,
+      spec: ReachFrequencyOptimizationSpec,
+  ) -> rf_pb.FrequencyOutcomeGrid:
+    """Creates a FrequencyOutcomeGrid proto."""
+    channel_cells = []
+    frequencies = optimal_frequency_dataset.coords[constants.FREQUENCY].data
+    channels = optimal_frequency_dataset.coords[constants.RF_CHANNEL].data
+    input_tensor_dims = "gtc"
+    output_tensor_dims = "c"
+
+    for channel in channels:
+      cells = []
+      for frequency in frequencies:
+        new_frequency = (
+            backend.ones_like(self._meridian.rf_tensors.frequency) * frequency
+        )
+        new_reach = (
+            self._meridian.rf_tensors.frequency
+            * self._meridian.rf_tensors.reach
+            / new_frequency
+        )
+        channel_mask = [c == channel for c in channels]
+        filtered_reach = backend.boolean_mask(new_reach, channel_mask, axis=2)
+        aggregated_reach = backend.einsum(
+            f"{input_tensor_dims}->...{output_tensor_dims}", filtered_reach
+        )
+        reach = aggregated_reach.numpy()[-1]
+
+        metric_data_array = optimal_frequency_dataset[constants.ROI].sel(
+            frequency=frequency, rf_channel=channel
+        )
+        outcome = common.to_estimate(metric_data_array, spec.confidence_level)
+
+        cell = rf_pb.FrequencyOutcomeGrid.Cell(
+            outcome=outcome,
+            reach_frequency=rf_pb.ReachFrequency(
+                reach=int(reach),
+                average_frequency=frequency,
+            ),
+        )
+        cells.append(cell)
+
+      channel_cell = rf_pb.FrequencyOutcomeGrid.ChannelCells(
+          channel_name=channel,
+          cells=cells,
+      )
+      channel_cells.append(channel_cell)
+
+    return rf_pb.FrequencyOutcomeGrid(
+        name=spec.grid_name,
+        frequency_step_size=_STEP_SIZE,
+        channel_cells=channel_cells,
+    )
+
+
+def _create_optimized_channel_frequencies(
+    optimal_frequency_dataset: xr.Dataset,
+) -> list[rf_pb.OptimizedChannelFrequency]:
+  """Creates an OptimizedChannelFrequency proto for each channel in the dataset."""
+  optimal_frequency_protos = []
+  optimal_frequency = optimal_frequency_dataset[constants.OPTIMAL_FREQUENCY]
+  channels = optimal_frequency.coords[constants.RF_CHANNEL].data
+
+  for channel in channels:
+    optimal_frequency_protos.append(
+        rf_pb.OptimizedChannelFrequency(
+            channel_name=channel,
+            optimal_average_frequency=optimal_frequency.sel(
+                rf_channel=channel
+            ).item(),
+        )
+    )
+  return optimal_frequency_protos
+
+
+def _to_outcome(
+    channel_optimal_frequency: xr.Dataset,
+    is_revenue_kpi: bool,
+) -> outcome_pb.Outcome:
+  """Returns an `Outcome` value for a given channel's optimized media analysis.
+
+  Args:
+    channel_optimal_frequency: A channel-selected dataset from
+      `Analyzer.optimal_freq()`.
+    is_revenue_kpi: Whether the KPI is revenue-based.
+  """
+  confidence_level = channel_optimal_frequency.attrs[constants.CONFIDENCE_LEVEL]
+  return outcome_pb.Outcome(
+      kpi_type=(
+          kpi_type_pb.REVENUE if is_revenue_kpi else kpi_type_pb.NON_REVENUE
+      ),
+      roi=common.to_estimate(
+          channel_optimal_frequency.optimized_roi, confidence_level
+      ),
+      marginal_roi=common.to_estimate(
+          channel_optimal_frequency.optimized_mroi_by_frequency,
+          confidence_level,
+      ),
+      cost_per_contribution=common.to_estimate(
+          channel_optimal_frequency.optimized_cpik,
+          confidence_level=confidence_level,
+      ),
+      contribution=outcome_pb.Contribution(
+          value=common.to_estimate(
+              channel_optimal_frequency.optimized_incremental_outcome,
+              confidence_level,
+          ),
+      ),
+      effectiveness=outcome_pb.Effectiveness(
+          media_unit=constants.IMPRESSIONS,
+          value=common.to_estimate(
+              channel_optimal_frequency.optimized_effectiveness,
+              confidence_level,
+          ),
+      ),
+  )
+
+
+def _compute_response_curve(
+    response_curve_dataset: xr.Dataset,
+    channel_name: str,
+) -> response_curve_pb2.ResponseCurve:
+  """Returns a ResponseCurve proto for the given channel.
+
+  Args:
+    response_curve_dataset: A dataset containing the data needed to generate a
+      response curve.
+    channel_name: The name of the channel to analyze.
+  """
+
+  spend_multiplier_list = response_curve_dataset.coords[
+      constants.SPEND_MULTIPLIER
+  ].data
+  response_points: list[response_curve_pb2.ResponsePoint] = []
+
+  for spend_multiplier in spend_multiplier_list:
+    spend = (
+        response_curve_dataset[constants.SPEND]
+        .sel(spend_multiplier=spend_multiplier, channel=channel_name)
+        .data.item()
+    )
+    incremental_outcome = (
+        response_curve_dataset[constants.INCREMENTAL_OUTCOME]
+        .sel(
+            spend_multiplier=spend_multiplier,
+            channel=channel_name,
+            metric=constants.MEAN,
+        )
+        .data.item()
+    )
+
+    response_point = response_curve_pb2.ResponsePoint(
+        input_value=spend,
+        incremental_kpi=incremental_outcome,
+    )
+    response_points.append(response_point)
+
+  return response_curve_pb2.ResponseCurve(
+      input_name=constants.SPEND,
+      response_points=response_points,
+  )

schema/serde/distribution.py

@@ -193,6 +193,8 @@ class DistributionSerde(
         return meridian_pb.TfpParameterValue(scalar_value=value)
       case int():
         return meridian_pb.TfpParameterValue(int_value=value)
+      # TODO: b/470407198 - case bool() has to be before int() because bool is a
+      # subtype of int.
       case bool():
         return meridian_pb.TfpParameterValue(bool_value=value)
       case str():
@@ -216,10 +218,6 @@ class DistributionSerde(
         return meridian_pb.TfpParameterValue(
             dict_value=meridian_pb.TfpParameterValue.Dict(value_map=dict_value)
         )
-      case backend.Tensor():
-        return meridian_pb.TfpParameterValue(
-            tensor_value=backend.make_tensor_proto(value)
-        )
       case backend.tfd.Distribution():
         return meridian_pb.TfpParameterValue(
             distribution_value=self._to_distribution_proto(value)
@@ -257,9 +255,16 @@ class DistributionSerde(
             f" {type(dist).__name__}, but not found in registry. Please"
             " add custom functions to registry when saving models."
         )
-
-    # Handle unsupported types.
-    raise TypeError(f"Unsupported type: {type(value)}, {value}")
+      case _:
+        # Handle unsupported types by attempting to convert to a tensor proto.
+        # This allows for more flexibility in handling types that are not
+        # explicitly handled above, such as numpy arrays or backend tensors.
+        try:
+          return meridian_pb.TfpParameterValue(
+              tensor_value=backend.make_tensor_proto(value)
+          )
+        except TypeError as e:
+          raise TypeError(f"Unsupported type: {type(value)}, {value!r}") from e
 
   def _from_distribution_proto(
       self,