increase code coverage

Author: drbenvincent

causalpy/tests/test_integration_pymc_examples.py

@@ -1132,9 +1132,26 @@ def test_transfer_function_bayesian_adstock_only(mock_pymc_sample):
     assert isinstance(fig, plt.Figure)
     assert len(ax) == 2  # Should have 2 subplots
 
+    # Test plot_transforms (Bayesian-specific path)
+    fig_trans, ax_trans = result.plot_transforms()
+    assert isinstance(fig_trans, plt.Figure)
+    assert len(ax_trans) >= 1  # Should have at least 1 panel (adstock)
+
     # Test summary (should not raise)
     result.summary()
 
+    # Test effect() method (Bayesian-specific path)
+    effect_result = result.effect(
+        window=(df.index[0], df.index[-1]), channels=["treatment"], scale=0.0
+    )
+    assert "effect_df" in effect_result
+    assert "total_effect" in effect_result
+
+    # Test plot_effect() (Bayesian-specific path)
+    fig_eff, ax_eff = result.plot_effect(effect_result)
+    assert isinstance(fig_eff, plt.Figure)
+    assert len(ax_eff) == 2
+
     # Test that half_life posterior is reasonable (should be positive)
     half_life_samples = az.extract(result.model.idata, var_names=["half_life"])
     assert (half_life_samples > 0).all(), "Half-life should be positive"
@@ -1251,9 +1268,26 @@ def test_transfer_function_ar_bayesian(mock_pymc_sample):
     assert isinstance(fig, plt.Figure)
     assert len(ax) == 2  # Should have 2 subplots
 
+    # Test plot_transforms (Bayesian-specific path)
+    fig_trans, ax_trans = result.plot_transforms()
+    assert isinstance(fig_trans, plt.Figure)
+    assert len(ax_trans) >= 1  # Should have at least 1 panel (adstock)
+
     # Test summary (should not raise)
     result.summary()
 
+    # Test effect() method (Bayesian-specific path)
+    effect_result = result.effect(
+        window=(df.index[0], df.index[-1]), channels=["treatment"], scale=0.0
+    )
+    assert "effect_df" in effect_result
+    assert "total_effect" in effect_result
+
+    # Test plot_effect() (Bayesian-specific path)
+    fig_eff, ax_eff = result.plot_effect(effect_result)
+    assert isinstance(fig_eff, plt.Figure)
+    assert len(ax_eff) == 2
+
     # Test that half_life posterior is reasonable (should be positive)
     half_life_samples = az.extract(result.model.idata, var_names=["half_life"])
     assert (half_life_samples > 0).all(), "Half-life should be positive"

causalpy/tests/test_transfer_function_its.py

@@ -1818,3 +1818,58 @@ def test_effect_with_arimax_model(self):
         fig, ax = result.plot_effect(effect_result)
         assert isinstance(fig, plt.Figure)
         plt.close(fig)
+
+
+class TestPlotTransformsEdgeCases:
+    """Test plot_transforms edge cases."""
+
+    def test_plot_transforms_with_lag(self):
+        """Test that lag transforms are applied correctly in build_treatment_matrix."""
+        np.random.seed(42)
+        n = 50
+        t = np.arange(n)
+        dates = pd.date_range("2020-01-01", periods=n, freq="W")
+
+        treatment_raw = 50 + np.random.uniform(-10, 10, n)
+        treatment_raw = np.maximum(treatment_raw, 0)
+        y = 100.0 + 0.5 * t + treatment_raw + np.random.normal(0, 5, n)
+
+        df = pd.DataFrame({"date": dates, "t": t, "y": y, "treatment": treatment_raw})
+        df = df.set_index("date")
+
+        model = TransferFunctionOLS(
+            saturation_type=None,
+            adstock_grid={"half_life": [3]},
+            estimation_method="grid",
+            error_model="hac",
+        )
+
+        result = GradedInterventionTimeSeries(
+            data=df,
+            y_column="y",
+            treatment_names=["treatment"],
+            base_formula="1 + t",
+            model=model,
+        )
+
+        # Manually add a lag to the treatment and test _build_treatment_matrix
+        from causalpy.transforms import DiscreteLag
+
+        treatments_with_lag = []
+        for treatment in result.treatments:
+            # Create a new treatment object with lag
+            treatment_lagged = Treatment(
+                name=treatment.name,
+                saturation=treatment.saturation,
+                adstock=treatment.adstock,
+                lag=DiscreteLag(k=1),  # Add 1-period lag
+            )
+            treatments_with_lag.append(treatment_lagged)
+
+        # Test _build_treatment_matrix with lag
+        Z, labels = result._build_treatment_matrix(df, treatments_with_lag)
+
+        assert Z.shape == (n, 1)
+        assert labels == ["treatment"]
+        # First value should be 0 due to lag
+        assert Z[0, 0] == 0