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

{google_meridian-1.1.2.dist-info → google_meridian-1.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-meridian
-Version: 1.1.2
+Version: 1.1.3
 Summary: Google's open source mixed marketing model library, helps you understand your return on investment and direct your ad spend with confidence. 
 Author-email: The Meridian Authors <no-reply@google.com>
 License: 
@@ -397,7 +397,7 @@ To cite this repository:
   author = {Google Meridian Marketing Mix Modeling Team},
   title = {Meridian: Marketing Mix Modeling},
   url = {https://github.com/google/meridian},
-  version = {1.1.2},
+  version = {1.1.3},
   year = {2025},
 }
 ```

{google_meridian-1.1.2.dist-info → google_meridian-1.1.3.dist-info}/RECORD

@@ -1,10 +1,11 @@
-google_meridian-1.1.2.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-meridian/__init__.py,sha256=rWkSMlr2TiRmH7Xf9z1Bj3grQiSbmrxl3dtGB9YGn9o,714
-meridian/constants.py,sha256=AWhDEP9VcyQtPCbZhM6cPXHeWuz19wjaqB5lGz6qBsw,17161
+google_meridian-1.1.3.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+meridian/__init__.py,sha256=XROKwHNVQvEa371QCXAHik5wN_YKObOdJQX9bJ2c4M4,832
+meridian/constants.py,sha256=VAVHyGfm9FyDd0dWomfqK5XYDUt9qJx7SAM4rzDh3RQ,17195
+meridian/version.py,sha256=CUTXDDaOfXFTukX_ywPK6Q3PiK9hMyJbmJRBeb5ez7c,644
 meridian/analysis/__init__.py,sha256=nGBYz7k9FVdadO_WVGMKJcfq7Yy_TuuP8zgee4i9pSA,836
-meridian/analysis/analyzer.py,sha256=AP2YJpM2R2qMJ-rwtMmgu-cM-xJLJCFodSaP9K8f0Do,204458
+meridian/analysis/analyzer.py,sha256=FY_SvnkmEqqCIS37UXB3bvaQi-U3BwLcSWhH1puTzdQ,206003
 meridian/analysis/formatter.py,sha256=ENIdR1CRiaVqIGEXx1HcnsA4ewgDD_nhsYCweJAThaw,7270
-meridian/analysis/optimizer.py,sha256=Se6_sg0O3A4p80vdVnRtDeyNaE5s-ywxKoU0CODQsWM,107608
+meridian/analysis/optimizer.py,sha256=P4uMcV9ByqMapqa1TEqcnu-3NyTH9fR8QLszdKxRAFc,107801
 meridian/analysis/summarizer.py,sha256=IthOUTMufGvAvbxiDhaKwe7uYCyiTyiQ8vgdmUtdevs,18855
 meridian/analysis/summary_text.py,sha256=I_smDkZJYp2j77ea-9AIbgeraDa7-qUYyb-IthP2qO4,12438
 meridian/analysis/test_utils.py,sha256=ES1r1akhRjD4pf2oTaGqzDfGNu9weAcLv6UZRuIkfEc,77699
@@ -22,25 +23,25 @@ meridian/data/__init__.py,sha256=StIe-wfYnnbfUbKtZHwnAQcRQUS8XCZk_PCaEzw90Ww,929
 meridian/data/arg_builder.py,sha256=Kqlt88bOqFj6D3xNwvWo4MBwNwcDFHzd-wMfEOmLoPU,3741
 meridian/data/data_frame_input_data_builder.py,sha256=3m6wrcC0psmD2ijsXk3R4uByA0Tu2gJxZBGaTS6Z7Io,22040
 meridian/data/input_data.py,sha256=teJPKTBfW-AzBWgf_fEO_S_Z1J_veqQkCvctINaid6I,39749
-meridian/data/input_data_builder.py,sha256=fFJTmUuIdTnTnZPtZNTiEf4_fsqR_haY7O9ZOFj47bE,25409
-meridian/data/load.py,sha256=cvvesjL6Dc7pYu2nOl558gUOZVAW_B69GirzHocyY3Q,42855
+meridian/data/input_data_builder.py,sha256=08E_MZLrCzwfjvjPWFVs7o_094vVJ5o6VmbTfrg4NUM,25602
+meridian/data/load.py,sha256=B-12fBhsghN7wj0A9IWyT7BVogIXjuUDDvR34JJFwPM,45157
 meridian/data/nd_array_input_data_builder.py,sha256=lfpmnENGuSGKyUd7bDGAwoLqHqteOKmHdKl0VI2wCQA,16341
 meridian/data/test_utils.py,sha256=6GJrPmeaF4uzMxxRgzERGv4g1XMUHwI0s7qDVMZUjuI,55565
 meridian/data/time_coordinates.py,sha256=C5A5fscSLjPH6G9YT8OspgIlCrkMY7y8dMFEt3tNSnE,9874
 meridian/mlflow/__init__.py,sha256=elwXUqPQYi7VF9PYjelU1tydfcUrmtuoq6eJCOnV9bk,693
-meridian/mlflow/autolog.py,sha256=Duubd_Z2Exlk_MJqGTfMIfFjaDUqG_YnsRGjzY4Hn84,1696
+meridian/mlflow/autolog.py,sha256=s240eLGAurzaNsulwRlyM1ZdBLvUzyr2eOMYgOyWAzk,6393
 meridian/model/__init__.py,sha256=9NFfqUE5WgFc-9lQMkbfkwwV-bQIz0tsQ_3Jyq0A4SU,982
 meridian/model/adstock_hill.py,sha256=20A_6rbDUAADEkkHspB7JpCm5tYfYS1FQ6hJMLu21Pk,9283
 meridian/model/knots.py,sha256=KPEgnb-UdQQ4QBugOYEke-zBgEghgTmeCMoeiJ30meY,8054
-meridian/model/media.py,sha256=R0LnMUNTuGzXD2lzNRRORA4-p21xpdhkVVsvFaWtEK0,13819
-meridian/model/model.py,sha256=KM2EU7eAK5UHDAn1jbUEI_SBrDkz-Bc93R8qRBEiic8,61500
+meridian/model/media.py,sha256=3BaPX8xYAFMEvf0mz3mBSCIDWViIs7M218nrCklc6Fk,14099
+meridian/model/model.py,sha256=BlLPyskHrEx5D71mUZFbNxS2VjkQgaiaE6hLKvQ5D3A,61489
 meridian/model/model_test_data.py,sha256=hDDTEzm72LknW9c5E_dNsy4Mm4Tfs6AirhGf_QxykFs,15552
-meridian/model/posterior_sampler.py,sha256=jjLqcYEAorVJ_2nmhpkVUjCGAyNUZYPTEXVTDHufbqA,27727
+meridian/model/posterior_sampler.py,sha256=K49zWTTelME2rL1JLeFAdMPzL0OwrBvyAXA3oR-kgSI,27801
 meridian/model/prior_distribution.py,sha256=IEDU1rabcmKNY8lxwbbO4OUAlMHPIMa7flM_zsu3DLM,42417
-meridian/model/prior_sampler.py,sha256=jSaxFmJzyN2OKqKyU059Ar4Yr565w4zlInPl4zxjGZk,23212
+meridian/model/prior_sampler.py,sha256=cmu6jG-bSEkYDkjVUxl3iSxrL7r-LN7a77cb2Vc0LoA,23218
 meridian/model/spec.py,sha256=0HNiMQUWQpYvWYOZr1_fj2ah8tH-bEyfEjoqgBZ9Lc0,18049
 meridian/model/transformers.py,sha256=nRjzq1fQG0ypldxboM7Gqok6WSAXAS1witRXoAzeH9Q,7763
-google_meridian-1.1.2.dist-info/METADATA,sha256=qoLl6RDBz8LxrnJ3c4-hDiIMcL5OC8G8X61rotn6PGs,22201
-google_meridian-1.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-google_meridian-1.1.2.dist-info/top_level.txt,sha256=nwaCebZvvU34EopTKZsjK0OMTFjVnkf4FfnBN_TAc0g,9
-google_meridian-1.1.2.dist-info/RECORD,,
+google_meridian-1.1.3.dist-info/METADATA,sha256=5W_XWui7q5gH68OC3Z-PXbDOeBftDbWuhqznNv7fOAk,22201
+google_meridian-1.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+google_meridian-1.1.3.dist-info/top_level.txt,sha256=nwaCebZvvU34EopTKZsjK0OMTFjVnkf4FfnBN_TAc0g,9
+google_meridian-1.1.3.dist-info/RECORD,,

meridian/__init__.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 """Meridian API."""
-
-__version__ = "1.1.2"
-
-
 from meridian import analysis
 from meridian import data
 from meridian import model
+from meridian.version import __version__
+
+try:
+  from meridian import mlflow  # pylint: disable=g-import-not-at-top
+except ImportError:
+  pass

meridian/analysis/analyzer.py

@@ -65,6 +65,8 @@ class DataTensors(tf.experimental.ExtensionType):
       time dimension `T`.
     frequency: Optional tensor with dimensions `(n_geos, T, n_rf_channels)` for
       any time dimension `T`.
+    rf_impressions: Optional tensor with dimensions `(n_geos, T, n_rf_channels)`
+      for any time dimension `T`.
     rf_spend: Optional tensor with dimensions `(n_geos, T, n_rf_channels)` for
       any time dimension `T`.
     organic_media: Optional tensor with dimensions `(n_geos, T,
@@ -86,6 +88,7 @@ class DataTensors(tf.experimental.ExtensionType):
   media_spend: Optional[tf.Tensor]
   reach: Optional[tf.Tensor]
   frequency: Optional[tf.Tensor]
+  rf_impressions: Optional[tf.Tensor]
   rf_spend: Optional[tf.Tensor]
   organic_media: Optional[tf.Tensor]
   organic_reach: Optional[tf.Tensor]
@@ -101,6 +104,7 @@ class DataTensors(tf.experimental.ExtensionType):
       media_spend: Optional[tf.Tensor] = None,
       reach: Optional[tf.Tensor] = None,
       frequency: Optional[tf.Tensor] = None,
+      rf_impressions: Optional[tf.Tensor] = None,
       rf_spend: Optional[tf.Tensor] = None,
       organic_media: Optional[tf.Tensor] = None,
       organic_reach: Optional[tf.Tensor] = None,
@@ -118,6 +122,11 @@ class DataTensors(tf.experimental.ExtensionType):
     self.frequency = (
         tf.cast(frequency, tf.float32) if frequency is not None else None
     )
+    self.rf_impressions = (
+        tf.cast(rf_impressions, tf.float32)
+        if rf_impressions is not None
+        else None
+    )
     self.rf_spend = (
         tf.cast(rf_spend, tf.float32) if rf_spend is not None else None
     )
@@ -189,7 +198,10 @@ class DataTensors(tf.experimental.ExtensionType):
     """
     for field in self._tf_extension_type_fields():
       new_tensor = getattr(self, field.name)
-      old_tensor = getattr(meridian.input_data, field.name)
+      if field.name == constants.RF_IMPRESSIONS:
+        old_tensor = getattr(meridian.rf_tensors, field.name)
+      else:
+        old_tensor = getattr(meridian.input_data, field.name)
       # The time dimension is always the second dimension, except for when spend
       # data is provided with only one dimension of (n_channels).
       if (
@@ -293,7 +305,13 @@ class DataTensors(tf.experimental.ExtensionType):
             "This is not supported and will be ignored."
         )
       if field.name in required_variables:
-        if getattr(meridian.input_data, field.name) is None:
+        if field.name == constants.RF_IMPRESSIONS:
+          if meridian.n_rf_channels == 0:
+            raise ValueError(
+                "New `rf_impressions` is not allowed because there are no R&F"
+                " channels in the Meridian model."
+            )
+        elif getattr(meridian.input_data, field.name) is None:
           raise ValueError(
               f"New `{field.name}` is not allowed because the input data to the"
               f" Meridian model does not contain `{field.name}`."
@@ -322,7 +340,10 @@ class DataTensors(tf.experimental.ExtensionType):
       if var_name in [constants.REVENUE_PER_KPI, constants.TIME]:
         continue
       new_tensor = getattr(self, var_name)
-      old_tensor = getattr(meridian.input_data, var_name)
+      if var_name == constants.RF_IMPRESSIONS:
+        old_tensor = getattr(meridian.rf_tensors, var_name)
+      else:
+        old_tensor = getattr(meridian.input_data, var_name)
       if new_tensor is not None:
         assert old_tensor is not None
         if new_tensor.shape[-1] != old_tensor.shape[-1]:
@@ -337,7 +358,10 @@ class DataTensors(tf.experimental.ExtensionType):
     """Validates the time dimension of the specified data variables."""
     for var_name in required_fields:
       new_tensor = getattr(self, var_name)
-      old_tensor = getattr(meridian.input_data, var_name)
+      if var_name == constants.RF_IMPRESSIONS:
+        old_tensor = getattr(meridian.rf_tensors, var_name)
+      else:
+        old_tensor = getattr(meridian.input_data, var_name)
 
       # Skip spend data with only 1 dimension of (n_channels).
       if (
@@ -375,7 +399,10 @@ class DataTensors(tf.experimental.ExtensionType):
     missing_params = []
     for var_name in required_fields:
       new_tensor = getattr(self, var_name)
-      old_tensor = getattr(meridian.input_data, var_name)
+      if var_name == constants.RF_IMPRESSIONS:
+        old_tensor = getattr(meridian.rf_tensors, var_name)
+      else:
+        old_tensor = getattr(meridian.input_data, var_name)
 
       if old_tensor is None:
         continue
@@ -3415,6 +3442,7 @@ class Analyzer:
   def optimal_freq(
       self,
       new_data: DataTensors | None = None,
+      max_frequency: float | None = None,
       freq_grid: Sequence[float] | None = None,
       use_posterior: bool = True,
       use_kpi: bool = False,
@@ -3443,7 +3471,7 @@ class Analyzer:
     ROI numerator is KPI units.
 
     Args:
-      new_data: Optional `DataTensors` object containing `reach`, `frequency`,
+      new_data: Optional `DataTensors` object containing `rf_impressions`,
         `rf_spend`, and `revenue_per_kpi`. If provided, the optimal frequency is
         calculated using the values of the tensors passed in `new_data` and the
         original values of all the remaining tensors. If `None`, the historical
@@ -3451,6 +3479,10 @@ class Analyzer:
         tensors in `new_data` is provided with a different number of time
         periods than in `InputData`, then all tensors must be provided with the
         same number of time periods.
+      max_frequency: Maximum frequency value used to calculate the frequency
+        grid. If `None`, the maximum frequency value is calculated from the
+        historic frequency (maximum value of Meridian.input_data, not
+        `new_data`). If `freq_grid` is provided, this argument has no effect.
       freq_grid: List of frequency values. The ROI of each channel is calculated
         for each frequency value in the list. By default, the list includes
         numbers from `1.0` to the maximum frequency in increments of `0.1`.
@@ -3506,7 +3538,11 @@ class Analyzer:
       )
 
     filled_data = new_data.validate_and_fill_missing_data(
-        constants.RF_DATA,
+        [
+            constants.RF_IMPRESSIONS,
+            constants.RF_SPEND,
+            constants.REVENUE_PER_KPI,
+        ],
         self._meridian,
     )
     # TODO: Once treatment type filtering is added, remove adding
@@ -3527,7 +3563,9 @@ class Analyzer:
         (self._meridian.n_geos, n_times, self._meridian.n_media_channels)
     )
 
-    max_freq = np.max(np.array(filled_data.frequency))
+    max_freq = max_frequency or np.max(
+        np.array(self._meridian.rf_tensors.frequency)
+    )
     if freq_grid is None:
       freq_grid = np.arange(1, max_freq, 0.1)
 
@@ -3537,8 +3575,8 @@ class Analyzer:
     metric_grid = np.zeros((len(freq_grid), self._meridian.n_rf_channels, 4))
 
     for i, freq in enumerate(freq_grid):
-      new_frequency = tf.ones_like(filled_data.frequency) * freq
-      new_reach = filled_data.frequency * filled_data.reach / new_frequency
+      new_frequency = tf.ones_like(filled_data.rf_impressions) * freq
+      new_reach = filled_data.rf_impressions / new_frequency
       new_roi_data = DataTensors(
           reach=new_reach,
           frequency=new_frequency,
@@ -3568,12 +3606,10 @@ class Analyzer:
 
     optimal_frequency = [freq_grid[i] for i in optimal_freq_idx]
     optimal_frequency_tensor = tf.convert_to_tensor(
-        tf.ones_like(filled_data.frequency) * optimal_frequency,
+        tf.ones_like(filled_data.rf_impressions) * optimal_frequency,
         tf.float32,
     )
-    optimal_reach = (
-        filled_data.frequency * filled_data.reach / optimal_frequency_tensor
-    )
+    optimal_reach = filled_data.rf_impressions / optimal_frequency_tensor
 
     new_summary_metrics_data = DataTensors(
         reach=optimal_reach,

meridian/analysis/optimizer.py

@@ -1863,9 +1863,14 @@ class BudgetOptimizer:
         )
     )
     if self._meridian.n_rf_channels > 0 and use_optimal_frequency:
+      opt_freq_data = analyzer.DataTensors(
+          rf_impressions=filled_data.reach * filled_data.frequency,
+          rf_spend=filled_data.rf_spend,
+          revenue_per_kpi=filled_data.revenue_per_kpi,
+      )
       optimal_frequency = tf.convert_to_tensor(
           self._analyzer.optimal_freq(
-              new_data=filled_data.filter_fields(c.RF_DATA),
+              new_data=opt_freq_data,
               use_posterior=use_posterior,
               selected_times=selected_times,
               use_kpi=use_kpi,

meridian/constants.py

@@ -63,6 +63,7 @@ CONTROLS = 'controls'
 POPULATION = 'population'
 REACH = 'reach'
 FREQUENCY = 'frequency'
+RF_IMPRESSIONS = 'rf_impressions'
 RF_SPEND = 'rf_spend'
 ORGANIC_MEDIA = 'organic_media'
 ORGANIC_REACH = 'organic_reach'

meridian/data/input_data_builder.py

@@ -646,12 +646,13 @@ class InputDataBuilder(abc.ABC):
     """Normalizes the given `DataArray`'s coordinates in Meridian convention.
 
     Validates that time values are in the conventional Meridian format and
-    that geos have national name if national.
+    that geos have national name if national. If geo coordinates are not string-
+    typed, they will be converted to strings.
 
     Args:
       da: The DataArray to normalize.
-      time_dimension_name: The name of the time dimension. If None, the
-        will skip time normalization.
+      time_dimension_name: The name of the time dimension. If None, the will
+        skip time normalization.
 
     Returns:
       The normalized DataArray.
@@ -686,6 +687,11 @@ class InputDataBuilder(abc.ABC):
       da = da.assign_coords(
           {constants.GEO: [constants.NATIONAL_MODEL_DEFAULT_GEO_NAME]},
       )
+    else:
+      da = da.assign_coords(
+          {constants.GEO: da.coords[constants.GEO].astype(str)}
+      )
+
     return da
 
   def _validate_set(self, component: str, da: xr.DataArray):

meridian/data/load.py

@@ -816,12 +816,55 @@ class DataFrameDataLoader(InputDataLoader):
         'organic_frequency': 'organic_frequency_to_channel',
     })
     for coord_name, channel_dict in required_mappings.items():
+      if getattr(self.coord_to_columns, coord_name, None) is not None:
+        if getattr(self, channel_dict, None) is None:
+          raise ValueError(
+              f"When {coord_name} data is provided, '{channel_dict}' is"
+              ' required.'
+          )
+        else:
+          if set(getattr(self, channel_dict)) != set(
+              getattr(self.coord_to_columns, coord_name)
+          ):
+            raise ValueError(
+                f'The {channel_dict} keys must have the same set of values as'
+                f' the {coord_name} columns.'
+            )
+    if (
+        self.media_to_channel is not None
+        and self.media_spend_to_channel is not None
+    ):
+      if set(self.media_to_channel.values()) != set(
+          self.media_spend_to_channel.values()
+      ):
+        raise ValueError(
+            'The media and media_spend columns must have the same set of'
+            ' channels.'
+        )
+    if (
+        self.reach_to_channel is not None
+        and self.frequency_to_channel is not None
+        and self.rf_spend_to_channel is not None
+    ):
       if (
-          getattr(self.coord_to_columns, coord_name, None) is not None
-          and getattr(self, channel_dict, None) is None
+          set(self.reach_to_channel.values())
+          != set(self.frequency_to_channel.values())
+          != set(self.rf_spend_to_channel.values())
       ):
         raise ValueError(
-            f"When {coord_name} data is provided, '{channel_dict}' is required."
+            'The reach, frequency, and rf_spend columns must have the same set'
+            ' of channels.'
+        )
+    if (
+        self.organic_reach_to_channel is not None
+        and self.organic_frequency_to_channel is not None
+    ):
+      if set(self.organic_reach_to_channel.values()) != set(
+          self.organic_frequency_to_channel.values()
+      ):
+        raise ValueError(
+            'The organic_reach and organic_frequency columns must have the'
+            ' same set of channels.'
         )
 
   def load(self) -> input_data.InputData:
@@ -861,28 +904,36 @@ class DataFrameDataLoader(InputDataLoader):
           self.coord_to_columns.geo,
       )
     if (
-        self.coord_to_columns.media is not None
-        and self.media_to_channel is not None
+        self.media_to_channel is not None
+        and self.media_spend_to_channel is not None
     ):
+      sorted_channels = sorted(self.media_to_channel.values())
+      inv_media_map = {v: k for k, v in self.media_to_channel.items()}
+      inv_spend_map = {v: k for k, v in self.media_spend_to_channel.items()}
+
       builder.with_media(
           self.df,
-          list(self.coord_to_columns.media),
-          list(self.coord_to_columns.media_spend),
-          list(self.media_to_channel.values()),
+          [inv_media_map[ch] for ch in sorted_channels],
+          [inv_spend_map[ch] for ch in sorted_channels],
+          sorted_channels,
           self.coord_to_columns.time,
           self.coord_to_columns.geo,
       )
-
     if (
-        self.coord_to_columns.reach is not None
-        and self.reach_to_channel is not None
+        self.reach_to_channel is not None
+        and self.frequency_to_channel is not None
+        and self.rf_spend_to_channel is not None
     ):
+      sorted_channels = sorted(self.reach_to_channel.values())
+      inv_reach_map = {v: k for k, v in self.reach_to_channel.items()}
+      inv_freq_map = {v: k for k, v in self.frequency_to_channel.items()}
+      inv_rf_spend_map = {v: k for k, v in self.rf_spend_to_channel.items()}
       builder.with_reach(
           self.df,
-          list(self.coord_to_columns.reach),
-          list(self.coord_to_columns.frequency),
-          list(self.coord_to_columns.rf_spend),
-          list(self.reach_to_channel.values()),
+          [inv_reach_map[ch] for ch in sorted_channels],
+          [inv_freq_map[ch] for ch in sorted_channels],
+          [inv_rf_spend_map[ch] for ch in sorted_channels],
+          sorted_channels,
           self.coord_to_columns.time,
           self.coord_to_columns.geo,
       )
@@ -895,14 +946,19 @@ class DataFrameDataLoader(InputDataLoader):
           self.coord_to_columns.geo,
       )
     if (
-        self.coord_to_columns.organic_reach is not None
-        and self.organic_reach_to_channel is not None
+        self.organic_reach_to_channel is not None
+        and self.organic_frequency_to_channel is not None
     ):
+      sorted_channels = sorted(self.organic_reach_to_channel.values())
+      inv_reach_map = {v: k for k, v in self.organic_reach_to_channel.items()}
+      inv_freq_map = {
+          v: k for k, v in self.organic_frequency_to_channel.items()
+      }
       builder.with_organic_reach(
           self.df,
-          list(self.coord_to_columns.organic_reach),
-          list(self.coord_to_columns.organic_frequency),
-          list(self.organic_reach_to_channel.values()),
+          [inv_reach_map[ch] for ch in sorted_channels],
+          [inv_freq_map[ch] for ch in sorted_channels],
+          sorted_channels,
           self.coord_to_columns.time,
           self.coord_to_columns.geo,
       )

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.3"