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

meridian/backend/config.py

@@ -0,0 +1,118 @@
+# 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.
+
+"""Backend configuration for Meridian."""
+
+import enum
+import os
+from typing import Union
+import warnings
+
+
+class Backend(enum.Enum):
+  TENSORFLOW = "tensorflow"
+  JAX = "jax"
+
+
+_DEFAULT_BACKEND = Backend.TENSORFLOW
+
+
+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 or a
+      valid string ('tensorflow', 'jax').
+
+  Raises:
+    ValueError: If the provided backend is not valid.
+  """
+  global _BACKEND
+
+  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_enum
+
+
+def get_backend() -> Backend:
+  """Returns the current backend for Meridian."""
+  return _BACKEND

meridian/backend/test_utils.py

@@ -0,0 +1,181 @@
+# 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.
+
+"""Common testing utilities for Meridian, designed to be backend-agnostic."""
+
+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.
+# We use `Any` here to avoid circular dependencies with the backend module
+# while still allowing the function to accept backend-specific tensor types.
+ArrayLike = Any
+
+
+def assert_allclose(
+    a: ArrayLike,
+    b: ArrayLike,
+    rtol: float = 1e-6,
+    atol: float = 1e-6,
+    err_msg: str = "",
+):
+  """Backend-agnostic assertion to check if two array-like objects are close.
+
+  This function converts both inputs to NumPy arrays before comparing them,
+  making it compatible with TensorFlow Tensors, JAX Arrays, and standard
+  Python lists or NumPy arrays.
+
+  Args:
+    a: The first array-like object to compare.
+    b: The second array-like object to compare.
+    rtol: The relative tolerance parameter.
+    atol: The absolute tolerance parameter.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the two arrays are not equal within the given tolerance.
+  """
+  np.testing.assert_allclose(
+      np.array(a), np.array(b), rtol=rtol, atol=atol, err_msg=err_msg
+  )
+
+
+def assert_allequal(a: ArrayLike, b: ArrayLike, err_msg: str = ""):
+  """Backend-agnostic assertion to check if two array-like objects are equal.
+
+  This function converts both inputs to NumPy arrays before comparing them.
+
+  Args:
+    a: The first array-like object to compare.
+    b: The second array-like object to compare.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the two arrays are not equal.
+  """
+  np.testing.assert_array_equal(np.array(a), np.array(b), err_msg=err_msg)
+
+
+def assert_all_finite(a: ArrayLike, err_msg: str = ""):
+  """Backend-agnostic assertion to check if all elements in an array are finite.
+
+  Args:
+    a: The array-like object to check.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the array contains non-finite values.
+  """
+  if not np.all(np.isfinite(np.array(a))):
+    raise AssertionError(err_msg or "Array contains non-finite values.")
+
+
+def assert_all_non_negative(a: ArrayLike, err_msg: str = ""):
+  """Backend-agnostic assertion to check if all elements are non-negative.
+
+  Args:
+    a: The array-like object to check.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the array contains negative values.
+  """
+  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,9 +67,10 @@ REACH = 'reach'
 FREQUENCY = 'frequency'
 RF_IMPRESSIONS = 'rf_impressions'
 RF_SPEND = 'rf_spend'
-ORGANIC_MEDIA = 'organic_media'
-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'
@@ -125,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'
 
@@ -143,8 +146,9 @@ MEDIA_CHANNEL = 'media_channel'
 RF_CHANNEL = 'rf_channel'
 CHANNEL = 'channel'
 RF = '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 = (
@@ -170,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'
@@ -212,9 +219,11 @@ NON_PAID_TREATMENT_PRIOR_TYPES = frozenset({
     TREATMENT_PRIOR_TYPE_COEFFICIENT,
     TREATMENT_PRIOR_TYPE_CONTRIBUTION,
 })
-PAID_MEDIA_ROI_PRIOR_TYPES = frozenset(
-    {TREATMENT_PRIOR_TYPE_ROI, TREATMENT_PRIOR_TYPE_MROI}
-)
+PAID_MEDIA_ROI_PRIOR_TYPES = frozenset({
+    TREATMENT_PRIOR_TYPE_ROI,
+    TREATMENT_PRIOR_TYPE_MROI,
+    TREATMENT_PRIOR_TYPE_CONTRIBUTION,
+})
 # Represents a 1% increase in spend.
 MROI_FACTOR = 1.01
 
@@ -315,6 +324,41 @@ RF_PARAMETER_NAMES = (
     BETA_RF,
     BETA_GRF,
 )
+ORGANIC_MEDIA_PARAMETER_NAMES = (
+    CONTRIBUTION_OM,
+    BETA_OM,
+    ETA_OM,
+    ALPHA_OM,
+    EC_OM,
+    SLOPE_OM,
+    BETA_GOM,
+)
+ORGANIC_RF_PARAMETER_NAMES = (
+    CONTRIBUTION_ORF,
+    BETA_ORF,
+    ETA_ORF,
+    ALPHA_ORF,
+    EC_ORF,
+    SLOPE_ORF,
+    BETA_GORF,
+)
+NON_MEDIA_PARAMETER_NAMES = (
+    CONTRIBUTION_N,
+    GAMMA_N,
+    XI_N,
+    GAMMA_GN,
+)
+ALL_NATIONAL_DETERMINISTIC_PARAMETER_NAMES = (
+    SLOPE_M,
+    SLOPE_OM,
+    XI_N,
+    XI_C,
+    ETA_M,
+    ETA_RF,
+    ETA_OM,
+    ETA_ORF,
+)
+
 
 MEDIA_PARAMETERS = (
     ROI_M,
@@ -501,10 +545,17 @@ ADSTOCK_HILL_FUNCTIONS = frozenset({
     'hill',
 })
 
+# Adstock decay functions.
+GEOMETRIC_DECAY = 'geometric'
+BINOMIAL_DECAY = 'binomial'
+
+ADSTOCK_DECAY_FUNCTIONS = frozenset({GEOMETRIC_DECAY, BINOMIAL_DECAY})
+ADSTOCK_CHANNELS = (MEDIA, RF, ORGANIC_MEDIA, ORGANIC_RF)
 
 # Distribution constants.
 DISTRIBUTION = 'distribution'
 DISTRIBUTION_TYPE = 'distribution_type'
+INDEPENDENT_MULTIVARIATE = 'IndependentMultivariate'
 PRIOR = 'prior'
 POSTERIOR = 'posterior'
 # Prior mean proportion of KPI incremental due to all media.
@@ -710,3 +761,13 @@ WEEKLY = 'weekly'
 QUARTERLY = 'quarterly'
 TIME_GRANULARITIES = frozenset({WEEKLY, QUARTERLY})
 QUARTERLY_SUMMARY_THRESHOLD_WEEKS = 52
+
+# Automatic Knot Selection constants
+KNOTS_SELECTED = 'knots_selected'
+SELECTION_COEFS = 'selection_coefs'
+MODEL = 'model'
+REGRESSION_COEFS = 'regression_coefs'
+SELECTED_MATRIX = 'selected_matrix'
+AIC = 'aic'
+BIC = 'bic'
+EBIC = 'ebic'

meridian/data/input_data.py

@@ -442,6 +442,59 @@ class InputData:
     """Checks whether the `rf_spend` array has a time dimension."""
     return self.rf_spend is not None and constants.TIME in self.rf_spend.coords
 
+  @property
+  def scaled_centered_kpi(self) -> np.ndarray:
+    """Calculates scaled and centered KPI values.
+
+    Returns:
+      An array of KPI values that have been population-scaled and
+    mean-centered by geo.
+    """
+    kpi = self.kpi.values
+    population = self.population.values[:, np.newaxis]
+
+    population_scaled_kpi = np.divide(
+        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, dtype=float
+        ),
+        where=(population_scaled_stdev != 0),
+    )
+    return kpi_scaled - np.mean(kpi_scaled, axis=1, keepdims=True)
+
+  def copy(self, deep: bool = True) -> "InputData":
+    """Returns a copy of the InputData instance.
+
+    Args:
+      deep: If True, a deep copy is made, meaning all xarray.DataArray objects
+        are also deepcopied. If False, a shallow copy is made.
+
+    Returns:
+      A new InputData instance.
+    """
+    if not deep:
+      return dataclasses.replace(self)
+
+    copied_fields = {}
+    for field in dataclasses.fields(self):
+      value = getattr(self, field.name)
+      if isinstance(value, xr.DataArray):
+        copied_fields[field.name] = value.copy(deep=True)
+      else:
+        # For other types, dataclasses.replace does a shallow copy.
+        copied_fields[field.name] = value
+
+    return InputData(**copied_fields)
+
   def _validate_scenarios(self):
     """Verifies that calibration and analysis is set correctly."""
     n_geos = len(self.kpi.coords[constants.GEO])
@@ -848,6 +901,32 @@ class InputData:
       raise ValueError("Both RF and media channel values are missing.")
     # pytype: enable=attribute-error
 
+  def get_all_adstock_hill_channels(self) -> np.ndarray:
+    """Returns all channel dimensions that adstock hill is applied to.
+
+    RF, organic media and organic RF channels are concatenated to the end of the
+    media channels if they are present.
+    """
+    adstock_hill_channels = []
+
+    if self.media_channel is not None:
+      adstock_hill_channels.append(self.media_channel.values)
+
+    if self.rf_channel is not None:
+      adstock_hill_channels.append(self.rf_channel.values)
+
+    if self.organic_media_channel is not None:
+      adstock_hill_channels.append(self.organic_media_channel.values)
+
+    if self.organic_rf_channel is not None:
+      adstock_hill_channels.append(self.organic_rf_channel.values)
+
+    if not adstock_hill_channels:
+      raise ValueError("Media, RF, organic media and organic RF channels are "
+                       "all missing.")
+
+    return np.concatenate(adstock_hill_channels, axis=None)
+
   def get_paid_channels_argument_builder(
       self,
   ) -> arg_builder.OrderedListArgumentBuilder:
@@ -870,6 +949,26 @@ class InputData:
       raise ValueError("There are no RF channels in the input data.")
     return arg_builder.OrderedListArgumentBuilder(self.rf_channel.values)
 
+  def get_organic_media_channels_argument_builder(
+      self
+  ) -> arg_builder.OrderedListArgumentBuilder:
+    """Returns an argument builder for *organic* media channels *only*."""
+    if self.organic_media_channel is None:
+      raise ValueError("There are no organic media channels in the input data.")
+    return arg_builder.OrderedListArgumentBuilder(
+        self.organic_media_channel.values
+        )
+
+  def get_organic_rf_channels_argument_builder(
+      self
+  ) -> arg_builder.OrderedListArgumentBuilder:
+    """Returns an argument builder for *organic* RF channels *only*."""
+    if self.organic_rf_channel is None:
+      raise ValueError("There are no organic RF channels in the input data.")
+    return arg_builder.OrderedListArgumentBuilder(
+        self.organic_rf_channel.values
+        )
+
   def get_all_channels(self) -> np.ndarray:
     """Returns all the channel dimensions.