google-meridian 1.4.0__py3-none-any.whl → 1.5.1__py3-none-any.whl

meridian/analysis/visualizer.py

@@ -48,7 +48,10 @@ class ModelDiagnostics:
 
   def __init__(self, meridian: model.Meridian, use_kpi: bool = False):
     self._meridian = meridian
-    self._analyzer = analyzer.Analyzer(meridian)
+    self._analyzer = analyzer.Analyzer(
+        model_context=meridian.model_context,
+        inference_data=meridian.inference_data,
+    )
     self._use_kpi = self._analyzer._use_kpi(use_kpi)
 
   @functools.lru_cache(maxsize=128)
@@ -271,11 +274,15 @@ class ModelDiagnostics:
           x=c.INDEPENDENT
       )
 
-    return plot.properties(
-        title=formatter.custom_title_params(
-            summary_text.PRIOR_POSTERIOR_DIST_CHART_TITLE
+    return (
+        plot.properties(
+            title=formatter.custom_title_params(
+                summary_text.PRIOR_POSTERIOR_DIST_CHART_TITLE
+            )
         )
-    ).configure_axis(**formatter.TEXT_CONFIG).interactive()
+        .configure_axis(**formatter.TEXT_CONFIG)
+        .interactive()
+    )
 
   def plot_rhat_boxplot(self) -> alt.Chart:
     """Plots the R-hat box plot.
@@ -387,7 +394,10 @@ class ModelFit:
         represented as a value between zero and one. Default is `0.9`.
     """
     self._meridian = meridian
-    self._analyzer = analyzer.Analyzer(meridian)
+    self._analyzer = analyzer.Analyzer(
+        model_context=meridian.model_context,
+        inference_data=meridian.inference_data,
+    )
     self._use_kpi = self._analyzer._use_kpi(use_kpi)
     self._model_fit_data = self._analyzer.expected_vs_actual_data(
         use_kpi=self._use_kpi, confidence_level=confidence_level
@@ -657,7 +667,10 @@ class ReachAndFrequency:
       use_kpi: If `True`, KPI is used instead of revenue.
     """
     self._meridian = meridian
-    self._analyzer = analyzer.Analyzer(meridian)
+    self._analyzer = analyzer.Analyzer(
+        model_context=meridian.model_context,
+        inference_data=meridian.inference_data,
+    )
     self._selected_times = selected_times
     self._use_kpi = self._analyzer._use_kpi(use_kpi)
     self._optimal_frequency_data = self._analyzer.optimal_freq(
@@ -857,7 +870,10 @@ class MediaEffects:
         the incremental revenue using the revenue per KPI (if available).
     """
     self._meridian = meridian
-    self._analyzer = analyzer.Analyzer(meridian)
+    self._analyzer = analyzer.Analyzer(
+        model_context=meridian.model_context,
+        inference_data=meridian.inference_data,
+    )
     self._by_reach = by_reach
     self._use_kpi = self._analyzer._use_kpi(use_kpi)
 
@@ -1431,7 +1447,10 @@ class MediaSummary:
       use_kpi: If `True`, use KPI instead of revenue.
     """
     self._meridian = meridian
-    self._analyzer = analyzer.Analyzer(meridian)
+    self._analyzer = analyzer.Analyzer(
+        model_context=meridian.model_context,
+        inference_data=meridian.inference_data,
+    )
     self._confidence_level = confidence_level
     self._selected_times = selected_times
     self._marginal_roi_by_reach = marginal_roi_by_reach

meridian/backend/__init__.py

@@ -909,6 +909,77 @@ if _BACKEND == config.Backend.JAX:
 
   xla_windowed_adaptive_nuts = _jax_xla_windowed_adaptive_nuts
 
+  def _jax_adstock_process(
+      media: "_jax.Array", weights: "_jax.Array", n_times_output: int
+  ) -> "_jax.Array":
+    """JAX implementation for adstock_process using convolution.
+
+    This function applies an adstock process to media spend data using a
+    convolutional approach. The weights represent the adstock decay over time.
+
+    Args:
+      media: A JAX array of media spend. Expected shape is
+        `(batch_dims, n_geos, n_times_in, n_channels)`.
+      weights: A JAX array of adstock weights. Expected shape is
+        `(batch_dims, n_channels, window_size)`, where `batch_dims` must be
+        broadcastable to the batch dimensions of `media`.
+      n_times_output: The number of time periods in the output. This corresponds
+        to `n_times_in - window_size + 1`.
+
+    Returns:
+      A JAX array representing the adstocked media, with shape
+      `(batch_dims, n_geos, n_times_output, n_channels)`.
+    """
+
+    batch_dims = weights.shape[:-2]
+    if media.shape[:-3] != batch_dims:
+      media = jax_ops.broadcast_to(media, batch_dims + media.shape[-3:])
+
+    n_geos = media.shape[-3]
+    n_times_in = media.shape[-2]
+    n_channels = media.shape[-1]
+    window_size = weights.shape[-1]
+
+    perm = list(range(media.ndim))
+    perm[-2], perm[-1] = perm[-1], perm[-2]
+    media_transposed = jax_ops.transpose(media, perm)
+    media_reshaped = jax_ops.reshape(media_transposed, (1, -1, n_times_in))
+
+    total_channels = media_reshaped.shape[1]
+    weights_expanded = jax_ops.expand_dims(weights, -3)
+    weights_tiled = jax_ops.broadcast_to(
+        weights_expanded, batch_dims + (n_geos, n_channels, window_size)
+    )
+    kernel_reshaped = jax_ops.reshape(
+        weights_tiled, (total_channels, 1, window_size)
+    )
+
+    dn = jax.lax.conv_dimension_numbers(
+        media_reshaped.shape, kernel_reshaped.shape, ("NCH", "OIH", "NCH")
+    )
+
+    out = jax.lax.conv_general_dilated(
+        lhs=media_reshaped,
+        rhs=kernel_reshaped,
+        window_strides=(1,),
+        padding="VALID",
+        lhs_dilation=(1,),
+        rhs_dilation=(1,),
+        dimension_numbers=dn,
+        feature_group_count=total_channels,
+        precision=jax.lax.Precision.HIGHEST,
+    )
+
+    t_out = out.shape[-1]
+    out_reshaped = jax_ops.reshape(
+        out, batch_dims + (n_geos, n_channels, t_out)
+    )
+    perm_back = list(range(out_reshaped.ndim))
+    perm_back[-2], perm_back[-1] = perm_back[-1], perm_back[-2]
+    out_final = jax_ops.transpose(out_reshaped, perm_back)
+
+    return out_final[..., :n_times_output, :]
+
   _ops = jax_ops
   errors = _JaxErrors()
   Tensor = jax.Array
@@ -920,6 +991,7 @@ if _BACKEND == config.Backend.JAX:
 
   # Standardized Public API
   absolute = _ops.abs
+  adstock_process = _jax_adstock_process
   allclose = _ops.allclose
   arange = _jax_arange
   argmax = _jax_argmax
@@ -1059,6 +1131,39 @@ elif _BACKEND == config.Backend.TENSORFLOW:
 
   xla_windowed_adaptive_nuts = _tf_xla_windowed_adaptive_nuts
 
+  def _tf_adstock_process(
+      media: "_tf.Tensor", weights: "_tf.Tensor", n_times_output: int
+  ) -> "_tf.Tensor":
+    """TensorFlow implementation for adstock_process using loop/einsum.
+
+    This function applies an adstock process to media spend data. It achieves
+    this by creating a windowed view of the `media` tensor and then using
+    `tf.einsum` to efficiently compute the weighted sum based on the provided
+    `weights`. The `weights` tensor defines the decay effect over a specific
+    `window_size`. The output is truncated to `n_times_output` periods.
+
+    Args:
+      media: Input media tensor. Expected shape is `(..., num_geos,
+        num_times_in, num_channels)`. The `...` represents optional batch
+        dimensions.
+      weights: Adstock weights tensor. Expected shape is `(..., num_channels,
+        window_size)`. The batch dimensions must be broadcast-compatible with
+        those in `media`.
+      n_times_output: The number of time periods to output. This should be less
+        than or equal to `num_times_in - window_size + 1`.
+
+    Returns:
+      A tensor of shape `(..., num_geos, n_times_output, num_channels)`
+      representing the adstocked media.
+    """
+
+    window_size = weights.shape[-1]
+    window_list = [
+        media[..., i : i + n_times_output, :] for i in range(window_size)
+    ]
+    windowed = tf_backend.stack(window_list)
+    return tf_backend.einsum("...cw,w...gtc->...gtc", weights, windowed)
+
   tfd = tfp.distributions
   bijectors = tfp.bijectors
   experimental = tfp.experimental
@@ -1067,6 +1172,7 @@ elif _BACKEND == config.Backend.TENSORFLOW:
 
   # Standardized Public API
   absolute = _ops.math.abs
+  adstock_process = _tf_adstock_process
   allclose = _ops.experimental.numpy.allclose
   arange = _tf_arange
   argmax = _tf_argmax

meridian/constants.py

@@ -392,6 +392,7 @@ ALL_NATIONAL_DETERMINISTIC_PARAMETER_NAMES = (
     ETA_RF,
     ETA_OM,
     ETA_ORF,
+    TAU_G,
 )
 
 MEDIA_PARAMETERS = (

meridian/data/input_data.py

@@ -20,13 +20,13 @@ The `InputData` class is used to store all the input data to the model.
 from collections import abc
 from collections.abc import Sequence
 import dataclasses
-import datetime as dt
 import functools
 import warnings
 
 from meridian import constants
 from meridian.data import arg_builder
 from meridian.data import time_coordinates as tc
+from meridian.data import validator
 import numpy as np
 import xarray as xr
 
@@ -298,6 +298,7 @@ class InputData:
     self._validate_time_formats()
     self._validate_times()
     self._validate_geos()
+    self._validate_no_negative_values()
 
   def _convert_geos_to_strings(self):
     """Converts geo coordinates to strings in all relevant DataArrays."""
@@ -542,17 +543,36 @@ class InputData:
           f" `{constants.REVENUE}` or `{constants.NON_REVENUE}`."
       )
 
-    if (self.kpi.values < 0).any():
-      raise ValueError("KPI values must be non-negative.")
-
     if (
         self.revenue_per_kpi is not None
-        and (self.revenue_per_kpi.values <= 0).all()
+        and (self.revenue_per_kpi.values == 0).all()
     ):
       raise ValueError(
-          "Revenue per KPI values must not be all zero or negative."
+          "All Revenue per KPI values are 0, which can break the ROI"
+          " computation. If this is not a data error, please consider setting"
+          " revenue_per_kpi to None or follow the instructions at"
+          " https://developers.google.com/meridian/docs/advanced-modeling/unknown-revenue-kpi-default#default-total-paid-media-contribution-prior."
       )
 
+  def _validate_no_negative_values(self) -> None:
+    """Validates no negative values for applicable fields."""
+
+    fields_to_loggable_name = {
+        constants.MEDIA_SPEND: "Media Spend",
+        constants.RF_SPEND: "RF Spend",
+        constants.REACH: "Reach",
+        constants.FREQUENCY: "Frequency",
+        constants.ORGANIC_REACH: "Organic Reach",
+        constants.ORGANIC_FREQUENCY: "Organic Frequency",
+        constants.REVENUE_PER_KPI: "Revenue per KPI",
+        constants.KPI: "KPI",
+    }
+
+    for field, loggable_field in fields_to_loggable_name.items():
+      da = getattr(self, field)
+      if da is not None and (da.values < 0).any():
+        raise ValueError(f"{loggable_field} values must be non-negative.")
+
   def _validate_names(self):
     """Verifies that the names of the data arrays are correct."""
     # Must match the order of constants.POSSIBLE_INPUT_DATA_ARRAY_NAMES!
@@ -762,52 +782,10 @@ class InputData:
 
   def _validate_time_formats(self):
     """Validates the time coordinate format for all variables."""
-    self._validate_time_coord_format(self.kpi)
-    self._validate_time_coord_format(self.revenue_per_kpi)
-    self._validate_time_coord_format(self.controls)
-    self._validate_time_coord_format(self.media)
-    self._validate_time_coord_format(self.media_spend)
-    self._validate_time_coord_format(self.reach)
-    self._validate_time_coord_format(self.frequency)
-    self._validate_time_coord_format(self.rf_spend)
-    self._validate_time_coord_format(self.organic_media)
-    self._validate_time_coord_format(self.organic_reach)
-    self._validate_time_coord_format(self.organic_frequency)
-    self._validate_time_coord_format(self.non_media_treatments)
-
-  def _validate_time_coord_format(self, array: xr.DataArray | None):
-    """Validates the `time` dimensions format of the selected DataArray.
-
-    The `time` dimension of the selected array must have labels that are
-    formatted in the Meridian conventional `"yyyy-mm-dd"` format.
-
-    Args:
-      array: An optional DataArray to validate.
-    """
-    if array is None:
-      return
-
-    time_values = array.coords.get(constants.TIME, None)
-    if time_values is not None:
-      for time in time_values:
-        try:
-          _ = dt.datetime.strptime(time.item(), constants.DATE_FORMAT)
-        except (TypeError, ValueError) as exc:
-          raise ValueError(
-              f"Invalid time label: {time.item()}. Expected format:"
-              f" {constants.DATE_FORMAT}"
-          ) from exc
-
-    media_time_values = array.coords.get(constants.MEDIA_TIME, None)
-    if media_time_values is not None:
-      for time in media_time_values:
-        try:
-          _ = dt.datetime.strptime(time.item(), constants.DATE_FORMAT)
-        except (TypeError, ValueError) as exc:
-          raise ValueError(
-              f"Invalid media_time label: {time.item()}. Expected format:"
-              f" {constants.DATE_FORMAT}"
-          ) from exc
+    for field in dataclasses.fields(self):
+      attr = getattr(self, field.name)
+      if field.name != constants.POPULATION and isinstance(attr, xr.DataArray):
+        validator.validate_time_coord_format(attr)
 
   def _check_unique_names(self, dim: str, array: xr.DataArray | None):
     """Checks if a DataArray contains unique names on the specified dimension."""

meridian/data/input_data_builder.py

@@ -21,11 +21,11 @@ validation logic and an overall final validation logic before a valid
 
 import abc
 from collections.abc import Sequence
-import datetime
 import warnings
 from meridian import constants
 from meridian.data import input_data
 from meridian.data import time_coordinates as tc
+from meridian.data import validator
 import natsort
 import numpy as np
 import xarray as xr
@@ -676,14 +676,7 @@ class InputDataBuilder(abc.ABC):
 
       # Assume that the time coordinate labels are date-formatted strings.
       # We don't currently support other, arbitrary object types in the builder.
-      for time in da.coords[time_dimension_name].values:
-        try:
-          _ = datetime.datetime.strptime(time, constants.DATE_FORMAT)
-        except ValueError as exc:
-          raise ValueError(
-              f"Invalid time label: '{time}'. Expected format:"
-              f" '{constants.DATE_FORMAT}'"
-          ) from exc
+      validator.validate_time_coord_format(da)
 
     if len(da.coords[constants.GEO].values.tolist()) == 1:
       da = da.assign_coords(

meridian/data/test_utils.py

@@ -1476,53 +1476,37 @@ def random_dataset(
       constant_value=constant_population_value,
   )
 
-  dataset = xr.combine_by_coords(
-      [kpi, population] + ([controls] if controls is not None else [])
-  )
+  to_merge = [kpi, population]
+  if controls is not None:
+    to_merge.append(controls)
   if revenue_per_kpi is not None:
-    dataset = xr.combine_by_coords([dataset, revenue_per_kpi])
+    to_merge.append(revenue_per_kpi)
   if media is not None:
-    media_renamed = (
-        media.rename({'media_time': 'time'}) if remove_media_time else media
-    )
-
-    dataset = xr.combine_by_coords([dataset, media_renamed, media_spend])
+    if remove_media_time:
+      media = media.rename({'media_time': 'time'})
+    to_merge.append(media)
+    to_merge.append(media_spend)
   if reach is not None:
-    reach_renamed = (
-        reach.rename({'media_time': 'time'}) if remove_media_time else reach
-    )
-    frequency_renamed = (
-        frequency.rename({'media_time': 'time'})
-        if remove_media_time
-        else frequency
-    )
-    dataset = xr.combine_by_coords(
-        [dataset, reach_renamed, frequency_renamed, rf_spend]
-    )
+    if remove_media_time:
+      reach = reach.rename({'media_time': 'time'})
+      frequency = frequency.rename({'media_time': 'time'})
+    to_merge.append(reach)
+    to_merge.append(frequency)
+    to_merge.append(rf_spend)
   if non_media_treatments is not None:
-    dataset = xr.combine_by_coords([dataset, non_media_treatments])
+    to_merge.append(non_media_treatments)
   if organic_media is not None:
-    organic_media_renamed = (
-        organic_media.rename({'media_time': 'time'})
-        if remove_media_time
-        else organic_media
-    )
-    dataset = xr.combine_by_coords([dataset, organic_media_renamed])
+    if remove_media_time:
+      organic_media = organic_media.rename({'media_time': 'time'})
+    to_merge.append(organic_media)
   if organic_reach is not None:
-    organic_reach_renamed = (
-        organic_reach.rename({'media_time': 'time'})
-        if remove_media_time
-        else organic_reach
-    )
-    organic_frequency_renamed = (
-        organic_frequency.rename({'media_time': 'time'})
-        if remove_media_time
-        else organic_frequency
-    )
-    dataset = xr.combine_by_coords(
-        [dataset, organic_reach_renamed, organic_frequency_renamed]
-    )
-  return dataset
+    if remove_media_time:
+      organic_reach = organic_reach.rename({'media_time': 'time'})
+      organic_frequency = organic_frequency.rename({'media_time': 'time'})
+    to_merge.append(organic_reach)
+    to_merge.append(organic_frequency)
+
+  return xr.merge(to_merge, join='outer', compat='no_conflicts')
 
 
 def dataset_to_dataframe(

meridian/data/validator.py

@@ -0,0 +1,48 @@
+# 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.
+
+"""This module contains common validation functions for Meridian data."""
+
+import datetime as dt
+from meridian import constants
+import xarray as xr
+
+
+def validate_time_coord_format(array: xr.DataArray | None):
+  """Validates the `time` dimensions format of the selected DataArray.
+
+  The `time` dimension of the selected array must have labels that are
+  formatted in the Meridian conventional `"yyyy-mm-dd"` format.
+
+  Args:
+    array: An optional DataArray to validate.
+  """
+  if array is None:
+    return
+
+  # The component data arrays from the input data builders that call this helper
+  # method should only have one of either `media_time` or `time` as its time
+  # dimension.
+  target_coords = [constants.TIME, constants.MEDIA_TIME]
+
+  for coord_name in target_coords:
+    if (values := array.coords.get(coord_name)) is not None:
+      for time in values:
+        try:
+          dt.datetime.strptime(time.item(), constants.DATE_FORMAT)
+        except (TypeError, ValueError) as exc:
+          raise ValueError(
+              f"Invalid {coord_name} label: {time.item()!r}. "
+              f"Expected format: '{constants.DATE_FORMAT}'"
+          ) from exc

meridian/mlflow/autolog.py

@@ -70,6 +70,7 @@ import dataclasses
 import inspect
 import json
 from typing import Any, Callable
+import warnings
 
 import arviz as az
 from meridian import backend
@@ -180,16 +181,25 @@ def autolog(
           f"sample_posterior.{param}", kwargs.get(param, "default")
       )
 
-    original(self, *args, **kwargs)
+    result = original(self, *args, **kwargs)
     if log_metrics:
-      model_diagnostics = visualizer.ModelDiagnostics(self.model)
-      df_diag = model_diagnostics.predictive_accuracy_table()
-
-      get_metric = lambda n: df_diag[df_diag.metric == n].value.to_list()[0]
-
-      mlflow.log_metric("R_Squared", get_metric("R_Squared"))
-      mlflow.log_metric("MAPE", get_metric("MAPE"))
-      mlflow.log_metric("wMAPE", get_metric("wMAPE"))
+      # TODO: Direct injection of `model.Meridian` object into
+      # `PosteriorMCMCSampler` is deprecated. Revisit patching method here.
+      if self.model is not None:
+        model_diagnostics = visualizer.ModelDiagnostics(self.model)
+        df_diag = model_diagnostics.predictive_accuracy_table()
+
+        get_metric = lambda n: df_diag[df_diag.metric == n].value.to_list()[0]
+
+        mlflow.log_metric("R_Squared", get_metric("R_Squared"))
+        mlflow.log_metric("MAPE", get_metric("MAPE"))
+        mlflow.log_metric("wMAPE", get_metric("wMAPE"))
+      else:
+        warnings.warn(
+            "log_metrics=True is not supported when PosteriorMCMCSampler is"
+            " initialized with model_context."
+        )
+    return result
 
   safe_patch(FLAVOR_NAME, model.Meridian, "__init__", patch_meridian_init)
   safe_patch(

meridian/model/adstock_hill.py

@@ -279,10 +279,6 @@ def _adstock(
     media = backend.concatenate([backend.zeros(pad_shape), media], axis=-2)
 
   # Adstock calculation.
-  window_list = [None] * window_size
-  for i in range(window_size):
-    window_list[i] = media[..., i : i + n_times_output, :]
-  windowed = backend.stack(window_list)
   l_range = backend.arange(window_size - 1, -1, -1, dtype=backend.float32)
   weights = compute_decay_weights(
       alpha=alpha,
@@ -291,7 +287,9 @@ def _adstock(
       decay_functions=decay_functions,
       normalize=True,
   )
-  return backend.einsum('...mw,w...gtm->...gtm', weights, windowed)
+  return backend.adstock_process(
+      media=media, weights=weights, n_times_output=n_times_output
+  )
 
 
 def _map_alpha_for_binomial_decay(x: backend.Tensor):