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

meridian/data/test_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.
@@ -409,10 +409,7 @@ DATASET_WITHOUT_TIME_VARIATION_IN_ORGANIC_MEDIA = xr.Dataset(
 )
 
 DATASET_WITHOUT_TIME_VARIATION_IN_ORGANIC_REACH = xr.Dataset(
-    coords=_REQUIRED_COORDS
-    | _MEDIA_COORDS
-    | _RF_COORDS
-    | _ORGANIC_RF_COORDS,
+    coords=_REQUIRED_COORDS | _MEDIA_COORDS | _RF_COORDS | _ORGANIC_RF_COORDS,
     data_vars=_REQUIRED_DATA_VARS
     | _MEDIA_DATA_VARS
     | _RF_DATA_VARS
@@ -818,11 +815,11 @@ def random_controls_da(
 
 def random_kpi_da(
     media: xr.DataArray,
-    controls: xr.DataArray,
     n_geos: int,
     n_times: int,
     n_media_channels: int,
-    n_controls: int,
+    n_controls: int | None = None,
+    controls: xr.DataArray | None = None,
     seed: int = 0,
     integer_geos: bool = False,
 ) -> xr.DataArray:
@@ -838,19 +835,26 @@ def random_kpi_da(
       n_media_channels,
       axis=2,
   )
-  control_geo_sd = abs(np.random.normal(0, 5, size=n_geos))
-  control_geo_sd = np.repeat(
-      np.repeat(control_geo_sd[:, np.newaxis], n_times, axis=1)[
-          ..., np.newaxis
-      ],
-      n_controls,
-      axis=2,
-  )
+  if n_controls:
+    control_geo_sd = abs(np.random.normal(0, 5, size=n_geos))
+    control_geo_sd = np.repeat(
+        np.repeat(control_geo_sd[:, np.newaxis], n_times, axis=1)[
+            ..., np.newaxis
+        ],
+        n_controls,
+        axis=2,
+    )
+  else:
+    control_geo_sd = 0
 
   # Simulates outcome which is the dependent variable. Typically this is the
   # number of units sold, but it can be any metric (e.g. revenue).
   media_portion = np.random.normal(media_common, media_geo_sd).sum(axis=2)
-  control_portion = np.random.normal(controls, control_geo_sd).sum(axis=2)
+  if controls is not None:
+    control_portion = np.random.normal(controls, control_geo_sd).sum(axis=2)
+  else:
+    control_portion = 0
+
   error = np.random.normal(0, 2, size=(n_geos, n_times))
   kpi = abs(media_portion + control_portion + error)
 
@@ -1161,7 +1165,7 @@ def random_dataset(
     n_geos: int,
     n_times: int,
     n_media_times: int,
-    n_controls: int,
+    n_controls: int | None = None,
     n_non_media_channels: int | None = None,
     n_organic_media_channels: int | None = None,
     n_organic_rf_channels: int | None = None,
@@ -1171,7 +1175,7 @@ def random_dataset(
     seed: int = 0,
     remove_media_time: bool = False,
     integer_geos: bool = False,
-):
+) -> xr.Dataset:
   """Generates a random dataset."""
   if n_media_channels:
     media = random_media_da(
@@ -1232,14 +1236,18 @@ def random_dataset(
   else:
     revenue_per_kpi = None
 
-  controls = random_controls_da(
-      media=media if n_media_channels else reach,
-      n_geos=n_geos,
-      n_times=n_times,
-      n_controls=n_controls,
-      seed=seed,
-      integer_geos=integer_geos,
-  )
+  if n_controls:
+    controls = random_controls_da(
+        media=media if n_media_channels else reach,
+        n_geos=n_geos,
+        n_times=n_times,
+        n_controls=n_controls,
+        seed=seed,
+        integer_geos=integer_geos,
+    )
+  else:
+    controls = None
+
   if n_non_media_channels:
     non_media_treatments = random_non_media_treatments_da(
         media=media if n_media_channels else reach,
@@ -1298,7 +1306,9 @@ def random_dataset(
       n_geos=n_geos, seed=seed, integer_geos=integer_geos
   )
 
-  dataset = xr.combine_by_coords([kpi, population, controls])
+  dataset = xr.combine_by_coords(
+      [kpi, population] + ([controls] if controls is not None else [])
+  )
   if revenue_per_kpi is not None:
     dataset = xr.combine_by_coords([dataset, revenue_per_kpi])
   if media is not None:
@@ -1347,7 +1357,7 @@ def random_dataset(
 
 def dataset_to_dataframe(
     dataset: xr.Dataset,
-    controls_column_names: list[str],
+    controls_column_names: list[str] | None = None,
     media_column_names: list[str] | None = None,
     media_spend_column_names: list[str] | None = None,
     reach_column_names: list[str] | None = None,
@@ -1396,10 +1406,15 @@ def dataset_to_dataframe(
   )
   population = dataset[c.POPULATION].to_dataframe(name=c.POPULATION)
 
-  controls = dataset[c.CONTROLS].to_dataframe(name=c.CONTROLS).unstack()
-  controls.columns = controls_column_names
+  if controls_column_names is not None:
+    controls = dataset[c.CONTROLS].to_dataframe(name=c.CONTROLS).unstack()
+    controls.columns = controls_column_names
+  else:
+    controls = None
 
-  result = kpi.join(revenue_per_kpi).join(population).join(controls)
+  result = kpi.join(revenue_per_kpi).join(population)
+  if controls is not None:
+    result = result.join(controls)
 
   if non_media_column_names is not None:
     non_media_treatments = (
@@ -1471,7 +1486,7 @@ def random_dataframe(
     n_geos,
     n_times,
     n_media_times,
-    n_controls,
+    n_controls=None,
     n_media_channels=None,
     n_rf_channels=None,
     seed=0,
@@ -1489,7 +1504,9 @@ def random_dataframe(
 
   return dataset_to_dataframe(
       dataset,
-      controls_column_names=_sample_names('control_', n_controls),
+      controls_column_names=(
+          _sample_names('control_', n_controls) if n_controls else None
+      ),
       media_column_names=_sample_names('media_', n_media_channels),
       media_spend_column_names=_sample_names('media_spend_', n_media_channels),
       reach_column_names=_sample_names('reach_', n_rf_channels),
@@ -1499,7 +1516,7 @@ def random_dataframe(
 
 
 def sample_coord_to_columns(
-    n_controls: int,
+    n_controls: int | None = None,
     n_media_channels: int | None = None,
     n_rf_channels: int | None = None,
     n_non_media_channels: int | None = None,
@@ -1550,7 +1567,7 @@ def sample_coord_to_columns(
       kpi=c.KPI,
       revenue_per_kpi=c.REVENUE_PER_KPI if include_revenue_per_kpi else None,
       population=c.POPULATION,
-      controls=_sample_names('control_', n_controls),
+      controls=(_sample_names('control_', n_controls) if n_controls else None),
       media=media,
       media_spend=media_spend,
       reach=reach,
@@ -1668,7 +1685,7 @@ def sample_input_data_non_revenue_revenue_per_kpi(
     n_geos: int = 10,
     n_times: int = 50,
     n_media_times: int = 53,
-    n_controls: int = 2,
+    n_controls: int | None = 2,
     n_non_media_channels: int | None = None,
     n_media_channels: int | None = None,
     n_rf_channels: int | None = None,
@@ -1694,10 +1711,10 @@ def sample_input_data_non_revenue_revenue_per_kpi(
       kpi_type=c.NON_REVENUE,
       revenue_per_kpi=dataset.revenue_per_kpi,
       population=dataset.population,
-      controls=dataset.controls,
-      non_media_treatments=dataset.non_media_treatments
-      if n_non_media_channels
-      else None,
+      controls=(dataset.controls if n_controls else None),
+      non_media_treatments=(
+          dataset.non_media_treatments if n_non_media_channels else None
+      ),
       media=dataset.media if n_media_channels else None,
       media_spend=dataset.media_spend if n_media_channels else None,
       reach=dataset.reach if n_rf_channels else None,
@@ -1705,9 +1722,9 @@ def sample_input_data_non_revenue_revenue_per_kpi(
       rf_spend=dataset.rf_spend if n_rf_channels else None,
       organic_media=dataset.organic_media if n_organic_media_channels else None,
       organic_reach=dataset.organic_reach if n_organic_rf_channels else None,
-      organic_frequency=dataset.organic_frequency
-      if n_organic_rf_channels
-      else None,
+      organic_frequency=(
+          dataset.organic_frequency if n_organic_rf_channels else None
+      ),
   )
 
 

meridian/data/time_coordinates.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.

meridian/mlflow/__init__.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.
+
+"""Meridian MLflow module that contains autologging functionality."""
+
+from meridian.mlflow import autolog

meridian/mlflow/autolog.py

@@ -0,0 +1,54 @@
+# 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.
+
+"""MLflow autologging integration for Meridian."""
+
+from typing import Any, Callable
+
+import arviz as az
+import meridian
+import mlflow
+from mlflow.utils.autologging_utils import autologging_integration, safe_patch
+from meridian.model import model
+
+FLAVOR_NAME = "meridian"
+
+
+def _log_versions() -> None:
+  """Logs Meridian and ArviZ versions."""
+  mlflow.log_param("meridian_version", meridian.__version__)
+  mlflow.log_param("arviz_version", az.__version__)
+
+
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    disable: bool = False,  # pylint: disable=unused-argument
+    silent: bool = False,  # pylint: disable=unused-argument
+) -> None:
+  """Enables MLflow tracking for Meridian.
+
+  See https://mlflow.org/docs/latest/tracking/
+
+  Args:
+    disable: Whether to disable autologging.
+    silent: Whether to suppress all event logs and warnings from MLflow.
+  """
+
+  def patch_meridian_init(
+      original: Callable[..., Any], *args, **kwargs
+  ) -> Callable[..., Any]:
+    _log_versions()
+    return original(*args, **kwargs)
+
+  safe_patch(FLAVOR_NAME, model.Meridian, "__init__", patch_meridian_init)

meridian/model/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.

meridian/model/adstock_hill.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.

meridian/model/knots.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.

meridian/model/media.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.

meridian/model/model.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.
@@ -214,7 +214,9 @@ class Meridian:
     )
 
   @functools.cached_property
-  def controls(self) -> tf.Tensor:
+  def controls(self) -> tf.Tensor | None:
+    if self.input_data.controls is None:
+      return None
     return tf.convert_to_tensor(self.input_data.controls, dtype=tf.float32)
 
   @functools.cached_property
@@ -271,6 +273,8 @@ class Meridian:
 
   @property
   def n_controls(self) -> int:
+    if self.input_data.control_variable is None:
+      return 0
     return len(self.input_data.control_variable)
 
   @property
@@ -304,7 +308,13 @@ class Meridian:
     )
 
   @functools.cached_property
-  def controls_transformer(self) -> transformers.CenteringAndScalingTransformer:
+  def controls_transformer(
+      self,
+  ) -> transformers.CenteringAndScalingTransformer | None:
+    """Returns a `CenteringAndScalingTransformer` for controls, if it exists."""
+    if self.controls is None:
+      return None
+
     if self.model_spec.control_population_scaling_id is not None:
       controls_population_scaling_id = tf.convert_to_tensor(
           self.model_spec.control_population_scaling_id, dtype=bool
@@ -343,8 +353,12 @@ class Meridian:
     return transformers.KpiTransformer(self.kpi, self.population)
 
   @functools.cached_property
-  def controls_scaled(self) -> tf.Tensor:
-    return self.controls_transformer.forward(self.controls)
+  def controls_scaled(self) -> tf.Tensor | None:
+    if self.controls is not None:
+      # If `controls` is defined, then `controls_transformer` is also defined.
+      return self.controls_transformer.forward(self.controls)  # pytype: disable=attribute-error
+    else:
+      return None
 
   @functools.cached_property
   def non_media_treatments_normalized(self) -> tf.Tensor | None:
@@ -894,11 +908,12 @@ class Meridian:
     if self.is_national:
       return
 
-    self._check_if_no_geo_variation(
-        self.controls_scaled,
-        constants.CONTROLS,
-        self.input_data.controls.coords[constants.CONTROL_VARIABLE].values,
-    )
+    if self.input_data.controls is not None:
+      self._check_if_no_geo_variation(
+          self.controls_scaled,
+          constants.CONTROLS,
+          self.input_data.controls.coords[constants.CONTROL_VARIABLE].values,
+      )
     if self.input_data.non_media_treatments is not None:
       self._check_if_no_geo_variation(
           self.non_media_treatments_normalized,
@@ -971,12 +986,12 @@ class Meridian:
 
   def _validate_time_invariants(self):
     """Validates model time invariants."""
-
-    self._check_if_no_time_variation(
-        self.controls_scaled,
-        constants.CONTROLS,
-        self.input_data.controls.coords[constants.CONTROL_VARIABLE].values,
-    )
+    if self.input_data.controls is not None:
+      self._check_if_no_time_variation(
+          self.controls_scaled,
+          constants.CONTROLS,
+          self.input_data.controls.coords[constants.CONTROL_VARIABLE].values,
+      )
     if self.input_data.non_media_treatments is not None:
       self._check_if_no_time_variation(
           self.non_media_treatments_normalized,
@@ -1031,16 +1046,24 @@ class Meridian:
     mask = tf.equal(counts, self.n_geos)
     col_idx_bad = tf.boolean_mask(col_idx_unique, mask)
     dims_bad = tf.gather(data_dims, col_idx_bad)
-
-    if col_idx_bad.shape[0] and not self.is_national:
-      raise ValueError(
-          f"The following {data_name} variables do not vary across time, making"
-          f" a model with geo main effects unidentifiable: {dims_bad}. This can"
-          " lead to poor model convergence. Since these variables only vary"
-          " across geo and not across time, they are collinear with geo and"
-          " redundant in a model with geo main effects. To address this, drop"
-          " the listed variables that do not vary across time."
-      )
+    if col_idx_bad.shape[0]:
+      if self.is_national:
+        raise ValueError(
+            f"The following {data_name} variables do not vary across time,"
+            " which is equivalent to no signal at all in a national model:"
+            f" {dims_bad}.  This can lead to poor model convergence. To address"
+            " this, drop the listed variables that do not vary across time."
+        )
+      else:
+        raise ValueError(
+            f"The following {data_name} variables do not vary across time,"
+            f" making a model with geo main effects unidentifiable: {dims_bad}."
+            " This can lead to poor model convergence. Since these variables"
+            " only vary across geo and not across time, they are collinear"
+            " with geo and redundant in a model with geo main effects. To"
+            " address this, drop the listed variables that do not vary across"
+            " time."
+        )
 
   def _validate_kpi_transformer(self):
     """Validates the KPI transformer."""
@@ -1356,31 +1379,36 @@ class Meridian:
       self, n_chains: int, n_draws: int
   ) -> Mapping[str, np.ndarray | Sequence[str]]:
     """Creates data coordinates for inference data."""
-    media_channel_values = (
+    media_channel_names = (
         self.input_data.media_channel
         if self.input_data.media_channel is not None
         else np.array([])
     )
-    rf_channel_values = (
+    rf_channel_names = (
         self.input_data.rf_channel
         if self.input_data.rf_channel is not None
         else np.array([])
     )
-    organic_media_channel_values = (
+    organic_media_channel_names = (
         self.input_data.organic_media_channel
         if self.input_data.organic_media_channel is not None
         else np.array([])
     )
-    organic_rf_channel_values = (
+    organic_rf_channel_names = (
         self.input_data.organic_rf_channel
         if self.input_data.organic_rf_channel is not None
         else np.array([])
     )
-    non_media_channel_values = (
+    non_media_channel_names = (
         self.input_data.non_media_channel
         if self.input_data.non_media_channel is not None
         else np.array([])
     )
+    control_variable_names = (
+        self.input_data.control_variable
+        if self.input_data.control_variable is not None
+        else np.array([])
+    )
     return {
         constants.CHAIN: np.arange(n_chains),
         constants.DRAW: np.arange(n_draws),
@@ -1388,12 +1416,12 @@ class Meridian:
         constants.TIME: self.input_data.time,
         constants.MEDIA_TIME: self.input_data.media_time,
         constants.KNOTS: np.arange(self.knot_info.n_knots),
-        constants.CONTROL_VARIABLE: self.input_data.control_variable,
-        constants.NON_MEDIA_CHANNEL: non_media_channel_values,
-        constants.MEDIA_CHANNEL: media_channel_values,
-        constants.RF_CHANNEL: rf_channel_values,
-        constants.ORGANIC_MEDIA_CHANNEL: organic_media_channel_values,
-        constants.ORGANIC_RF_CHANNEL: organic_rf_channel_values,
+        constants.CONTROL_VARIABLE: control_variable_names,
+        constants.NON_MEDIA_CHANNEL: non_media_channel_names,
+        constants.MEDIA_CHANNEL: media_channel_names,
+        constants.RF_CHANNEL: rf_channel_names,
+        constants.ORGANIC_MEDIA_CHANNEL: organic_media_channel_names,
+        constants.ORGANIC_RF_CHANNEL: organic_rf_channel_names,
     }
 
   def create_inference_data_dims(self) -> Mapping[str, Sequence[str]]:

meridian/model/model_test_data.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The Meridian Authors.
+# 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.
@@ -70,6 +70,10 @@ class WithInputDataSamples:
       _TEST_DIR,
       "sample_prior_media_only.nc",
   )
+  _TEST_SAMPLE_PRIOR_MEDIA_ONLY_NO_CONTROLS_PATH = os.path.join(
+      _TEST_DIR,
+      "sample_prior_media_only_no_controls.nc",
+  )
   _TEST_SAMPLE_PRIOR_RF_ONLY_PATH = os.path.join(
       _TEST_DIR,
       "sample_prior_rf_only.nc",
@@ -82,6 +86,10 @@ class WithInputDataSamples:
       _TEST_DIR,
       "sample_posterior_media_only.nc",
   )
+  _TEST_SAMPLE_POSTERIOR_MEDIA_ONLY_NO_CONTROLS_PATH = os.path.join(
+      _TEST_DIR,
+      "sample_posterior_media_only_no_controls.nc",
+  )
   _TEST_SAMPLE_POSTERIOR_RF_ONLY_PATH = os.path.join(
       _TEST_DIR,
       "sample_posterior_rf_only.nc",
@@ -172,6 +180,17 @@ class WithInputDataSamples:
             seed=0,
         )
     )
+    self.input_data_with_media_and_rf_no_controls = (
+        test_utils.sample_input_data_non_revenue_revenue_per_kpi(
+            n_geos=self._N_GEOS,
+            n_times=self._N_TIMES,
+            n_media_times=self._N_MEDIA_TIMES,
+            n_controls=None,
+            n_media_channels=self._N_MEDIA_CHANNELS,
+            n_rf_channels=self._N_RF_CHANNELS,
+            seed=0,
+        )
+    )
     self.short_input_data_with_media_only = (
         test_utils.sample_input_data_non_revenue_revenue_per_kpi(
             n_geos=self._N_GEOS,
@@ -182,6 +201,16 @@ class WithInputDataSamples:
             seed=0,
         )
     )
+    self.short_input_data_with_media_only_no_controls = (
+        test_utils.sample_input_data_non_revenue_revenue_per_kpi(
+            n_geos=self._N_GEOS,
+            n_times=self._N_TIMES_SHORT,
+            n_media_times=self._N_MEDIA_TIMES_SHORT,
+            n_controls=0,
+            n_media_channels=self._N_MEDIA_CHANNELS,
+            seed=0,
+        )
+    )
     self.short_input_data_with_rf_only = (
         test_utils.sample_input_data_non_revenue_revenue_per_kpi(
             n_geos=self._N_GEOS,
@@ -231,6 +260,9 @@ class WithInputDataSamples:
     test_prior_media_only = xr.open_dataset(
         self._TEST_SAMPLE_PRIOR_MEDIA_ONLY_PATH
     )
+    test_prior_media_only_no_controls = xr.open_dataset(
+        self._TEST_SAMPLE_PRIOR_MEDIA_ONLY_NO_CONTROLS_PATH
+    )
     test_prior_rf_only = xr.open_dataset(self._TEST_SAMPLE_PRIOR_RF_ONLY_PATH)
     self.test_dist_media_and_rf = collections.OrderedDict({
         param: tf.convert_to_tensor(test_prior_media_and_rf[param])
@@ -243,6 +275,18 @@ class WithInputDataSamples:
         for param in constants.COMMON_PARAMETER_NAMES
         + constants.MEDIA_PARAMETER_NAMES
     })
+    self.test_dist_media_only_no_controls = collections.OrderedDict({
+        param: tf.convert_to_tensor(test_prior_media_only_no_controls[param])
+        for param in (
+            set(
+                constants.COMMON_PARAMETER_NAMES
+                + constants.MEDIA_PARAMETER_NAMES
+            )
+            - set(
+                constants.CONTROL_PARAMETERS + constants.GEO_CONTROL_PARAMETERS
+            )
+        )
+    })
     self.test_dist_rf_only = collections.OrderedDict({
         param: tf.convert_to_tensor(test_prior_rf_only[param])
         for param in constants.COMMON_PARAMETER_NAMES
@@ -255,6 +299,9 @@ class WithInputDataSamples:
     test_posterior_media_only = xr.open_dataset(
         self._TEST_SAMPLE_POSTERIOR_MEDIA_ONLY_PATH
     )
+    test_posterior_media_only_no_controls = xr.open_dataset(
+        self._TEST_SAMPLE_POSTERIOR_MEDIA_ONLY_NO_CONTROLS_PATH
+    )
     test_posterior_rf_only = xr.open_dataset(
         self._TEST_SAMPLE_POSTERIOR_RF_ONLY_PATH
     )
@@ -273,6 +320,21 @@ class WithInputDataSamples:
         for param in constants.COMMON_PARAMETER_NAMES
         + constants.MEDIA_PARAMETER_NAMES
     }
+    posterior_params_to_tensors_media_only_no_controls = {
+        param: _convert_with_swap(
+            test_posterior_media_only_no_controls[param],
+            n_burnin=self._N_BURNIN,
+        )
+        for param in (
+            set(
+                constants.COMMON_PARAMETER_NAMES
+                + constants.MEDIA_PARAMETER_NAMES
+            )
+            - set(
+                constants.CONTROL_PARAMETERS + constants.GEO_CONTROL_PARAMETERS
+            )
+        )
+    }
     posterior_params_to_tensors_rf_only = {
         param: _convert_with_swap(
             test_posterior_rf_only[param], n_burnin=self._N_BURNIN
@@ -290,6 +352,18 @@ class WithInputDataSamples:
         "StructTuple",
         constants.COMMON_PARAMETER_NAMES + constants.MEDIA_PARAMETER_NAMES,
     )(**posterior_params_to_tensors_media_only)
+    self.test_posterior_states_media_only_no_controls = collections.namedtuple(
+        "StructTuple",
+        (
+            set(
+                constants.COMMON_PARAMETER_NAMES
+                + constants.MEDIA_PARAMETER_NAMES
+            )
+            - set(
+                constants.CONTROL_PARAMETERS + constants.GEO_CONTROL_PARAMETERS
+            )
+        ),
+    )(**posterior_params_to_tensors_media_only_no_controls)
     self.test_posterior_states_rf_only = collections.namedtuple(
         "StructTuple",
         constants.COMMON_PARAMETER_NAMES + constants.RF_PARAMETER_NAMES,