google-meridian 1.2.1__py3-none-any.whl → 1.3.1__py3-none-any.whl

meridian/backend/test_utils.py

@@ -15,11 +15,21 @@
 """Common testing utilities for Meridian, designed to be backend-agnostic."""
 
 from typing import Any, Optional
+
 from absl.testing import parameterized
+from google.protobuf import descriptor
+from google.protobuf import message
 from meridian import backend
 from meridian.backend import config
 import numpy as np
 
+from tensorflow.python.util.protobuf import compare
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.core.framework import tensor_pb2
+# pylint: enable=g-direct-tensorflow-import
+
+FieldDescriptor = descriptor.FieldDescriptor
+
 # A type alias for backend-agnostic array-like objects.
 # We use `Any` here to avoid circular dependencies with the backend module
 # while still allowing the function to accept backend-specific tensor types.
@@ -70,6 +80,39 @@ def assert_allequal(a: ArrayLike, b: ArrayLike, err_msg: str = ""):
   np.testing.assert_array_equal(np.array(a), np.array(b), err_msg=err_msg)
 
 
+def assert_seed_allequal(a: Any, b: Any, err_msg: str = ""):
+  """Backend-agnostic assertion to check if two seed objects are equal."""
+  data_a = backend.get_seed_data(a)
+  data_b = backend.get_seed_data(b)
+  if data_a is None and data_b is None:
+    return
+  np.testing.assert_array_equal(data_a, data_b, err_msg=err_msg)
+
+
+def assert_not_allequal(a: ArrayLike, b: ArrayLike, err_msg: str = ""):
+  """Asserts that two objects are not element-wise equal."""
+  np.testing.assert_(
+      not np.array_equal(np.array(a), np.array(b)),
+      msg=f"Arrays are unexpectedly equal.\n{err_msg}",
+  )
+
+
+def assert_seed_not_allequal(a: Any, b: Any, err_msg: str = ""):
+  """Asserts that two seed objects are not element-wise equal."""
+  data_a = backend.get_seed_data(a)
+  data_b = backend.get_seed_data(b)
+  if data_a is None and data_b is None:
+    raise AssertionError(
+        f"Seeds are unexpectedly equal (both are None). {err_msg}"
+    )
+  if data_a is None or data_b is None:
+    return
+  np.testing.assert_(
+      not np.array_equal(data_a, data_b),
+      msg=f"Seeds are unexpectedly equal.\n{err_msg}",
+  )
+
+
 def assert_all_finite(a: ArrayLike, err_msg: str = ""):
   """Backend-agnostic assertion to check if all elements in an array are finite.
 
@@ -98,6 +141,118 @@ def assert_all_non_negative(a: ArrayLike, err_msg: str = ""):
     raise AssertionError(err_msg or "Array contains negative values.")
 
 
+# --- Proto Utilities ---
+def normalize_tensor_protos(proto: message.Message):
+  """Recursively normalizes TensorProto messages within a proto (In-place).
+
+  This ensures a consistent serialization format across different backends
+  (e.g., JAX vs TF) by repacking TensorProtos using the current backend's
+  canonical method (backend.make_tensor_proto). This handles differences
+  like using `bool_val` versus `tensor_content` for boolean tensors.
+
+  Args:
+    proto: The protobuf message object to normalize. This object is modified in
+      place.
+  """
+  if not isinstance(proto, message.Message):
+    return
+
+  for desc, value in proto.ListFields():
+    if desc.type != FieldDescriptor.TYPE_MESSAGE:
+      continue
+
+    # A map is defined as a repeated field whose message type has the
+    # map_entry option set.
+    is_map = (
+        desc.label == FieldDescriptor.LABEL_REPEATED
+        and desc.message_type.has_options
+        and desc.message_type.GetOptions().map_entry
+    )
+
+    if is_map:
+      for item in value.values():
+        # Helper checks if values are scalars or messages.
+        _process_message_for_normalization(item)
+
+    elif desc.label == FieldDescriptor.LABEL_REPEATED:
+      # Handle standard repeated message fields.
+      for item in value:
+        _process_message_for_normalization(item)
+    else:
+      # Handle singular message fields.
+      _process_message_for_normalization(value)
+
+
+def _process_message_for_normalization(msg: Any):
+  """Helper to process a potential message during normalization traversal."""
+  # Ensure we only process message objects.
+  # If msg is a scalar (e.g., string from map<string, string>), stop recursion.
+  if not isinstance(msg, message.Message):
+    return
+
+  if isinstance(msg, tensor_pb2.TensorProto):
+    _repack_tensor_proto(msg)
+  else:
+    # If it's another message type, recurse into its fields.
+    normalize_tensor_protos(msg)
+
+
+def _repack_tensor_proto(tensor_proto: "tensor_pb2.TensorProto"):
+  """Repacks a TensorProto in place to use a consistent serialization format."""
+  if not tensor_proto.ByteSize():
+    return
+
+  try:
+    data_array = backend.make_ndarray(tensor_proto)
+  except Exception as e:
+    raise ValueError(
+        "Failed to deserialize TensorProto during normalization:"
+        f" {e}\nProto content:\n{tensor_proto}"
+    ) from e
+
+  new_tensor_proto = backend.make_tensor_proto(data_array)
+
+  tensor_proto.Clear()
+  tensor_proto.CopyFrom(new_tensor_proto)
+
+
+def assert_normalized_proto_equal(
+    test_case: parameterized.TestCase,
+    expected: message.Message,
+    actual: message.Message,
+    msg: Optional[str] = None,
+    **kwargs: Any,
+):
+  """Compares two protos after normalizing TensorProto fields.
+
+  Use this instead of compare.assertProtoEqual when protos contain tensors
+  to ensure backend-agnostic comparison.
+
+  Args:
+    test_case: The TestCase instance (self).
+    expected: The expected protobuf message.
+    actual: The actual protobuf message.
+    msg: An optional message to display on failure.
+    **kwargs: Additional keyword arguments passed to assertProto2Equal (e.g.,
+      precision).
+  """
+  # Work on copies to avoid mutating the original objects
+  expected_copy = expected.__class__()
+  expected_copy.CopyFrom(expected)
+  actual_copy = actual.__class__()
+  actual_copy.CopyFrom(actual)
+
+  try:
+    normalize_tensor_protos(expected_copy)
+    normalize_tensor_protos(actual_copy)
+  except ValueError as e:
+    test_case.fail(f"Proto normalization failed: {e}. {msg}")
+
+  compare.assertProtoEqual(
+      test_case, expected_copy, actual_copy, msg=msg, **kwargs
+  )
+
+
 class MeridianTestCase(parameterized.TestCase):
   """Base test class for Meridian providing backend-aware utilities.
 
@@ -106,6 +261,13 @@ class MeridianTestCase(parameterized.TestCase):
   number generation across backends (Stateful TF vs Stateless JAX).
   """
 
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    # Enforce determinism for TensorFlow tests before any tests are run.
+    # This is a no-op with a warning for the JAX backend.
+    backend.enable_op_determinism()
+
   def setUp(self):
     super().setUp()
     # Default seed, can be overridden by subclasses before calling

meridian/constants.py

@@ -55,6 +55,7 @@ DATE_FORMAT = '%Y-%m-%d'
 QUARTER_FORMAT = '%Y %b'
 
 ORGANIC_PREFIX = 'organic_'
+NATIONAL_PREFIX = 'national_'
 
 # Input data variables.
 KPI = 'kpi'
@@ -123,7 +124,6 @@ RF_DATA = (
     RF_SPEND,
     REVENUE_PER_KPI,
 )
-NON_REVENUE_DATA = IMPRESSIONS_DATA + (CONTROLS,)
 
 # Scaled input data variables.
 MEDIA_SCALED = 'media_scaled'
@@ -132,6 +132,9 @@ ORGANIC_MEDIA_SCALED = ORGANIC_PREFIX + MEDIA_SCALED
 ORGANIC_REACH_SCALED = ORGANIC_PREFIX + REACH_SCALED
 NON_MEDIA_TREATMENTS_SCALED = 'non_media_treatments_scaled'
 CONTROLS_SCALED = 'controls_scaled'
+KPI_SCALED = f'{KPI}_scaled'
+POPULATION_SCALED_KPI = f'{POPULATION}_scaled_{KPI}'
+RF_IMPRESSIONS_SCALED = f'{RF_IMPRESSIONS}_scaled'
 
 # Non-media treatments baseline value constants.
 NON_MEDIA_BASELINE_MIN = 'min'
@@ -174,8 +177,40 @@ POSSIBLE_INPUT_DATA_COORDS_AND_ARRAYS_SET = frozenset(
     POSSIBLE_INPUT_DATA_COORD_NAMES + POSSIBLE_INPUT_DATA_ARRAY_NAMES
 )
 
-# EDA property constants
+# EDA Engine properties
 ORGANIC_RF_IMPRESSIONS = ORGANIC_PREFIX + RF_IMPRESSIONS
+ORGANIC_RF_IMPRESSIONS_SCALED = f'{ORGANIC_RF_IMPRESSIONS}_scaled'
+TREATMENT_CONTROL_SCALED = 'treatment_control_scaled'
+NATIONAL_TREATMENT_CONTROL_SCALED = (
+    f'{NATIONAL_PREFIX}{TREATMENT_CONTROL_SCALED}'
+)
+NATIONAL_CONTROLS_SCALED = f'{NATIONAL_PREFIX}{CONTROLS_SCALED}'
+NATIONAL_MEDIA_SPEND = f'{NATIONAL_PREFIX}{MEDIA_SPEND}'
+NATIONAL_MEDIA = f'{NATIONAL_PREFIX}{MEDIA}'
+NATIONAL_MEDIA_SCALED = f'{NATIONAL_PREFIX}{MEDIA_SCALED}'
+NATIONAL_ORGANIC_MEDIA = f'{NATIONAL_PREFIX}{ORGANIC_MEDIA}'
+NATIONAL_ORGANIC_MEDIA_SCALED = f'{NATIONAL_PREFIX}{ORGANIC_MEDIA_SCALED}'
+NATIONAL_NON_MEDIA_TREATMENTS_SCALED = (
+    f'{NATIONAL_PREFIX}{NON_MEDIA_TREATMENTS_SCALED}'
+)
+NATIONAL_RF_SPEND = f'{NATIONAL_PREFIX}{RF_SPEND}'
+NATIONAL_KPI_SCALED = f'{NATIONAL_PREFIX}{KPI_SCALED}'
+NATIONAL_REACH = f'{NATIONAL_PREFIX}{REACH}'
+NATIONAL_REACH_SCALED = f'{NATIONAL_PREFIX}{REACH_SCALED}'
+NATIONAL_ORGANIC_REACH = f'{NATIONAL_PREFIX}{ORGANIC_REACH}'
+NATIONAL_ORGANIC_REACH_SCALED = f'{NATIONAL_PREFIX}{ORGANIC_REACH_SCALED}'
+NATIONAL_FREQUENCY = f'{NATIONAL_PREFIX}{FREQUENCY}'
+NATIONAL_ORGANIC_FREQUENCY = f'{NATIONAL_PREFIX}{ORGANIC_FREQUENCY}'
+NATIONAL_RF_IMPRESSIONS = f'{NATIONAL_PREFIX}{RF_IMPRESSIONS}'
+NATIONAL_ORGANIC_RF_IMPRESSIONS = f'{NATIONAL_PREFIX}{ORGANIC_RF_IMPRESSIONS}'
+NATIONAL_RF_IMPRESSIONS_SCALED = f'{NATIONAL_PREFIX}{RF_IMPRESSIONS_SCALED}'
+NATIONAL_ORGANIC_RF_IMPRESSIONS_SCALED = (
+    f'{NATIONAL_PREFIX}{ORGANIC_RF_IMPRESSIONS_SCALED}'
+)
+ALL_REACH_SCALED = 'all_reach_scaled'
+ALL_FREQUENCY = 'all_frequency'
+NATIONAL_ALL_REACH_SCALED = f'{NATIONAL_PREFIX}{ALL_REACH_SCALED}'
+NATIONAL_ALL_FREQUENCY = f'{NATIONAL_PREFIX}{ALL_FREQUENCY}'
 
 
 # National model constants.
@@ -359,7 +394,6 @@ ALL_NATIONAL_DETERMINISTIC_PARAMETER_NAMES = (
     ETA_ORF,
 )
 
-
 MEDIA_PARAMETERS = (
     ROI_M,
     MROI_M,
@@ -747,6 +781,8 @@ HILL_NUM_STEPS = 500
 # Summary template params.
 START_DATE = 'start_date'
 END_DATE = 'end_date'
+DEFAULT_CURRENCY = '$'
+SELECTED_GEOS = 'selected_geos'
 CARD_INSIGHTS = 'insights'
 CARD_CHARTS = 'charts'
 CARD_STATS = 'stats'

meridian/model/__init__.py

@@ -15,6 +15,7 @@
 """The Meridian API module that models the data."""
 
 from meridian.model import adstock_hill
+from meridian.model import eda
 from meridian.model import knots
 from meridian.model import media
 from meridian.model import model

meridian/model/eda/__init__.py

@@ -15,3 +15,6 @@
 """The Meridian API module that performs EDA checks."""
 
 from meridian.model.eda import eda_engine
+from meridian.model.eda import eda_outcome
+from meridian.model.eda import eda_spec
+from meridian.model.eda import meridian_eda

meridian/model/eda/constants.py

@@ -0,0 +1,21 @@
+# 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.
+
+"""Constants specific to MeridianEDA."""
+
+# EDA Plotting properties
+VARIABLE_1 = 'var1'
+VARIABLE_2 = 'var2'
+VARIABLE = 'var'
+CORRELATION = 'correlation'