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

schema/processors/model_processor.py

@@ -0,0 +1,415 @@
+# 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 common and base classes for processing trained Meridian model to an MMM schema."""
+
+import abc
+from collections.abc import Sequence
+import dataclasses
+import datetime
+import functools
+from typing import Generic, TypeVar
+
+from google.protobuf import message
+from meridian import constants as c
+from meridian.analysis import analyzer
+from meridian.analysis import optimizer
+from meridian.analysis import visualizer
+from meridian.data import time_coordinates as tc
+from meridian.model import model
+from mmm.v1 import mmm_pb2 as pb
+from mmm.v1.common import date_interval_pb2
+from schema.utils import time_record
+from typing_extensions import override
+
+
+__all__ = [
+    'ModelProcessor',
+    'TrainedModel',
+    'DatedSpec',
+    'DatedSpecResolver',
+    'OptimizationSpec',
+    'ensure_trained_model',
+]
+
+
+class TrainedModel(abc.ABC):
+  """Encapsulates a trained MMM model."""
+
+  def __init__(self, mmm: model.Meridian):
+    """Initializes the TrainedModel with a Meridian model.
+
+    Args:
+      mmm: A Meridian model that has been fitted (posterior samples drawn).
+
+    Raises:
+      ValueError: If the model has not been fitted (posterior samples drawn).
+    """
+    # Ideally, this could be encoded in the model type itself, and we won't need
+    # this extra runtime check.
+    if mmm.inference_data.prior is None or mmm.inference_data.posterior is None:  # pytype: disable=attribute-error
+      raise ValueError('MMM model has not been fitted.')
+    self._mmm = mmm
+
+  @property
+  def mmm(self) -> model.Meridian:
+    return self._mmm
+
+  @property
+  def time_coordinates(self) -> tc.TimeCoordinates:
+    return self._mmm.input_data.time_coordinates
+
+  @functools.cached_property
+  def internal_analyzer(self) -> analyzer.Analyzer:
+    """Returns an internal `Analyzer`  bound to this trained model."""
+    return analyzer.Analyzer(
+        model_context=self.mmm.model_context,
+        inference_data=self.mmm.inference_data,
+    )
+
+  @functools.cached_property
+  def internal_optimizer(self) -> optimizer.BudgetOptimizer:
+    """Returns an internal `BudgetOptimizer` bound to this trained model."""
+    return optimizer.BudgetOptimizer(self.mmm)
+
+  @functools.cached_property
+  def internal_model_diagnostics(self) -> visualizer.ModelDiagnostics:
+    """Returns an internal `ModelDiagnostics` bound to this trained model."""
+    return visualizer.ModelDiagnostics(self.mmm)
+
+
+ModelType = model.Meridian | TrainedModel
+
+
+def ensure_trained_model(model_input: ModelType) -> TrainedModel:
+  """Ensure the given model is a trained model, and wrap it in a TrainedModel."""
+  if isinstance(model_input, TrainedModel):
+    return model_input
+  return TrainedModel(model_input)
+
+
+class Spec(abc.ABC):
+  """Contains parameters needed for model-based analysis/optimization."""
+
+  @abc.abstractmethod
+  def validate(self):
+    """Checks whether each parameter in the Spec has a valid value.
+
+    Raises:
+      ValueError: If any parameter in the Spec has an invalid value.
+    """
+
+  def __post_init__(self):
+    self.validate()
+
+
+@dataclasses.dataclass(frozen=True)
+class DatedSpec(Spec):
+  """A spec with a `[start_date, end_date)` closed-open date range semantic.
+
+  Attrs:
+    start_date: The start date of the analysis/optimization. If left as `None`,
+      then this will eventually resolve to a model's first time coordinate.
+    end_date: The end date of the analysis/optimization. If left as `None`, then
+      this will eventually resolve to a model's last time coordinate. When
+      specified, this end date is exclusive.
+    date_interval_tag: An optional tag that identifies the date interval.
+  """
+
+  start_date: datetime.date | None = None
+  end_date: datetime.date | None = None
+  date_interval_tag: str = ''
+
+  @override
+  def validate(self):
+    """Overrides the Spec.validate() method to check that dates are valid."""
+    if (
+        self.start_date is not None
+        and self.end_date is not None
+        and self.start_date > self.end_date
+    ):
+      raise ValueError('Start date must be before end date.')
+
+  def resolver(
+      self, time_coordinates: tc.TimeCoordinates
+  ) -> 'DatedSpecResolver':
+    """Returns a date resolver for this spec, with the given Meridian model."""
+    return DatedSpecResolver(self, time_coordinates)
+
+
+class DatedSpecResolver:
+  """Resolves date parameters in specs based on a model's time coordinates."""
+
+  def __init__(self, spec: DatedSpec, time_coordinates: tc.TimeCoordinates):
+    self._spec = spec
+    self._time_coordinates = time_coordinates
+
+  @property
+  def _interval_days(self) -> int:
+    return self._time_coordinates.interval_days
+
+  @property
+  def time_coordinates(self) -> tc.TimeCoordinates:
+    return self._time_coordinates
+
+  def to_closed_date_interval_tuple(
+      self,
+  ) -> tuple[str | None, str | None]:
+    """Transforms given spec into a closed `[start, end]` date interval tuple.
+
+    For each of the bookends in the tuple, `None` value indicates a time
+    coordinate default (first or last time coordinate, respectively).
+
+    Returns:
+      A **closed** `[start, end]` date interval tuple.
+    """
+    start, end = (None, None)
+
+    if self._spec.start_date is not None:
+      start = self._spec.start_date.strftime(c.DATE_FORMAT)
+    if self._spec.end_date is not None:
+      inclusive_end_date = self._spec.end_date - datetime.timedelta(
+          days=self._interval_days
+      )
+      end = inclusive_end_date.strftime(c.DATE_FORMAT)
+
+    return (start, end)
+
+  def resolve_to_enumerated_selected_times(self) -> list[str] | None:
+    """Resolves the given spec into an enumerated list of time coordinates.
+
+    Returns:
+      An enumerated list of time coordinates, or None (semantic "All") if the
+      bound spec is also None.
+    """
+    start, end = self.to_closed_date_interval_tuple()
+    expanded = self._time_coordinates.expand_selected_time_dims(
+        start_date=start, end_date=end
+    )
+    if expanded is None:
+      return None
+    return [date.strftime(c.DATE_FORMAT) for date in expanded]
+
+  def resolve_to_bool_selected_times(self) -> list[bool] | None:
+    """Resolves the given spec into a list of booleans indicating selected times.
+
+    Returns:
+      A list of booleans indicating selected times, or None (semantic "All") if
+      the bound spec is also None.
+    """
+    selected_times = self.resolve_to_enumerated_selected_times()
+    if selected_times is None:
+      return None
+    return [
+        time in selected_times for time in self._time_coordinates.all_dates_str
+    ]
+
+  def collapse_to_date_interval_proto(self) -> date_interval_pb2.DateInterval:
+    """Collapses the given spec into a `DateInterval` proto.
+
+    If the spec's date range is unbounded, then the DateInterval proto will have
+    the semantic "All", and we resolve it by consulting the time coordinates of
+    the model bound to this resolver.
+
+    Note that the exclusive end date semantic will be preserved in the returned
+    proto.
+
+    Returns:
+      A `DateInterval` proto the represents the date interval specified by the
+      spec.
+    """
+    selected_times = self.resolve_to_enumerated_selected_times()
+    if selected_times is None:
+      start_date = self._time_coordinates.all_dates[0]
+      end_date = self._time_coordinates.all_dates[-1]
+    else:
+      normalized_selected_times = [
+          tc.normalize_date(date) for date in selected_times
+      ]
+      start_date = normalized_selected_times[0]
+      end_date = normalized_selected_times[-1]
+
+    # Adjust end_date to make it exclusive.
+    end_date += datetime.timedelta(days=self._interval_days)
+
+    return time_record.create_date_interval_pb(
+        start_date, end_date, tag=self._spec.date_interval_tag
+    )
+
+  def transform_to_date_interval_protos(
+      self,
+  ) -> list[date_interval_pb2.DateInterval]:
+    """Transforms the given spec into `DateInterval` protos.
+
+    If the spec's date range is unbounded, then the DateInterval proto will have
+    the semantic "All", and we resolve it by consulting the time coordinates of
+    the model bound to this resolver.
+
+    Note that the exclusive end date semantic will be preserved in the returned
+    proto.
+
+    Returns:
+      A list of `DateInterval` protos the represents the date intervals
+      specified by the spec.
+    """
+    selected_times = self.resolve_to_enumerated_selected_times()
+    if selected_times is None:
+      times_list = self._time_coordinates.all_dates
+    else:
+      times_list = [tc.normalize_date(date) for date in selected_times]
+
+    date_intervals = []
+    for start_date in times_list:
+      date_interval = time_record.create_date_interval_pb(
+          start_date=start_date,
+          end_date=start_date + datetime.timedelta(days=self._interval_days),
+          tag=self._spec.date_interval_tag,
+      )
+      date_intervals.append(date_interval)
+
+    return date_intervals
+
+  def resolve_to_date_interval_open_end(
+      self,
+  ) -> tuple[datetime.date, datetime.date]:
+    """Resolves given spec into an open-ended `[start, end)` date interval."""
+    start = self._spec.start_date or self._time_coordinates.all_dates[0]
+    end = self._spec.end_date
+    if end is None:
+      end = self._time_coordinates.all_dates[-1]
+      # Adjust `end` to make it exclusive, but only if we pulled it from the
+      # time coordinates.
+      end += datetime.timedelta(days=self._interval_days)
+    return (start, end)
+
+  def resolve_to_date_interval_proto(self) -> date_interval_pb2.DateInterval:
+    """Resolves the given spec into a fully specified `DateInterval` proto.
+
+    If either `start_date` or `end_date` is None in the bound spec, then we
+    resolve it by consulting the time coordinates of the model bound to this
+    resolver. They are resolved to the first and last time coordinates (plus
+    interval length), respectively.
+
+    Note that the exclusive end date semantic will be preserved in the returned
+    proto.
+
+    Returns:
+      A resolved `DateInterval` proto the represents the date interval specified
+      by the bound spec.
+    """
+    start_date, end_date = self.resolve_to_date_interval_open_end()
+    return time_record.create_date_interval_pb(
+        start_date, end_date, tag=self._spec.date_interval_tag
+    )
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class OptimizationSpec(DatedSpec):
+  """A dated spec for optimization.
+
+  Attrs:
+    optimization_name: The name of the optimization in this spec.
+    grid_name: The name of the optimization grid.
+    group_id: An optional group ID for linking related optimizations.
+    confidence_level: The threshold for computing confidence intervals. Defaults
+      to 0.9. Must be a number between 0 and 1.
+  """
+
+  optimization_name: str
+  grid_name: str
+  group_id: str | None = None
+  confidence_level: float = c.DEFAULT_CONFIDENCE_LEVEL
+
+  @override
+  def validate(self):
+    """Check optimization parameters are valid."""
+    super().validate()
+
+    if not self.optimization_name or self.optimization_name.isspace():
+      raise ValueError('Optimization name must not be empty or blank.')
+
+    if not self.grid_name or self.grid_name.isspace():
+      raise ValueError('Grid name must not be empty or blank.')
+
+    if self.confidence_level < 0 or self.confidence_level > 1:
+      raise ValueError('Confidence level must be between 0 and 1.')
+
+
+S = TypeVar('S', bound=Spec)
+M = TypeVar('M', bound=message.Message)
+
+
+class ModelProcessor(abc.ABC, Generic[S, M]):
+  """Performs model-based analysis or optimization."""
+
+  @classmethod
+  @abc.abstractmethod
+  def spec_type(cls) -> type[S]:
+    """Returns the concrete Spec type that this ModelProcessor operates on."""
+    raise NotImplementedError()
+
+  @classmethod
+  @abc.abstractmethod
+  def output_type(cls) -> type[M]:
+    """Returns the concrete output type that this ModelProcessor produces."""
+    raise NotImplementedError()
+
+  @abc.abstractmethod
+  def execute(self, specs: Sequence[S]) -> M:
+    """Runs an analysis/optimization on the model using the given specs.
+
+    Args:
+      specs: Sequence of Specs containing parameters needed for the
+        analysis/optimization. The specs must all be of the same type as
+        `self.spec_type()` for this processor
+
+    Returns:
+      A proto containing the results of the analysis/optimization.
+    """
+    raise NotImplementedError()
+
+  @abc.abstractmethod
+  def _set_output(self, output: pb.Mmm, result: M):
+    """Sets the output field in the given `MmmOutput` proto.
+
+    A model consumer that orchestrated this processor will indirectly call this
+    method (via `__call__`) to attach the output of `execute()` (a
+    processor-defined message `M`) into a partially built `MmmOutput` proto that
+    the model consumer manages.
+
+    Args:
+      output: The container output proto to which the given result message
+        should be attached.
+      result: An output of `execute()`.
+    """
+    raise NotImplementedError()
+
+  def __call__(self, specs: Sequence[S], output: pb.Mmm):
+    """Runs an analysis/optimization on the model using the given specs.
+
+    This also sets the appropriate output field in the given MmmOutput proto.
+
+    Args:
+      specs: Sequence of Specs containing parameters needed for the
+        analysis/optimization. The specs must all be of the same type as
+        `self.spec_type()` for this processor
+      output: The output proto to which the results of the analysis/optimization
+        should be attached.
+
+    Raises:
+      ValueError: If any spec is not of the same type as `self.spec_type()`.
+    """
+    if not all([isinstance(spec, self.spec_type()) for spec in specs]):
+      raise ValueError('Not all specs are of type %s' % self.spec_type())
+    self._set_output(output, self.execute(specs))