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

schema/serde/hyperparameters.py

@@ -16,12 +16,14 @@
 
 import warnings
 
+import bidict
 from meridian import backend
 from meridian import constants as c
 from meridian.model import spec
 from mmm.v1.model.meridian import meridian_model_pb2 as meridian_pb
 from schema.serde import constants as sc
 from schema.serde import serde
+from schema.utils import proto_enum_converter
 import numpy as np
 
 _MediaEffectsDist = meridian_pb.MediaEffectsDistribution
@@ -31,101 +33,44 @@ _NonMediaBaselineFunction = (
     meridian_pb.NonMediaBaselineValue.NonMediaBaselineFunction
 )
 
+media_effects_converter = proto_enum_converter.ProtoEnumConverter(
+    enum_display_name="Media effects distribution",
+    enum_message=_MediaEffectsDist,
+    mapping=bidict.bidict({
+        c.MEDIA_EFFECTS_LOG_NORMAL: "LOG_NORMAL",
+        c.MEDIA_EFFECTS_NORMAL: "NORMAL",
+    }),
+    enum_unspecified=_MediaEffectsDist.MEDIA_EFFECTS_DISTRIBUTION_UNSPECIFIED,
+    default_when_unspecified=c.MEDIA_EFFECTS_LOG_NORMAL,
+)
 
-def _media_effects_dist_to_proto_enum(
-    media_effect_dict: str,
-) -> _MediaEffectsDist:
-  match media_effect_dict:
-    case c.MEDIA_EFFECTS_LOG_NORMAL:
-      return _MediaEffectsDist.LOG_NORMAL
-    case c.MEDIA_EFFECTS_NORMAL:
-      return _MediaEffectsDist.NORMAL
-    case _:
-      return _MediaEffectsDist.MEDIA_EFFECTS_DISTRIBUTION_UNSPECIFIED
-
-
-def _proto_enum_to_media_effects_dist(
-    proto_enum: _MediaEffectsDist,
-) -> str:
-  """Converts a `_MediaEffectsDist` enum to its string representation."""
-  match proto_enum:
-    case _MediaEffectsDist.LOG_NORMAL:
-      return c.MEDIA_EFFECTS_LOG_NORMAL
-    case _MediaEffectsDist.NORMAL:
-      return c.MEDIA_EFFECTS_NORMAL
-    case _MediaEffectsDist.MEDIA_EFFECTS_DISTRIBUTION_UNSPECIFIED:
-      warnings.warn(
-          "Media effects distribution is unspecified. Resolving to"
-          " 'log-normal'."
-      )
-      return c.MEDIA_EFFECTS_LOG_NORMAL
-    case _:
-      raise ValueError(
-          "Unsupported MediaEffectsDistribution proto enum value:"
-          f" {proto_enum}."
-      )
-
-
-def _paid_media_prior_type_to_proto_enum(
-    paid_media_prior_type: str | None,
-) -> _PaidMediaPriorType:
-  """Converts a paid media prior type string to its proto enum."""
-  if paid_media_prior_type is None:
-    return _PaidMediaPriorType.PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED
-  try:
-    return _PaidMediaPriorType.Value(paid_media_prior_type.upper())
-  except ValueError:
-    warnings.warn(
-        f"Invalid paid media prior type: {paid_media_prior_type}. Resolving to"
-        " PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED."
-    )
-    return _PaidMediaPriorType.PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED
-
-
-def _proto_enum_to_paid_media_prior_type(
-    proto_enum: _PaidMediaPriorType,
-) -> str | None:
-  """Converts a `_PaidMediaPriorType` enum to its string representation."""
-  if proto_enum == _PaidMediaPriorType.PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED:
-    return None
-  return _PaidMediaPriorType.Name(proto_enum).lower()
-
-
-def _non_paid_prior_type_to_proto_enum(
-    non_paid_prior_type: str,
-) -> _NonPaidTreatmentsPriorType:
-  """Converts a non-paid prior type string to its proto enum."""
-  try:
-    return _NonPaidTreatmentsPriorType.Value(
-        f"NON_PAID_TREATMENTS_PRIOR_TYPE_{non_paid_prior_type.upper()}"
-    )
-  except ValueError:
-    warnings.warn(
-        f"Invalid non-paid prior type: {non_paid_prior_type}. Resolving to"
-        " NON_PAID_TREATMENTS_PRIOR_TYPE_CONTRIBUTION."
-    )
-    return (
-        _NonPaidTreatmentsPriorType.NON_PAID_TREATMENTS_PRIOR_TYPE_CONTRIBUTION
-    )
-
+paid_media_prior_type_converter = proto_enum_converter.ProtoEnumConverter(
+    enum_display_name="Paid media prior type",
+    enum_message=_PaidMediaPriorType,
+    mapping=bidict.bidict({
+        c.TREATMENT_PRIOR_TYPE_ROI: "ROI",
+        c.TREATMENT_PRIOR_TYPE_MROI: "MROI",
+        c.TREATMENT_PRIOR_TYPE_COEFFICIENT: "COEFFICIENT",
+        c.TREATMENT_PRIOR_TYPE_CONTRIBUTION: "CONTRIBUTION",
+    }),
+    enum_unspecified=_PaidMediaPriorType.PAID_MEDIA_PRIOR_TYPE_UNSPECIFIED,
+    default_when_unspecified=None,
+)
 
-def _proto_enum_to_non_paid_prior_type(
-    proto_enum: _NonPaidTreatmentsPriorType,
-) -> str:
-  """Converts a `_NonPaidTreatmentsPriorType` enum to its string representation."""
-  if (
-      proto_enum
-      == _NonPaidTreatmentsPriorType.NON_PAID_TREATMENTS_PRIOR_TYPE_UNSPECIFIED
-  ):
-    warnings.warn(
-        "Non-paid prior type is unspecified. Resolving to 'contribution'."
-    )
-    return c.TREATMENT_PRIOR_TYPE_CONTRIBUTION
-  return (
-      _NonPaidTreatmentsPriorType.Name(proto_enum)
-      .replace("NON_PAID_TREATMENTS_PRIOR_TYPE_", "")
-      .lower()
-  )
+non_paid_treatments_prior_type_converter = proto_enum_converter.ProtoEnumConverter(
+    enum_display_name="Non-paid treatments prior type",
+    enum_message=_NonPaidTreatmentsPriorType,
+    mapping=bidict.bidict({
+        c.TREATMENT_PRIOR_TYPE_COEFFICIENT: (
+            "NON_PAID_TREATMENTS_PRIOR_TYPE_COEFFICIENT"
+        ),
+        c.TREATMENT_PRIOR_TYPE_CONTRIBUTION: (
+            "NON_PAID_TREATMENTS_PRIOR_TYPE_CONTRIBUTION"
+        ),
+    }),
+    enum_unspecified=_NonPaidTreatmentsPriorType.NON_PAID_TREATMENTS_PRIOR_TYPE_UNSPECIFIED,
+    default_when_unspecified=c.TREATMENT_PRIOR_TYPE_CONTRIBUTION,
+)
 
 
 class HyperparametersSerde(
@@ -141,25 +86,27 @@ class HyperparametersSerde(
   def serialize(self, obj: spec.ModelSpec) -> meridian_pb.Hyperparameters:
     """Serializes the given ModelSpec into a `Hyperparameters` proto."""
     hyperparameters_proto = meridian_pb.Hyperparameters(
-        media_effects_dist=_media_effects_dist_to_proto_enum(
+        media_effects_dist=media_effects_converter.to_proto(
             obj.media_effects_dist
         ),
         hill_before_adstock=obj.hill_before_adstock,
         unique_sigma_for_each_geo=obj.unique_sigma_for_each_geo,
-        media_prior_type=_paid_media_prior_type_to_proto_enum(
+        media_prior_type=paid_media_prior_type_converter.to_proto(
             obj.media_prior_type
         ),
-        rf_prior_type=_paid_media_prior_type_to_proto_enum(obj.rf_prior_type),
-        paid_media_prior_type=_paid_media_prior_type_to_proto_enum(
+        rf_prior_type=paid_media_prior_type_converter.to_proto(
+            obj.rf_prior_type
+        ),
+        paid_media_prior_type=paid_media_prior_type_converter.to_proto(
             obj.paid_media_prior_type
         ),
-        organic_media_prior_type=_non_paid_prior_type_to_proto_enum(
+        organic_media_prior_type=non_paid_treatments_prior_type_converter.to_proto(
             obj.organic_media_prior_type
         ),
-        organic_rf_prior_type=_non_paid_prior_type_to_proto_enum(
+        organic_rf_prior_type=non_paid_treatments_prior_type_converter.to_proto(
             obj.organic_rf_prior_type
         ),
-        non_media_treatments_prior_type=_non_paid_prior_type_to_proto_enum(
+        non_media_treatments_prior_type=non_paid_treatments_prior_type_converter.to_proto(
             obj.non_media_treatments_prior_type
         ),
         enable_aks=obj.enable_aks,
@@ -326,28 +273,28 @@ class HyperparametersSerde(
       adstock_decay_spec = sc.DEFAULT_DECAY
 
     return spec.ModelSpec(
-        media_effects_dist=_proto_enum_to_media_effects_dist(
+        media_effects_dist=media_effects_converter.from_proto(
             serialized.media_effects_dist
         ),
         hill_before_adstock=serialized.hill_before_adstock,
         max_lag=max_lag,
         unique_sigma_for_each_geo=serialized.unique_sigma_for_each_geo,
-        media_prior_type=_proto_enum_to_paid_media_prior_type(
+        media_prior_type=paid_media_prior_type_converter.from_proto(
             serialized.media_prior_type
         ),
-        rf_prior_type=_proto_enum_to_paid_media_prior_type(
+        rf_prior_type=paid_media_prior_type_converter.from_proto(
             serialized.rf_prior_type
         ),
-        paid_media_prior_type=_proto_enum_to_paid_media_prior_type(
+        paid_media_prior_type=paid_media_prior_type_converter.from_proto(
             serialized.paid_media_prior_type
         ),
-        organic_media_prior_type=_proto_enum_to_non_paid_prior_type(
+        organic_media_prior_type=non_paid_treatments_prior_type_converter.from_proto(
             serialized.organic_media_prior_type
         ),
-        organic_rf_prior_type=_proto_enum_to_non_paid_prior_type(
+        organic_rf_prior_type=non_paid_treatments_prior_type_converter.from_proto(
             serialized.organic_rf_prior_type
         ),
-        non_media_treatments_prior_type=_proto_enum_to_non_paid_prior_type(
+        non_media_treatments_prior_type=non_paid_treatments_prior_type_converter.from_proto(
             serialized.non_media_treatments_prior_type
         ),
         non_media_baseline_values=non_media_baseline_values,

schema/serde/meridian_serde.py

@@ -43,6 +43,7 @@ import dataclasses
 import os
 import warnings
 
+import arviz as az
 from google.protobuf import text_format
 import meridian
 from meridian import backend
@@ -165,6 +166,7 @@ class MeridianSerde(serde.Serde[kernel_pb.MmmKernel, model.Meridian]):
         inference_data=inference_data.InferenceDataSerde().serialize(
             mmm.inference_data
         ),
+        arviz_version=az.__version__,
     )
     # For backwards compatibility, only serialize EDA spec if it exists.
     if hasattr(mmm, 'eda_spec'):
@@ -190,7 +192,10 @@ class MeridianSerde(serde.Serde[kernel_pb.MmmKernel, model.Meridian]):
       # NotFittedModelError can be raised below. If raised,
       # return None. Otherwise, set convergence status based on
       # MCMCSamplingError (caught in the except block).
-      rhats = analyzer.Analyzer(mmm).get_rhat()
+      rhats = analyzer.Analyzer(
+          model_context=mmm.model_context,
+          inference_data=mmm.inference_data,
+      ).get_rhat()
       rhat_proto = meridian_pb.RHatDiagnostic()
       for name, tensor in rhats.items():
         rhat_proto.parameter_r_hats.add(

schema/test_data.py

@@ -0,0 +1,380 @@
+# 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.
+
+"""Test data for MMM proto generator."""
+
+from collections.abc import Sequence
+import datetime
+
+from mmm.v1 import mmm_pb2 as mmm_pb
+from mmm.v1.common import date_interval_pb2 as date_interval_pb
+from mmm.v1.fit import model_fit_pb2 as fit_pb
+from mmm.v1.marketing.analysis import marketing_analysis_pb2
+from mmm.v1.marketing.optimization import budget_optimization_pb2 as budget_pb
+from mmm.v1.marketing.optimization import reach_frequency_optimization_pb2 as rf_pb
+from schema.processors import budget_optimization_processor
+from schema.processors import marketing_processor
+from schema.processors import model_fit_processor
+from schema.processors import model_processor
+from schema.processors import reach_frequency_optimization_processor as rf_opt_processor
+
+from google.type import date_pb2
+
+# Weekly dates from 2022-11-21 to 2024-01-01.
+ALL_TIMES_IN_MERIDIAN = (
+    '2022-11-21',
+    '2022-11-28',
+    '2022-12-05',
+    '2022-12-12',
+    '2022-12-19',
+    '2022-12-26',
+    '2023-01-02',
+    '2023-01-09',
+    '2023-01-16',
+    '2023-01-23',
+    '2023-01-30',
+    '2023-02-06',
+    '2023-02-13',
+    '2023-02-20',
+    '2023-02-27',
+    '2023-03-06',
+    '2023-03-13',
+    '2023-03-20',
+    '2023-03-27',
+    '2023-04-03',
+    '2023-04-10',
+    '2023-04-17',
+    '2023-04-24',
+    '2023-05-01',
+    '2023-05-08',
+    '2023-05-15',
+    '2023-05-22',
+    '2023-05-29',
+    '2023-06-05',
+    '2023-06-12',
+    '2023-06-19',
+    '2023-06-26',
+    '2023-07-03',
+    '2023-07-10',
+    '2023-07-17',
+    '2023-07-24',
+    '2023-07-31',
+    '2023-08-07',
+    '2023-08-14',
+    '2023-08-21',
+    '2023-08-28',
+    '2023-09-04',
+    '2023-09-11',
+    '2023-09-18',
+    '2023-09-25',
+    '2023-10-02',
+    '2023-10-09',
+    '2023-10-16',
+    '2023-10-23',
+    '2023-10-30',
+    '2023-11-06',
+    '2023-11-13',
+    '2023-11-20',
+    '2023-11-27',
+    '2023-12-04',
+    '2023-12-11',
+    '2023-12-18',
+    '2023-12-25',
+    '2024-01-01',
+)
+
+ALL_TIME_BUCKET_DATED_SPECS = (
+    # All
+    model_processor.DatedSpec(
+        start_date=datetime.date(2022, 11, 21),
+        end_date=datetime.date(2024, 1, 8),
+        date_interval_tag='ALL',
+    ),
+    # Monthly buckets
+    model_processor.DatedSpec(
+        start_date=datetime.date(2022, 12, 5),
+        end_date=datetime.date(2023, 1, 2),
+        date_interval_tag='Y2022 Dec',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 1, 2),
+        end_date=datetime.date(2023, 2, 6),
+        date_interval_tag='Y2023 Jan',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 2, 6),
+        end_date=datetime.date(2023, 3, 6),
+        date_interval_tag='Y2023 Feb',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 3, 6),
+        end_date=datetime.date(2023, 4, 3),
+        date_interval_tag='Y2023 Mar',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 4, 3),
+        end_date=datetime.date(2023, 5, 1),
+        date_interval_tag='Y2023 Apr',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 5, 1),
+        end_date=datetime.date(2023, 6, 5),
+        date_interval_tag='Y2023 May',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 6, 5),
+        end_date=datetime.date(2023, 7, 3),
+        date_interval_tag='Y2023 Jun',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 7, 3),
+        end_date=datetime.date(2023, 8, 7),
+        date_interval_tag='Y2023 Jul',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 8, 7),
+        end_date=datetime.date(2023, 9, 4),
+        date_interval_tag='Y2023 Aug',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 9, 4),
+        end_date=datetime.date(2023, 10, 2),
+        date_interval_tag='Y2023 Sep',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 10, 2),
+        end_date=datetime.date(2023, 11, 6),
+        date_interval_tag='Y2023 Oct',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 11, 6),
+        end_date=datetime.date(2023, 12, 4),
+        date_interval_tag='Y2023 Nov',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 12, 4),
+        end_date=datetime.date(2024, 1, 1),
+        date_interval_tag='Y2023 Dec',
+    ),
+    # Quarterly buckets
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 1, 2),
+        end_date=datetime.date(2023, 4, 3),
+        date_interval_tag='Y2023 Q1',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 4, 3),
+        end_date=datetime.date(2023, 7, 3),
+        date_interval_tag='Y2023 Q2',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 7, 3),
+        end_date=datetime.date(2023, 10, 2),
+        date_interval_tag='Y2023 Q3',
+    ),
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 10, 2),
+        end_date=datetime.date(2024, 1, 1),
+        date_interval_tag='Y2023 Q4',
+    ),
+    # Yearly buckets
+    model_processor.DatedSpec(
+        start_date=datetime.date(2023, 1, 2),
+        end_date=datetime.date(2024, 1, 1),
+        date_interval_tag='Y2023',
+    ),
+)
+
+
+def _dated_spec_to_date_interval(
+    spec: model_processor.DatedSpec,
+) -> date_interval_pb.DateInterval:
+  if spec.start_date is None or spec.end_date is None:
+    raise ValueError('Start date or end date is None.')
+
+  return date_interval_pb.DateInterval(
+      start_date=date_pb2.Date(
+          year=spec.start_date.year,
+          month=spec.start_date.month,
+          day=spec.start_date.day,
+      ),
+      end_date=date_pb2.Date(
+          year=spec.end_date.year,
+          month=spec.end_date.month,
+          day=spec.end_date.day,
+      ),
+      tag=spec.date_interval_tag,
+  )
+
+
+class FakeModelFitProcessor(
+    model_processor.ModelProcessor[
+        model_fit_processor.ModelFitSpec, fit_pb.ModelFit
+    ]
+):
+  """Fake ModelFitProcessor for testing."""
+
+  def __init__(self, trained_model: model_processor.TrainedModel):
+    self._trained_model = trained_model
+
+  @classmethod
+  def spec_type(cls):
+    return model_fit_processor.ModelFitSpec
+
+  @classmethod
+  def output_type(cls):
+    return fit_pb.ModelFit
+
+  def execute(
+      self, specs: Sequence[model_fit_processor.ModelFitSpec]
+  ) -> fit_pb.ModelFit:
+    return fit_pb.ModelFit()
+
+  def _set_output(self, output: mmm_pb.Mmm, result: fit_pb.ModelFit):
+    output.model_fit.CopyFrom(result)
+
+
+class FakeBudgetOptimizationProcessor(
+    model_processor.ModelProcessor[
+        budget_optimization_processor.BudgetOptimizationSpec,
+        budget_pb.BudgetOptimization,
+    ]
+):
+  """Fake BudgetOptimizationProcessor for testing."""
+
+  def __init__(self, trained_model: model_processor.TrainedModel):
+    self._trained_model = trained_model
+
+  @classmethod
+  def spec_type(cls):
+    return budget_optimization_processor.BudgetOptimizationSpec
+
+  @classmethod
+  def output_type(cls):
+    return budget_pb.BudgetOptimization
+
+  def execute(
+      self,
+      specs: Sequence[budget_optimization_processor.BudgetOptimizationSpec],
+  ) -> budget_pb.BudgetOptimization:
+    results = []
+    for spec in specs:
+      result = budget_pb.BudgetOptimizationResult(
+          name=spec.optimization_name,
+          spec=budget_pb.BudgetOptimizationSpec(
+              date_interval=_dated_spec_to_date_interval(spec)
+          ),
+          incremental_outcome_grid=budget_pb.IncrementalOutcomeGrid(
+              name=spec.grid_name
+          ),
+      )
+      if spec.group_id:
+        result.group_id = spec.group_id
+      results.append(result)
+
+    return budget_pb.BudgetOptimization(results=results)
+
+  def _set_output(
+      self, output: mmm_pb.Mmm, result: budget_pb.BudgetOptimization
+  ):
+    output.marketing_optimization.budget_optimization.CopyFrom(result)
+
+
+class FakeReachFrequencyOptimizationProcessor(
+    model_processor.ModelProcessor[
+        rf_opt_processor.ReachFrequencyOptimizationSpec,
+        rf_pb.ReachFrequencyOptimization,
+    ]
+):
+  """Fake ReachFrequencyOptimizationProcessor for testing."""
+
+  def __init__(self, trained_model: model_processor.TrainedModel):
+    self._trained_model = trained_model
+
+  @classmethod
+  def spec_type(cls):
+    return rf_opt_processor.ReachFrequencyOptimizationSpec
+
+  @classmethod
+  def output_type(cls):
+    return rf_pb.ReachFrequencyOptimization
+
+  def execute(
+      self,
+      specs: Sequence[rf_opt_processor.ReachFrequencyOptimizationSpec],
+  ) -> rf_pb.ReachFrequencyOptimization:
+    results = []
+    for spec in specs:
+      result = rf_pb.ReachFrequencyOptimizationResult(
+          name=spec.optimization_name,
+          spec=rf_pb.ReachFrequencyOptimizationSpec(
+              date_interval=_dated_spec_to_date_interval(spec)
+          ),
+          frequency_outcome_grid=rf_pb.FrequencyOutcomeGrid(
+              name=spec.grid_name
+          ),
+      )
+      if spec.group_id:
+        result.group_id = spec.group_id
+      results.append(result)
+
+    return rf_pb.ReachFrequencyOptimization(results=results)
+
+  def _set_output(
+      self,
+      output: mmm_pb.Mmm,
+      result: rf_pb.ReachFrequencyOptimization,
+  ):
+    output.marketing_optimization.reach_frequency_optimization.CopyFrom(result)
+
+
+class FakeMarketingProcessor(
+    model_processor.ModelProcessor[
+        marketing_processor.MarketingAnalysisSpec,
+        marketing_analysis_pb2.MarketingAnalysisList,
+    ]
+):
+  """Fake MarketingProcessor for testing."""
+
+  def __init__(self, trained_model: model_processor.TrainedModel):
+    self._trained_model = trained_model
+
+  @classmethod
+  def spec_type(cls):
+    return marketing_processor.MarketingAnalysisSpec
+
+  @classmethod
+  def output_type(cls):
+    return marketing_analysis_pb2.MarketingAnalysisList
+
+  def execute(
+      self, specs: Sequence[marketing_processor.MarketingAnalysisSpec]
+  ) -> marketing_analysis_pb2.MarketingAnalysisList:
+    marketing_analyses = []
+    for spec in specs:
+      marketing_analysis = marketing_analysis_pb2.MarketingAnalysis(
+          date_interval=_dated_spec_to_date_interval(spec)
+      )
+      marketing_analyses.append(marketing_analysis)
+
+    return marketing_analysis_pb2.MarketingAnalysisList(
+        marketing_analyses=marketing_analyses
+    )
+
+  def _set_output(
+      self,
+      output: mmm_pb.Mmm,
+      result: marketing_analysis_pb2.MarketingAnalysisList,
+  ):
+    output.marketing_analysis_list.CopyFrom(result)

schema/utils/__init__.py

@@ -14,4 +14,6 @@
 
 """Module containing MMM schema util functions."""
 
+from schema.utils import date_range_bucketing
+from schema.utils import proto_enum_converter
 from schema.utils import time_record