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

meridian/backend/config.py

@@ -15,6 +15,8 @@
 """Backend configuration for Meridian."""
 
 import enum
+import os
+from typing import Union
 import warnings
 
 
@@ -23,35 +25,92 @@ class Backend(enum.Enum):
   JAX = "jax"
 
 
-_BACKEND = Backend.TENSORFLOW
+_DEFAULT_BACKEND = Backend.TENSORFLOW
 
 
-def set_backend(backend: Backend) -> None:
+def _warn_jax_experimental() -> None:
+  """Issues a warning that the JAX backend is experimental."""
+  warnings.warn(
+      (
+          "The JAX backend is currently under development and is not yet"
+          " functional. It is intended for internal testing only and should"
+          " not be used. Please use the TensorFlow backend."
+      ),
+      UserWarning,
+      # Set stacklevel=2 so the warning points to the caller of set_backend
+      # or the location where the module is imported if initialized via env var.
+      stacklevel=2,
+  )
+
+
+def _initialize_backend() -> Backend:
+  """Initializes the backend based on environment variables or defaults."""
+  env_backend_str = os.environ.get("MERIDIAN_BACKEND")
+
+  if not env_backend_str:
+    return _DEFAULT_BACKEND
+
+  try:
+    backend = Backend(env_backend_str.lower())
+    if backend == Backend.JAX:
+      _warn_jax_experimental()
+    return backend
+  except ValueError:
+    warnings.warn(
+        (
+            "Invalid MERIDIAN_BACKEND environment variable:"
+            f" '{env_backend_str}'. Supported values are 'tensorflow' and"
+            f" 'jax'. Defaulting to {_DEFAULT_BACKEND.value}."
+        ),
+        RuntimeWarning,
+    )
+    return _DEFAULT_BACKEND
+
+
+_BACKEND = _initialize_backend()
+
+
+def set_backend(backend: Union[Backend, str]) -> None:
   """Sets the backend for Meridian.
 
+  **Warning:** This function should ideally be called at the beginning of your
+  program, before any other Meridian modules are imported or used.
+
+  Changing the backend after Meridian's functions or classes have been
+  imported can lead to unpredictable behavior. This is because already-imported
+  modules will not reflect the backend change.
+
   Note: The JAX backend is currently under development and should not be used.
 
+  Changing the backend at runtime requires reloading the `meridian.backend`
+  module for the changes to take effect globally.
+
   Args:
-    backend: The backend to use, must be a member of the `Backend` enum.
+    backend: The backend to use, must be a member of the `Backend` enum or a
+      valid string ('tensorflow', 'jax').
 
   Raises:
-    ValueError: If the provided backend is not a valid `Backend` enum member.
+    ValueError: If the provided backend is not valid.
   """
   global _BACKEND
-  if not isinstance(backend, Backend):
-    raise ValueError("Backend must be a member of the Backend enum.")
 
-  if backend == Backend.JAX:
-    warnings.warn(
-        (
-            "The JAX backend is currently under development and is not yet"
-            " functional. It is intended for internal testing only and should"
-            " not be used. Please use the TensorFlow backend."
-        ),
-        UserWarning,
-    )
+  if isinstance(backend, str):
+    try:
+      backend_enum = Backend(backend.lower())
+    except ValueError as exc:
+      raise ValueError(
+          f"Invalid backend string '{backend}'. Must be one of: "
+          f"{[b.value for b in Backend]}"
+      ) from exc
+  elif isinstance(backend, Backend):
+    backend_enum = backend
+  else:
+    raise ValueError("Backend must be a Backend enum member or a string.")
+
+  if backend_enum == Backend.JAX and _BACKEND != Backend.JAX:
+    _warn_jax_experimental()
 
-  _BACKEND = backend
+  _BACKEND = backend_enum
 
 
 def get_backend() -> Backend:

meridian/backend/test_utils.py

@@ -14,7 +14,10 @@
 
 """Common testing utilities for Meridian, designed to be backend-agnostic."""
 
-from typing import Any
+from typing import Any, Optional
+from absl.testing import parameterized
+from meridian import backend
+from meridian.backend import config
 import numpy as np
 
 # A type alias for backend-agnostic array-like objects.
@@ -93,3 +96,86 @@ def assert_all_non_negative(a: ArrayLike, err_msg: str = ""):
   """
   if not np.all(np.array(a) >= 0):
     raise AssertionError(err_msg or "Array contains negative values.")
+
+
+class MeridianTestCase(parameterized.TestCase):
+  """Base test class for Meridian providing backend-aware utilities.
+
+  This class handles initialization timing issues (crucial for JAX by forcing
+  tensor operations into setUp) and provides a unified way to handle random
+  number generation across backends (Stateful TF vs Stateless JAX).
+  """
+
+  def setUp(self):
+    super().setUp()
+    # Default seed, can be overridden by subclasses before calling
+    # _initialize_rng().
+    self.seed = 42
+    self._jax_key = None
+    self._initialize_rng()
+
+  def _initialize_rng(self):
+    """Initializes the RNG state or key based on self.seed."""
+    current_backend = config.get_backend()
+
+    if current_backend == config.Backend.TENSORFLOW:
+      # In TF, we use the global stateful seed for test reproducibility.
+      try:
+        backend.set_random_seed(self.seed)
+      except NotImplementedError:
+        # Handle cases where backend might be misconfigured during transition.
+        pass
+    elif current_backend == config.Backend.JAX:
+      # In JAX, we must manage PRNGKeys explicitly.
+      # Import JAX locally to avoid hard dependency if TF is the active backend,
+      # and to ensure initialization happens after absltest.main() starts.
+      # pylint: disable=g-import-not-at-top
+      import jax
+      # pylint: enable=g-import-not-at-top
+      self._jax_key = jax.random.PRNGKey(self.seed)
+    else:
+      raise ValueError(f"Unknown backend: {current_backend}")
+
+  def get_next_rng_seed_or_key(self) -> Optional[Any]:
+    """Gets the next available seed or key for backend operations.
+
+    This should be passed to the `seed` argument of TFP sampling methods.
+
+    Returns:
+      A JAX PRNGKey if the backend is JAX (splitting the internal key).
+      None if the backend is TensorFlow (relying on the global state).
+    """
+    if self._jax_key is not None:
+      # JAX requires splitting the key for each use.
+      # pylint: disable=g-import-not-at-top
+      import jax
+      # pylint: enable=g-import-not-at-top
+      self._jax_key, subkey = jax.random.split(self._jax_key)
+      return subkey
+    else:
+      # For stateful TF, returning None allows TFP/TF to use the global seed.
+      return None
+
+  def sample(
+      self,
+      distribution: backend.tfd.Distribution,
+      sample_shape: Any = (),
+      **kwargs: Any,
+  ) -> backend.Tensor:
+    """Performs a backend-agnostic sample from a distribution.
+
+    This method abstracts away the need for explicit seed management in JAX.
+    When the JAX backend is active, it automatically provides a PRNGKey from
+    the test's managed key state. In TensorFlow, it performs a standard sample.
+
+    Args:
+      distribution: The TFP distribution object to sample from.
+      sample_shape: The shape of the desired sample.
+      **kwargs: Additional keyword arguments to pass to the underlying `sample`
+        method (e.g., `name`).
+
+    Returns:
+      A tensor containing the sampled values.
+    """
+    seed = self.get_next_rng_seed_or_key()
+    return distribution.sample(sample_shape=sample_shape, seed=seed, **kwargs)

meridian/constants.py

@@ -54,6 +54,8 @@ DATE_FORMAT = '%Y-%m-%d'
 # Example: "2024 Apr"
 QUARTER_FORMAT = '%Y %b'
 
+ORGANIC_PREFIX = 'organic_'
+
 # Input data variables.
 KPI = 'kpi'
 REVENUE_PER_KPI = 'revenue_per_kpi'
@@ -65,10 +67,10 @@ REACH = 'reach'
 FREQUENCY = 'frequency'
 RF_IMPRESSIONS = 'rf_impressions'
 RF_SPEND = 'rf_spend'
-ORGANIC_MEDIA = 'organic_media'
-ORGANIC_RF = 'organic_rf'
-ORGANIC_REACH = 'organic_reach'
-ORGANIC_FREQUENCY = 'organic_frequency'
+ORGANIC_MEDIA = ORGANIC_PREFIX + MEDIA
+# ORGANIC_RF is defined below.
+ORGANIC_REACH = ORGANIC_PREFIX + REACH
+ORGANIC_FREQUENCY = ORGANIC_PREFIX + FREQUENCY
 NON_MEDIA_TREATMENTS = 'non_media_treatments'
 REVENUE = 'revenue'
 NON_REVENUE = 'non_revenue'
@@ -126,8 +128,8 @@ NON_REVENUE_DATA = IMPRESSIONS_DATA + (CONTROLS,)
 # Scaled input data variables.
 MEDIA_SCALED = 'media_scaled'
 REACH_SCALED = 'reach_scaled'
-ORGANIC_MEDIA_SCALED = 'organic_media_scaled'
-ORGANIC_REACH_SCALED = 'organic_reach_scaled'
+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'
 
@@ -144,9 +146,9 @@ MEDIA_CHANNEL = 'media_channel'
 RF_CHANNEL = 'rf_channel'
 CHANNEL = 'channel'
 RF = 'rf'
-ORGANIC_RF = 'organic_rf'
-ORGANIC_MEDIA_CHANNEL = 'organic_media_channel'
-ORGANIC_RF_CHANNEL = 'organic_rf_channel'
+ORGANIC_RF = ORGANIC_PREFIX + RF
+ORGANIC_MEDIA_CHANNEL = ORGANIC_PREFIX + MEDIA_CHANNEL
+ORGANIC_RF_CHANNEL = ORGANIC_PREFIX + RF_CHANNEL
 NON_MEDIA_CHANNEL = 'non_media_channel'
 CONTROL_VARIABLE = 'control_variable'
 REQUIRED_INPUT_DATA_COORD_NAMES = (
@@ -172,6 +174,9 @@ POSSIBLE_INPUT_DATA_COORDS_AND_ARRAYS_SET = frozenset(
     POSSIBLE_INPUT_DATA_COORD_NAMES + POSSIBLE_INPUT_DATA_ARRAY_NAMES
 )
 
+# EDA property constants
+ORGANIC_RF_IMPRESSIONS = ORGANIC_PREFIX + RF_IMPRESSIONS
+
 
 # National model constants.
 NATIONAL = 'national'

meridian/data/input_data.py

@@ -454,14 +454,19 @@ class InputData:
     population = self.population.values[:, np.newaxis]
 
     population_scaled_kpi = np.divide(
-        kpi, population, out=np.zeros_like(kpi), where=(population != 0)
+        kpi,
+        population,
+        out=np.zeros_like(kpi, dtype=float),
+        where=(population != 0),
     )
     population_scaled_mean = np.mean(population_scaled_kpi)
     population_scaled_stdev = np.std(population_scaled_kpi)
     kpi_scaled = np.divide(
         population_scaled_kpi - population_scaled_mean,
         population_scaled_stdev,
-        out=np.zeros_like(population_scaled_kpi - population_scaled_mean),
+        out=np.zeros_like(
+            population_scaled_kpi - population_scaled_mean, dtype=float
+        ),
         where=(population_scaled_stdev != 0),
     )
     return kpi_scaled - np.mean(kpi_scaled, axis=1, keepdims=True)

meridian/data/test_utils.py

@@ -1898,10 +1898,12 @@ def sample_input_data_for_aks_with_expected_knot_info() -> (
       'non_revenue',
   )
   expected_knot_info = knots.KnotInfo(
-      n_knots=6,
-      knot_locations=np.array([38, 39, 41, 48, 50, 55]),
+      n_knots=13,
+      knot_locations=np.array(
+          [11, 14, 38, 39, 41, 43, 45, 48, 50, 55, 87, 89, 90]
+      ),
       weights=knots.l1_distance_weights(
-          117, np.array([38, 39, 41, 48, 50, 55])
+          117, np.array([11, 14, 38, 39, 41, 43, 45, 48, 50, 55, 87, 89, 90])
       ),
   )
   return data, expected_knot_info

meridian/mlflow/autolog.py

@@ -72,6 +72,7 @@ import json
 from typing import Any, Callable
 
 import arviz as az
+from meridian import backend
 from meridian.analysis import visualizer
 import mlflow
 from mlflow.utils.autologging_utils import autologging_integration, safe_patch
@@ -81,7 +82,6 @@ from meridian.model import prior_sampler
 from meridian.model import spec
 from meridian.version import __version__
 import numpy as np
-import tensorflow_probability as tfp
 
 
 FLAVOR_NAME = "meridian"
@@ -123,7 +123,7 @@ def _log_priors(model_spec: spec.ModelSpec) -> None:
     field_value = getattr(priors, field.name)
 
     # Stringify Distributions and numpy arrays.
-    if isinstance(field_value, tfp.distributions.Distribution):
+    if isinstance(field_value, backend.tfd.Distribution):
       field_value = str(field_value)
     elif isinstance(field_value, np.ndarray):
       field_value = json.dumps(field_value.tolist())

meridian/model/adstock_hill.py

@@ -145,8 +145,9 @@ def compute_decay_weights(
       alpha, l_range, window_size, constants.GEOMETRIC_DECAY, normalize,
   )
 
+  is_binomial = [s == constants.BINOMIAL_DECAY for s in decay_functions]
   binomial_decay_mask = backend.reshape(
-      backend.to_tensor(decay_functions) == constants.BINOMIAL_DECAY,
+      backend.to_tensor(is_binomial, dtype=backend.bool_),
       (-1, 1),
   )
 
@@ -189,14 +190,14 @@ def _compute_single_decay_function_weights(
       A tensor of weights with a shape of `(*alpha.shape, len(l_range))`.
   """
   expanded_alpha = backend.expand_dims(alpha, -1)
-  match decay_function:
-    case constants.GEOMETRIC_DECAY:
-      weights = expanded_alpha**l_range
-    case constants.BINOMIAL_DECAY:
-      mapped_alpha_binomial = _map_alpha_for_binomial_decay(expanded_alpha)
-      weights = (1 - l_range / window_size) ** mapped_alpha_binomial
-    case _:
-      raise ValueError(f'Unsupported decay function: {decay_function}')
+
+  if decay_function == constants.GEOMETRIC_DECAY:
+    weights = expanded_alpha**l_range
+  elif decay_function == constants.BINOMIAL_DECAY:
+    mapped_alpha_binomial = _map_alpha_for_binomial_decay(expanded_alpha)
+    weights = (1 - l_range / window_size) ** mapped_alpha_binomial
+  else:
+    raise ValueError(f'Unsupported decay function: {decay_function}')
 
   if normalize:
     normalization_factors = backend.reduce_sum(weights, axis=-1, keepdims=True)