more tests

Author: juanitorduz

tests/mmm/test_mmm.py

@@ -1054,6 +1054,138 @@ def test_mmm_missing_dag_or_nodes(self):
         assert mmm.treatment_nodes is None, "Treatment nodes should default to None."
         assert mmm.outcome_node is None, "Outcome node should default to None."
 
+    def test_scaling_dict_with_missing_channel(
+        self, toy_X: pd.DataFrame, toy_y: pd.Series
+    ):
+        """Test scaling dict that only specifies target."""
+        from pymc_marketing.mmm.scaling import VariableScaling
+
+        mmm = MMM(
+            date_column="date",
+            channel_columns=["channel_1", "channel_2"],
+            adstock=GeometricAdstock(l_max=4),
+            saturation=LogisticSaturation(),
+            scaling={"target": VariableScaling(method="mean", dims=())},
+        )
+
+        mmm.build_model(X=toy_X, y=toy_y)
+        assert hasattr(mmm, "model")
+        assert mmm.scaling.channel.method == "max"
+
+    def test_scaling_dict_with_missing_target(
+        self, toy_X: pd.DataFrame, toy_y: pd.Series
+    ):
+        """Test scaling dict that only specifies channel."""
+        from pymc_marketing.mmm.scaling import VariableScaling
+
+        mmm = MMM(
+            date_column="date",
+            channel_columns=["channel_1", "channel_2"],
+            adstock=GeometricAdstock(l_max=4),
+            saturation=LogisticSaturation(),
+            scaling={"channel": VariableScaling(method="mean", dims=())},
+        )
+
+        mmm.build_model(X=toy_X, y=toy_y)
+        assert hasattr(mmm, "model")
+        assert mmm.scaling.target.method == "max"
+
+    def test_mean_scaling_method(self, toy_X: pd.DataFrame, toy_y: pd.Series):
+        """Test using mean scaling instead of max."""
+        from pymc_marketing.mmm.scaling import Scaling, VariableScaling
+
+        mmm = MMM(
+            date_column="date",
+            channel_columns=["channel_1", "channel_2"],
+            adstock=GeometricAdstock(l_max=4),
+            saturation=LogisticSaturation(),
+            scaling=Scaling(
+                target=VariableScaling(method="mean", dims=()),
+                channel=VariableScaling(method="mean", dims=()),
+            ),
+        )
+
+        mmm.build_model(X=toy_X, y=toy_y)
+        assert hasattr(mmm, "model")
+        assert mmm.target_scale > 0
+
+    def test_validation_disabled(self, toy_X: pd.DataFrame, toy_y: pd.Series):
+        """Test model with validation disabled."""
+        mmm = MMM(
+            date_column="date",
+            channel_columns=["channel_1", "channel_2"],
+            adstock=GeometricAdstock(l_max=4),
+            saturation=LogisticSaturation(),
+            validate_data=False,
+        )
+
+        mmm.build_model(X=toy_X, y=toy_y)
+        assert hasattr(mmm, "model")
+
+    def test_adstock_first_false(self, toy_X: pd.DataFrame, toy_y: pd.Series):
+        """Test model with saturation applied before adstock."""
+        mmm = MMM(
+            date_column="date",
+            channel_columns=["channel_1", "channel_2"],
+            adstock=GeometricAdstock(l_max=4),
+            saturation=LogisticSaturation(),
+            adstock_first=False,
+        )
+
+        mmm.build_model(X=toy_X, y=toy_y)
+        assert hasattr(mmm, "model")
+        assert mmm.adstock_first is False
+
+    def test_serializable_model_config_with_ndarray(self, mmm_fitted: MMM):
+        """Test _serializable_model_config converts ndarrays to lists."""
+        mmm_fitted.model_config["test_array"] = np.array([1, 2, 3])
+        mmm_fitted.model_config["nested"] = {"array": np.array([4, 5])}
+
+        serializable = mmm_fitted._serializable_model_config
+
+        assert isinstance(serializable["test_array"], list)
+        assert isinstance(serializable["nested"]["array"], list)
+
+    def test_model_config_formatting_with_nested_dicts(self):
+        """Test _model_config_formatting with nested dictionaries."""
+        config = {
+            "param1": [1, 2, 3],
+            "param2": {"nested": [4, 5, 6], "dims": ["a", "b"]},
+            "param3": {"deep": {"nested": [7, 8], "dims": ["c"]}},
+        }
+
+        formatted = MMM._model_config_formatting(config)
+
+        assert isinstance(formatted["param2"]["dims"], tuple)
+        assert isinstance(formatted["param3"]["deep"]["dims"], tuple)
+        assert isinstance(formatted["param1"], np.ndarray)
+        assert isinstance(formatted["param2"]["nested"], np.ndarray)
+
+    def test_deserialize_scaling_none(self):
+        """Test _deserialize_scaling with None input."""
+        result = MMM._deserialize_scaling(None)
+        assert result is None
+
+    def test_attrs_to_init_kwargs_with_defaults(self):
+        """Test attrs_to_init_kwargs with missing optional attributes."""
+        attrs = {
+            "model_config": "{}",
+            "date_column": '"date"',
+            "control_columns": "null",
+            "channel_columns": '["ch1"]',
+            "adstock": '{"lookup_name": "geometric", "l_max": 4, "normalize": true, "mode": "After"}',
+            "saturation": '{"lookup_name": "logistic"}',
+            "yearly_seasonality": "null",
+            "validate_data": "true",
+            "sampler_config": "{}",
+        }
+
+        kwargs = MMM.attrs_to_init_kwargs(attrs)
+
+        assert kwargs["adstock_first"] is True
+        assert kwargs["time_varying_intercept"] is False
+        assert kwargs["scaling"] is None
+
 
 def new_date_ranges_to_test():
     yield from [
@@ -2227,6 +2359,160 @@ def test_mixed_positive_negative_channels(self, toy_y: pd.Series):
         mmm.build_model(X=X, y=toy_y)
         assert hasattr(mmm, "model")
 
+    def test_date_conversion_error(self, toy_y: pd.Series):
+        """Test error when date cannot be converted."""
+        X = pd.DataFrame(
+            {
+                "date": [object() for _ in range(50)],
+                "channel_1": rng.integers(100, 1000, size=50),
+                "channel_2": rng.integers(100, 1000, size=50),
+            }
+        )
+
+        mmm = MMM(
+            date_column="date",
+            channel_columns=["channel_1", "channel_2"],
+            adstock=GeometricAdstock(l_max=4),
+            saturation=LogisticSaturation(),
+        )
+
+        with pytest.raises(ValueError, match="Could not convert"):
+            mmm.build_model(X=X, y=toy_y)
+
+    def test_transformer_fitting_failure(
+        self, toy_X: pd.DataFrame, toy_y: pd.Series, monkeypatch
+    ):
+        """Test handling when transformer fitting fails."""
+
+        def mock_failing_transformer(*args, **kwargs):
+            raise RuntimeError("Transformer failure")
+
+        mmm = MMM(
+            date_column="date",
+            channel_columns=["channel_1", "channel_2"],
+            adstock=GeometricAdstock(l_max=4),
+            saturation=LogisticSaturation(),
+        )
+
+        monkeypatch.setattr(mmm, "max_abs_scale_target_data", mock_failing_transformer)
+
+        with pytest.warns(UserWarning, match="Failed to fit transformers"):
+            mmm.build_model(X=toy_X, y=toy_y)
+
+    def test_causal_graph_with_warnings(self, toy_X: pd.DataFrame, toy_y: pd.Series):
+        """Test causal graph initialization with warnings when treatment nodes not specified."""
+        dag = """
+        digraph {
+            channel_1 -> y;
+            channel_2 -> y;
+            control_1 -> y;
+        }
+        """
+
+        with pytest.warns(UserWarning, match="No treatment nodes provided"):
+            mmm = MMM(
+                date_column="date",
+                channel_columns=["channel_1", "channel_2"],
+                control_columns=["control_1"],
+                adstock=GeometricAdstock(l_max=4),
+                saturation=LogisticSaturation(),
+                dag=dag,
+                outcome_node="y",
+            )
+
+        assert mmm.treatment_nodes == ["channel_1", "channel_2"]
+
+    def test_causal_graph_excludes_seasonality(
+        self, toy_X: pd.DataFrame, toy_y: pd.Series
+    ):
+        """Test that seasonality is excluded when not in adjustment set."""
+        dag = """
+        digraph {
+            channel_1 -> y;
+            channel_2 -> y;
+            control_1 -> y;
+        }
+        """
+
+        with pytest.warns(UserWarning, match="Yearly seasonality excluded"):
+            mmm = MMM(
+                date_column="date",
+                channel_columns=["channel_1", "channel_2"],
+                control_columns=["control_1"],
+                adstock=GeometricAdstock(l_max=4),
+                saturation=LogisticSaturation(),
+                yearly_seasonality=2,
+                dag=dag,
+                treatment_nodes=["channel_1", "channel_2"],
+                outcome_node="y",
+            )
+
+        assert mmm.yearly_seasonality is None
+
+    def test_x_data_not_dataframe_error(self, toy_y: pd.Series):
+        """Test error when X data is not a DataFrame."""
+        mmm = MMM(
+            date_column="date",
+            channel_columns=["channel_1", "channel_2"],
+            adstock=GeometricAdstock(l_max=4),
+            saturation=LogisticSaturation(),
+        )
+
+        mmm.preprocessed_data = {"X": np.array([[1, 2], [3, 4]]), "y": toy_y}
+
+        with pytest.raises(TypeError, match="must be a DataFrame"):
+            mmm._compute_scales()
+
+    def test_invalid_synth_dataset_granularity(
+        self, mmm_fitted: MMM, toy_X: pd.DataFrame
+    ):
+        """Test invalid time granularity in synthetic dataset creation."""
+        from xarray import DataArray
+
+        allocation = DataArray(
+            [1000, 2000],
+            dims=["channel"],
+            coords={"channel": mmm_fitted.channel_columns},
+        )
+
+        with pytest.raises(ValueError, match="Unsupported time granularity"):
+            mmm_fitted._create_synth_dataset(
+                df=toy_X,
+                date_column="date",
+                allocation_strategy=allocation,
+                channels=mmm_fitted.channel_columns,
+                controls=None,
+                target_col="y",
+                time_granularity="invalid_granularity",
+                time_length=10,
+                lag=4,
+            )
+
+    def test_invalid_allocation_dims(self, mmm_fitted: MMM, toy_X: pd.DataFrame):
+        """Test error when allocation strategy has wrong dimensions."""
+        from xarray import DataArray
+
+        allocation = DataArray(
+            [1000, 2000],
+            dims=["wrong_dim"],
+            coords={"wrong_dim": ["a", "b"]},
+        )
+
+        with pytest.raises(
+            ValueError, match="must have a single dimension named 'channel'"
+        ):
+            mmm_fitted._create_synth_dataset(
+                df=toy_X,
+                date_column="date",
+                allocation_strategy=allocation,
+                channels=mmm_fitted.channel_columns,
+                controls=None,
+                target_col="y",
+                time_granularity="weekly",
+                time_length=10,
+                lag=4,
+            )
+
 
 class TestMMMHelperMethods:
     """Tests for internal helper methods added during refactoring."""
@@ -2630,3 +2916,13 @@ def test_add_original_scale_deterministics(
         assert "control_contribution_original_scale" in deterministic_names
         assert "yearly_seasonality_contribution_original_scale" in deterministic_names
         assert "y_original_scale" in deterministic_names
+
+    def test_prepare_target_data_with_none_and_no_preprocessed(self, mmm_fitted: MMM):
+        """Test _prepare_target_data when y is None but preprocessed_data doesn't match expected types."""
+        mmm_fitted.preprocessed_data["y"] = "not_a_series_or_array"  # type: ignore
+
+        result = mmm_fitted._prepare_target_data(y=None, n_rows=10)
+
+        assert "target_data" in result or "target" in result
+        if "target_data" in result:
+            assert len(result["target_data"]) == 10