Handle 0-dim channel_xr in create_zero_dataset (#1890)

* Handle 0-dim channel_xr in create_zero_dataset

Update create_zero_dataset to support channel_xr Datasets with no dimensions by broadcasting scalar values for each channel across all rows. Add tests to verify correct behavior when channel_xr is 0-dim and when some channels are missing.

* Test missing lines

Author: cetagostini

pymc_marketing/mmm/utils.py

@@ -349,26 +349,41 @@ def create_zero_dataset(
         if date_col in channel_xr.dims:
             raise ValueError("`channel_xr` must NOT include the date dimension.")
 
-        # --- 4.3 Convert to DataFrame & merge ----------------------------------
-        channel_df = channel_xr.to_dataframe().reset_index()
-
-        # Left-join on every dimension; suffix prevents collisions during merge
-        pred_df = pred_df.merge(
-            channel_df,
-            on=dim_cols,
-            how="left",
-            suffixes=("", "_chan"),
-        )
+        # --- 4.3 Inject constants ----------------------------------------------
+        # Special-case: when there are NO dims (e.g., only channel dimension in the
+        # allocation which was pivoted into variables), xarray can't create an index
+        # for to_dataframe(). In this scenario, simply broadcast scalar values
+        # across all rows.
+        if len(channel_xr.dims) == 0:
+            for ch in channel_cols:
+                if ch in channel_xr.data_vars:
+                    # assign scalar value across all rows
+                    try:
+                        pred_df[ch] = channel_xr[ch].item()
+                    except Exception:
+                        pred_df[ch] = channel_xr[ch].values
+        else:
+            # Convert to DataFrame & merge when dims are present
+            channel_df = channel_xr.to_dataframe().reset_index()
+
+            # Left-join on every dimension; suffix prevents collisions during merge
+            pred_df = pred_df.merge(
+                channel_df,
+                on=dim_cols,
+                how="left",
+                suffixes=("", "_chan"),
+            )
 
         # --- 4.4 Copy merged values into official channel columns --------------
-        for ch in channel_cols:
-            chan_col = f"{ch}_chan"
-            if chan_col in pred_df.columns:
-                pred_df[ch] = pred_df[chan_col]
-                pred_df.drop(columns=chan_col, inplace=True)
-
-        # Replace any remaining NaNs introduced by the merge
-        pred_df[channel_cols] = pred_df[channel_cols].fillna(0.0)
+        if len(channel_xr.dims) != 0:
+            for ch in channel_cols:
+                chan_col = f"{ch}_chan"
+                if chan_col in pred_df.columns:
+                    pred_df[ch] = pred_df[chan_col]
+                    pred_df.drop(columns=chan_col, inplace=True)
+
+            # Replace any remaining NaNs introduced by the merge
+            pred_df[channel_cols] = pred_df[channel_cols].fillna(0.0)
 
     # ---- 5. Bring in any “other” columns from the training data ----------------
     other_cols = [

tests/mmm/test_utils.py

@@ -507,6 +507,47 @@ def test_create_zero_dataset_error_cases(self):
         ):
             create_zero_dataset(model, start_date, end_date, date_dim_xr)
 
+    def test_create_zero_dataset_channel_xr_includes_date_specific_error(self):
+        """Ensure we hit the explicit date-dimension error when date is an allowed model dim."""
+
+        class FakeMMM_DateDim:
+            def __init__(self):
+                dates = pd.date_range("2022-01-01", "2022-01-10", freq="D")
+                self.X = pd.DataFrame(
+                    {
+                        "date": dates,
+                        "channel1": np.random.rand(10) * 10,
+                        "channel2": np.random.rand(10) * 5,
+                    }
+                )
+                self.date_column = "date"
+                self.channel_columns = ["channel1", "channel2"]
+                self.control_columns = []
+                # Include 'date' as a model dim so the invalid-dims check passes,
+                # and we can assert on the specific date-dimension error.
+                self.dims = ["date"]
+
+                class FakeAdstock:
+                    l_max = 1
+
+                self.adstock = FakeAdstock()
+
+        model = FakeMMM_DateDim()
+        start_date = "2022-02-01"
+        end_date = "2022-02-03"
+
+        channel_with_date = xr.Dataset(
+            data_vars={
+                "channel1": ("date", np.array([1.0, 2.0])),
+            },
+            coords={"date": pd.date_range("2022-01-01", periods=2, freq="D")},
+        )
+
+        with pytest.raises(
+            ValueError, match=r"`channel_xr` must NOT include the date dimension\."
+        ):
+            create_zero_dataset(model, start_date, end_date, channel_with_date)
+
     def test_create_zero_dataset_no_dims(self):
         """Test create_zero_dataset with a model that has no dimensions."""
 
@@ -549,3 +590,89 @@ def test_create_zero_dataset_empty_date_range_error(self):
 
         with pytest.raises(ValueError, match="Generated date range is empty"):
             create_zero_dataset(model, start_date, end_date)
+
+    def test_create_zero_dataset_channel_xr_no_dims_all_channels(self):
+        """Channel-only allocation: channel_xr is a 0-dim Dataset with per-channel scalars."""
+
+        class FakeMMM_NoDims:
+            def __init__(self):
+                dates = pd.date_range("2022-01-01", "2022-01-10", freq="D")
+                self.X = pd.DataFrame(
+                    {
+                        "date": dates,
+                        "channel1": np.random.rand(10) * 10,
+                        "channel2": np.random.rand(10) * 5,
+                    }
+                )
+                self.date_column = "date"
+                self.channel_columns = ["channel1", "channel2"]
+                self.control_columns = []
+                self.dims = []  # No dimensions
+
+                class FakeAdstock:
+                    l_max = 3
+
+                self.adstock = FakeAdstock()
+
+        model = FakeMMM_NoDims()
+        start_date = "2022-02-01"
+        end_date = "2022-02-05"
+
+        # 0-dim Dataset: variables are channels with scalar values
+        channel_values = xr.Dataset(
+            data_vars={
+                "channel1": 100.0,
+                "channel2": 200.0,
+            }
+        )
+
+        result = create_zero_dataset(model, start_date, end_date, channel_values)
+
+        # (5 + 3) days = 8 rows
+        assert len(result) == 8
+        assert np.all(result["channel1"] == 100.0)
+        assert np.all(result["channel2"] == 200.0)
+
+    def test_create_zero_dataset_channel_xr_no_dims_missing_channel(self):
+        """Channel-only allocation with missing channel var should warn and leave others at 0."""
+
+        class FakeMMM_NoDims:
+            def __init__(self):
+                dates = pd.date_range("2022-01-01", "2022-01-10", freq="D")
+                self.X = pd.DataFrame(
+                    {
+                        "date": dates,
+                        "channel1": np.random.rand(10) * 10,
+                        "channel2": np.random.rand(10) * 5,
+                    }
+                )
+                self.date_column = "date"
+                self.channel_columns = ["channel1", "channel2"]
+                self.control_columns = []
+                self.dims = []
+
+                class FakeAdstock:
+                    l_max = 2
+
+                self.adstock = FakeAdstock()
+
+        model = FakeMMM_NoDims()
+        start_date = "2022-02-01"
+        end_date = "2022-02-03"
+
+        # Provide only one channel as scalar variable in 0-dim Dataset
+        channel_values = xr.Dataset(
+            data_vars={
+                "channel1": 50.0,
+            }
+        )
+
+        with pytest.warns(
+            UserWarning, match="does not supply values for \\['channel2'\\]"
+        ):
+            result = create_zero_dataset(model, start_date, end_date, channel_values)
+
+        # (3 + 2) days = 5 rows
+        assert len(result) == 5
+        assert np.all(result["channel1"] == 50.0)
+        assert np.all(result["channel2"] == 0.0)