google-meridian 1.4.0__py3-none-any.whl → 1.5.1__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(
@@ -352,7 +357,9 @@ def save_meridian(
   if not _file_exists(os.path.dirname(file_path)):
     _make_dirs(os.path.dirname(file_path))
 
-  with _file_open(file_path, 'wb') as f:
+  mode = 'wb' if file_path.endswith('.binpb') else 'w'
+
+  with _file_open(file_path, mode) as f:
     # Creates an MmmKernel.
     serialized_kernel = MeridianSerde().serialize(
         mmm,
@@ -397,7 +404,9 @@ def load_meridian(
   Returns:
     Model object loaded from the file path.
   """
-  with _file_open(file_path, 'rb') as f:
+  mode = 'rb' if file_path.endswith('.binpb') else 'r'
+
+  with _file_open(file_path, mode) as f:
     if file_path.endswith('.binpb'):
       serialized_model = kernel_pb.MmmKernel.FromString(f.read())
     elif file_path.endswith('.textproto') or file_path.endswith('.txtpb'):

schema/utils/__init__.py

@@ -15,4 +15,5 @@
 """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

schema/utils/proto_enum_converter.py

@@ -0,0 +1,127 @@
+# 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.
+
+"""A generic class for converting between Protobuf enums and strings."""
+
+from typing import Generic, Type, TypeVar
+import warnings
+
+import bidict
+
+
+EnumType = TypeVar("EnumType")
+DefaultType = TypeVar("DefaultType")
+
+
+class ProtoEnumConverter(Generic[EnumType, DefaultType]):
+  """Class for converting between proto enums and strings."""
+
+  def __init__(
+      self,
+      enum_message: Type[EnumType],
+      enum_display_name: str,
+      mapping: bidict.bidict,
+      enum_unspecified: EnumType,
+      default_when_unspecified: DefaultType,
+  ):
+    """Initializes the ProtoEnumConverter.
+
+    Arguments:
+      enum_message: The proto enum message definition.
+      enum_display_name: The loggable proto enum message name.
+      mapping: The mapping between the proto enum name and the string
+        representation.
+      enum_unspecified: The enum value that corresponds to unspecified.
+      default_when_unspecified: The default value that should be returned when
+        the proto enum is unspecified.
+    """
+    self.enum_message = enum_message
+    self.enum_display_name = enum_display_name
+    self.mapping = mapping
+    self.enum_unspecified = enum_unspecified
+    self.default_when_unspecified = default_when_unspecified
+
+  def to_proto(self, string_value: str | None) -> EnumType:
+    """Converts a string to its corresponding proto enum.
+
+    Args:
+      string_value: The string to convert to a proto enum.
+
+    Returns:
+      The corresponding proto enum or enum_unspecified when the enum message
+      doesn't exist.
+
+    Raises:
+      ValueError when given string is not found in the mapping.
+    """
+    if string_value is None:
+      return self.enum_unspecified
+
+    proto_name = self.mapping.get(string_value)
+    if proto_name:
+      try:
+        return self.enum_message.Value(proto_name)
+      except ValueError:
+        warnings.warn(
+            "Invalid %s value: %s. Resolving to %s."
+            % (
+                self.enum_message.DESCRIPTOR.name,
+                string_value,
+                self.enum_message.Name(self.enum_unspecified),
+            )
+        )
+        return self.enum_unspecified
+    else:
+      raise ValueError(
+          f"Unmatched {self.enum_message.DESCRIPTOR.name} value:"
+          f" {string_value}."
+      )
+
+  def from_proto(self, proto_enum: EnumType) -> str | DefaultType:
+    """Converts a proto enum to its string representation.
+
+    Args:
+      proto_enum: The enum value to convert to its string representation
+
+    Returns:
+      The string representation of the given proto_enum or the default value
+      when the proto enum is unspecified.
+
+    Raises:
+      ValueError when given proto enum is not found in the mapping.
+    """
+    if proto_enum == self.enum_unspecified:
+      warnings.warn(
+          "%s is unspecified. Resolving to default: %s."
+          % (
+              self.enum_display_name,
+              self.enum_message.Name(self.enum_unspecified),
+          )
+      )
+      return self.default_when_unspecified
+
+    try:
+      proto_name = self.enum_message.Name(proto_enum)
+    except ValueError as e:
+      raise ValueError(
+          f"Invalid {self.enum_message.DESCRIPTOR.name} proto enum value:"
+          f" {proto_enum}."
+      ) from e
+
+    try:
+      return self.mapping.inv[proto_name]
+    except KeyError as e:
+      raise KeyError(
+          f"Protobuf enum name '{proto_name}' is not configured in the mapping."
+      ) from e