google-meridian 1.1.2__py3-none-any.whl → 1.1.4__py3-none-any.whl

meridian/mlflow/autolog.py

@@ -12,29 +12,130 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""MLflow autologging integration for Meridian."""
+"""MLflow autologging integration for Meridian.
 
+This module enables MLflow tracking for Meridian. When enabled via `autolog()`,
+parameters, metrics, and other metadata will be automatically logged to MLflow,
+allowing for improved experiment tracking and analysis.
+
+To enable MLflow autologging for your Meridian workflows, simply call
+`autolog.autolog()` once before your model run.
+
+Example usage:
+
+```python
+import mlflow
+from meridian.data import load
+from meridian.mlflow import autolog
+from meridian.model import model
+
+# Enable autologging (call this once per session)
+autolog.autolog(log_metrics=True)
+
+# Start an MLflow run (optionally name it for better grouping)
+with mlflow.start_run(run_name="my_run"):
+  # Load data
+  data = load.CsvDataLoader(...).load()
+
+  # Initialize Meridian model
+  mmm = model.Meridian(input_data=data)
+
+  # Run Meridian sampling processes
+  mmm.sample_prior(n_draws=100, seed=123)
+  mmm.sample_posterior(n_chains=7, n_adapt=500, n_burnin=500, n_keep=1000,
+  seed=1)
+
+# After the run completes, you can retrieve run results using the MLflow client.
+client = mlflow.tracking.MlflowClient()
+
+# Get the experiment ID for the run you just launched
+experiment_id = "0"
+
+# Search for runs matching the run name
+runs = client.search_runs(
+    experiment_id,
+    max_results=1000,
+    filter_string=f"attributes.run_name = 'my_run'"
+)
+
+# Print details of the run
+if runs:
+  print(runs[0])
+else:
+  print("No runs found.")
+```
+"""
+
+import dataclasses
+import inspect
+import json
 from typing import Any, Callable
 
 import arviz as az
-import meridian
+from meridian.analysis import visualizer
 import mlflow
 from mlflow.utils.autologging_utils import autologging_integration, safe_patch
 from meridian.model import model
+from meridian.model import posterior_sampler
+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"
 
 
+__all__ = ["autolog"]
+
+
 def _log_versions() -> None:
   """Logs Meridian and ArviZ versions."""
-  mlflow.log_param("meridian_version", meridian.__version__)
+  mlflow.log_param("meridian_version", __version__)
   mlflow.log_param("arviz_version", az.__version__)
 
 
+def _log_model_spec(model_spec: spec.ModelSpec) -> None:
+  """Logs the `ModelSpec` object."""
+  # TODO: Replace with serde api when it's available.
+  # PriorDistribution is logged separately.
+  excluded_fields = ["prior"]
+
+  for field in dataclasses.fields(model_spec):
+    if field.name in excluded_fields:
+      continue
+
+    field_value = getattr(model_spec, field.name)
+
+    # Stringify numpy arrays before logging.
+    if isinstance(field_value, np.ndarray):
+      field_value = json.dumps(field_value.tolist())
+
+    mlflow.log_param(f"spec.{field.name}", field_value)
+
+
+def _log_priors(model_spec: spec.ModelSpec) -> None:
+  """Logs the `PriorDistribution` object."""
+  # TODO: Replace with serde api when it's available.
+  priors = model_spec.prior
+  for field in dataclasses.fields(priors):
+    field_value = getattr(priors, field.name)
+
+    # Stringify Distributions and numpy arrays.
+    if isinstance(field_value, tfp.distributions.Distribution):
+      field_value = str(field_value)
+    elif isinstance(field_value, np.ndarray):
+      field_value = json.dumps(field_value.tolist())
+
+    mlflow.log_param(f"prior.{field.name}", field_value)
+
+
 @autologging_integration(FLAVOR_NAME)
 def autolog(
     disable: bool = False,  # pylint: disable=unused-argument
     silent: bool = False,  # pylint: disable=unused-argument
+    log_metrics: bool = False,
 ) -> None:
   """Enables MLflow tracking for Meridian.
 
@@ -43,12 +144,63 @@ def autolog(
   Args:
     disable: Whether to disable autologging.
     silent: Whether to suppress all event logs and warnings from MLflow.
+    log_metrics: Whether model metrics should be logged. Enabling this option
+      involves the creation of post-modeling objects to compute relevant
+      performance metrics. Metrics include R-Squared, MAPE, and wMAPE values.
   """
 
   def patch_meridian_init(
-      original: Callable[..., Any], *args, **kwargs
-  ) -> Callable[..., Any]:
+      original: Callable[..., Any], self, *args, **kwargs
+  ) -> model.Meridian:
     _log_versions()
-    return original(*args, **kwargs)
+    mmm = original(self, *args, **kwargs)
+    _log_model_spec(self.model_spec)
+    _log_priors(self.model_spec)
+    return mmm
+
+  def patch_prior_sampling(original: Callable[..., Any], self, *args, **kwargs):
+    mlflow.log_param("sample_prior.n_draws", kwargs.get("n_draws", "default"))
+    mlflow.log_param("sample_prior.seed", kwargs.get("seed", "default"))
+    return original(self, *args, **kwargs)
+
+  def patch_posterior_sampling(
+      original: Callable[..., Any], self, *args, **kwargs
+  ):
+    excluded_fields = ["current_state", "pins"]
+    params = [
+        name
+        for name, value in inspect.signature(original).parameters.items()
+        if name != "self"
+        and value.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
+        and name not in excluded_fields
+    ]
+
+    for param in params:
+      mlflow.log_param(
+          f"sample_posterior.{param}", kwargs.get(param, "default")
+      )
+
+    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"))
 
   safe_patch(FLAVOR_NAME, model.Meridian, "__init__", patch_meridian_init)
+  safe_patch(
+      FLAVOR_NAME,
+      prior_sampler.PriorDistributionSampler,
+      "__call__",
+      patch_prior_sampling,
+  )
+  safe_patch(
+      FLAVOR_NAME,
+      posterior_sampler.PosteriorMCMCSampler,
+      "__call__",
+      patch_posterior_sampling,
+  )

meridian/model/media.py

@@ -207,6 +207,8 @@ class RfTensors:
   Attributes:
     reach: A tensor constructed from `InputData.reach`.
     frequency: A tensor constructed from `InputData.frequency`.
+    rf_impressions: A tensor constructed from `InputData.reach` *
+      `InputData.frequency`.
     rf_spend: A tensor constructed from `InputData.rf_spend`.
     reach_transformer: A `MediaTransformer` to scale RF tensors using the
       model's RF data.
@@ -233,6 +235,7 @@ class RfTensors:
 
   reach: tf.Tensor | None = None
   frequency: tf.Tensor | None = None
+  rf_impressions: tf.Tensor | None = None
   rf_spend: tf.Tensor | None = None
   reach_transformer: transformers.MediaTransformer | None = None
   reach_scaled: tf.Tensor | None = None
@@ -250,6 +253,9 @@ def build_rf_tensors(
 
   reach = tf.convert_to_tensor(input_data.reach, dtype=tf.float32)
   frequency = tf.convert_to_tensor(input_data.frequency, dtype=tf.float32)
+  rf_impressions = (
+      reach * frequency if reach is not None and frequency is not None else None
+  )
   rf_spend = tf.convert_to_tensor(input_data.rf_spend, dtype=tf.float32)
   reach_transformer = transformers.MediaTransformer(
       reach, tf.convert_to_tensor(input_data.population, dtype=tf.float32)
@@ -292,6 +298,7 @@ def build_rf_tensors(
   return RfTensors(
       reach=reach,
       frequency=frequency,
+      rf_impressions=rf_impressions,
       rf_spend=rf_spend,
       reach_transformer=reach_transformer,
       reach_scaled=reach_scaled,

meridian/model/model.py

@@ -1447,8 +1447,7 @@ class Meridian:
         see [PRNGS and seeds]
         (https://github.com/tensorflow/probability/blob/main/PRNGS.md).
     """
-    prior_inference_data = self.prior_sampler_callable(n_draws, seed)
-    self.inference_data.extend(prior_inference_data, join="right")
+    self.prior_sampler_callable(n_draws=n_draws, seed=seed)
 
   def sample_posterior(
       self,
@@ -1527,22 +1526,21 @@ class Meridian:
         [ResourceExhaustedError when running Meridian.sample_posterior]
         (https://developers.google.com/meridian/docs/advanced-modeling/model-debugging#gpu-oom-error).
     """
-    posterior_inference_data = self.posterior_sampler_callable(
-        n_chains,
-        n_adapt,
-        n_burnin,
-        n_keep,
-        current_state,
-        init_step_size,
-        dual_averaging_kwargs,
-        max_tree_depth,
-        max_energy_diff,
-        unrolled_leapfrog_steps,
-        parallel_iterations,
-        seed,
+    self.posterior_sampler_callable(
+        n_chains=n_chains,
+        n_adapt=n_adapt,
+        n_burnin=n_burnin,
+        n_keep=n_keep,
+        current_state=current_state,
+        init_step_size=init_step_size,
+        dual_averaging_kwargs=dual_averaging_kwargs,
+        max_tree_depth=max_tree_depth,
+        max_energy_diff=max_energy_diff,
+        unrolled_leapfrog_steps=unrolled_leapfrog_steps,
+        parallel_iterations=parallel_iterations,
+        seed=seed,
         **pins,
     )
-    self.inference_data.extend(posterior_inference_data, join="right")
 
 
 def save_mmm(mmm: Meridian, file_path: str):

meridian/model/posterior_sampler.py

@@ -85,9 +85,13 @@ class PosteriorMCMCSampler:
   def __init__(self, meridian: "model.Meridian"):
     self._meridian = meridian
 
+  @property
+  def model(self) -> "model.Meridian":
+    return self._meridian
+
   def _get_joint_dist_unpinned(self) -> tfp.distributions.Distribution:
     """Returns a `JointDistributionCoroutineAutoBatched` function for MCMC."""
-    mmm = self._meridian
+    mmm = self.model
     mmm.populate_cached_properties()
 
     # This lists all the derived properties and states of this Meridian object
@@ -453,7 +457,7 @@ class PosteriorMCMCSampler:
     return joint_dist_unpinned
 
   def _get_joint_dist(self) -> tfp.distributions.Distribution:
-    mmm = self._meridian
+    mmm = self.model
     y = (
         tf.where(mmm.holdout_id, 0.0, mmm.kpi_scaled)
         if mmm.holdout_id is not None
@@ -476,7 +480,7 @@ class PosteriorMCMCSampler:
       parallel_iterations: int = 10,
       seed: Sequence[int] | int | None = None,
       **pins,
-  ) -> az.InferenceData:
+  ) -> None:
     """Runs Markov Chain Monte Carlo (MCMC) sampling of posterior distributions.
 
     For more information about the arguments, see [`windowed_adaptive_nuts`]
@@ -529,9 +533,6 @@ class PosteriorMCMCSampler:
       **pins: These are used to condition the provided joint distribution, and
         are passed directly to `joint_dist.experimental_pin(**pins)`.
 
-    Returns:
-      An Arviz `InferenceData` object containing posterior samples only.
-
     Throws:
       MCMCOOMError: If the model is out of memory. Try reducing `n_keep` or pass
         a list of integers as `n_chains` to sample chains serially. For more
@@ -589,10 +590,10 @@ class PosteriorMCMCSampler:
         if k not in constants.UNSAVED_PARAMETERS
     }
     # Create Arviz InferenceData for posterior draws.
-    posterior_coords = self._meridian.create_inference_data_coords(
+    posterior_coords = self.model.create_inference_data_coords(
         total_chains, n_keep
     )
-    posterior_dims = self._meridian.create_inference_data_dims()
+    posterior_dims = self.model.create_inference_data_dims()
     infdata_posterior = az.convert_to_inference_data(
         mcmc_states, coords=posterior_coords, dims=posterior_dims
     )
@@ -654,4 +655,7 @@ class PosteriorMCMCSampler:
         dims=sample_stats_dims,
         group="sample_stats",
     )
-    return az.concat(infdata_posterior, infdata_trace, infdata_sample_stats)
+    posterior_inference_data = az.concat(
+        infdata_posterior, infdata_trace, infdata_sample_stats
+    )
+    self.model.inference_data.extend(posterior_inference_data, join="right")

meridian/model/prior_sampler.py

@@ -588,22 +588,20 @@ class PriorDistributionSampler:
         | non_media_treatments_vars
     )
 
-  def __call__(self, n_draws: int, seed: int | None = None) -> az.InferenceData:
+  def __call__(self, n_draws: int, seed: int | None = None) -> None:
     """Draws samples from prior distributions.
 
-    Returns:
-      An Arviz `InferenceData` object containing prior samples only.
-
     Args:
       n_draws: Number of samples drawn from the prior distribution.
       seed: Used to set the seed for reproducible results. For more information,
         see [PRNGS and seeds]
         (https://github.com/tensorflow/probability/blob/main/PRNGS.md).
     """
-    prior_draws = self._sample_prior(n_draws, seed=seed)
+    prior_draws = self._sample_prior(n_draws=n_draws, seed=seed)
     # Create Arviz InferenceData for prior draws.
     prior_coords = self._meridian.create_inference_data_coords(1, n_draws)
     prior_dims = self._meridian.create_inference_data_dims()
-    return az.convert_to_inference_data(
+    prior_inference_data = az.convert_to_inference_data(
         prior_draws, coords=prior_coords, dims=prior_dims, group=constants.PRIOR
     )
+    self._meridian.inference_data.extend(prior_inference_data, join="right")

meridian/version.py

@@ -0,0 +1,17 @@
+# 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.
+
+"""Module for Meridian version."""
+
+__version__ = "1.1.4"