mypy

Author: juanitorduz

pymc_marketing/mmm/mmm.py

@@ -667,16 +667,16 @@ def _build_control_contribution(self) -> pt.TensorVariable | None:
         pt.TensorVariable | None
             Control contribution variable or None if no controls
         """
+        if self.control_columns is None or len(self.control_columns) == 0:
+            return None
+
         X_data = self.preprocessed_data["X"]
-        has_controls = (
-            self.control_columns is not None
-            and len(self.control_columns) > 0
-            and isinstance(X_data, pd.DataFrame)
-            and all(column in X_data.columns for column in self.control_columns)
-        )
+        if not isinstance(X_data, pd.DataFrame):
+            raise TypeError("X data must be a DataFrame for control contribution")
 
-        if not has_controls:
-            return None
+        if not all(column in X_data.columns for column in self.control_columns):
+            missing_cols = set(self.control_columns) - set(X_data.columns)
+            raise ValueError(f"Control columns {missing_cols} not found in X data")
 
         if self.model_config["gamma_control"].dims != ("control",):
             self.model_config["gamma_control"].dims = "control"
@@ -858,10 +858,12 @@ def build_model(  # type: ignore[override]
             if not isinstance(X_data, pd.DataFrame):
                 raise TypeError("X data must be a DataFrame")
 
+            y_data = self.preprocessed_data["y"]
+            if not isinstance(y_data, (pd.Series, np.ndarray)):
+                raise TypeError("y data must be a Series or ndarray")
+
             channel_data_scaled, target_scaled, _, target_scale_ = (
-                self._create_scaled_data_variables(
-                    X_data[self.channel_columns], self.preprocessed_data["y"]
-                )
+                self._create_scaled_data_variables(X_data[self.channel_columns], y_data)
             )
 
             # Create time index if needed
@@ -1152,8 +1154,13 @@ def _prepare_target_data(
             Dictionary with target data ready for pm.set_data
         """
         if y is None:
-            dtype = self.preprocessed_data["y"].dtype  # type: ignore
-            y_data = np.zeros(n_rows, dtype=dtype)
+            # When y is None, create zeros array matching the type of preprocessed y
+            y_preprocessed = self.preprocessed_data["y"]
+            if isinstance(y_preprocessed, (pd.Series, np.ndarray)):
+                y_data = np.zeros(n_rows, dtype=np.asarray(y_preprocessed).dtype)
+            else:
+                # Default to float64 if type is unknown
+                y_data = np.zeros(n_rows, dtype="float64")
         elif isinstance(y, pd.Series):
             y_data = y.to_numpy()
         elif isinstance(y, np.ndarray):