Support adstock-only and saturation-only transfer functions

Refactored transfer function parameter estimation and plotting to allow models with only adstock or only saturation transforms (or both). Updated validation logic, grid/optimize routines, and plotting to handle optional transforms. Added comprehensive tests for all transform configurations and clarified documentation and notebook to demonstrate adstock-only use case.

Author: drbenvincent

causalpy/experiments/graded_intervention_its.py

@@ -764,11 +764,11 @@ def plot_transforms(
         x_range=None,
         **kwargs,
     ) -> Tuple[plt.Figure, np.ndarray]:
-        """Plot estimated saturation and adstock transformation curves.
+        """Plot estimated transformation curves (saturation and/or adstock).
 
-        Creates a 2-panel figure showing:
-        1. Saturation curve (input exposure -> saturated exposure)
-        2. Adstock weights over time (lag distribution)
+        Creates a figure with 1-2 panels depending on which transforms are present:
+        - Saturation curve (input exposure -> saturated exposure) if saturation exists
+        - Adstock weights over time (lag distribution) if adstock exists
 
         Parameters
         ----------
@@ -784,8 +784,8 @@ def plot_transforms(
         Returns
         -------
         fig : matplotlib.figure.Figure
-        ax : array of matplotlib.axes.Axes
-            Array of 2 axes objects (left: saturation, right: adstock).
+        ax : list of matplotlib.axes.Axes
+            List of axes objects (1 or 2 panels depending on which transforms exist).
 
         Examples
         --------
@@ -810,13 +810,33 @@ def plot_transforms(
         est_saturation = treatment.saturation
         est_adstock = treatment.adstock
 
-        # Create 2-panel subplot
-        fig, axes = plt.subplots(1, 2, figsize=(14, 5))
+        # Check which transforms exist
+        has_saturation = est_saturation is not None
+        has_adstock = est_adstock is not None
+
+        if not has_saturation and not has_adstock:
+            raise ValueError(
+                "No transforms to plot (both saturation and adstock are None). "
+                "At least one transform must be specified."
+            )
+
+        # Determine number of panels based on available transforms
+        n_panels = int(has_saturation) + int(has_adstock)
+
+        # Create subplot with appropriate number of panels
+        fig, axes = plt.subplots(1, n_panels, figsize=(7 * n_panels, 5))
+
+        # Make axes a list for consistent indexing
+        if n_panels == 1:
+            axes = [axes]
+
+        panel_idx = 0
 
         # ============================================================================
-        # LEFT PLOT: Saturation curves
+        # SATURATION PLOT (if present)
         # ============================================================================
         if est_saturation is not None:
+            ax = axes[panel_idx]
             # Determine x range
             if x_range is None:
                 # Use range from data
@@ -831,7 +851,7 @@ def plot_transforms(
             # Plot true saturation if provided
             if true_saturation is not None:
                 y_true_sat = true_saturation.apply(x_sat)
-                axes[0].plot(
+                ax.plot(
                     x_sat,
                     y_true_sat,
                     "k--",
@@ -842,13 +862,13 @@ def plot_transforms(
 
             # Plot estimated saturation
             y_est_sat = est_saturation.apply(x_sat)
-            axes[0].plot(x_sat, y_est_sat, "C0-", linewidth=2.5, label="Estimated")
+            ax.plot(x_sat, y_est_sat, "C0-", linewidth=2.5, label="Estimated")
 
-            axes[0].set_xlabel(f"{treatment.name} (raw)", fontsize=11)
-            axes[0].set_ylabel("Saturated Value", fontsize=11)
-            axes[0].set_title("Saturation Function", fontsize=12, fontweight="bold")
-            axes[0].legend(fontsize=LEGEND_FONT_SIZE, framealpha=0.9)
-            axes[0].grid(True, alpha=0.3)
+            ax.set_xlabel(f"{treatment.name} (raw)", fontsize=11)
+            ax.set_ylabel("Saturated Value", fontsize=11)
+            ax.set_title("Saturation Function", fontsize=12, fontweight="bold")
+            ax.legend(fontsize=LEGEND_FONT_SIZE, framealpha=0.9)
+            ax.grid(True, alpha=0.3)
 
             # Add parameter text
             est_params = est_saturation.get_params()
@@ -864,30 +884,22 @@ def plot_transforms(
                     if key not in ["alpha", "l_max", "normalize"]:
                         param_text += f"  {key}={val:.2f}\n"
 
-            axes[0].text(
+            ax.text(
                 0.05,
                 0.95,
                 param_text.strip(),
-                transform=axes[0].transAxes,
+                transform=ax.transAxes,
                 fontsize=9,
                 verticalalignment="top",
                 bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
             )
-        else:
-            axes[0].text(
-                0.5,
-                0.5,
-                "No saturation transform",
-                ha="center",
-                va="center",
-                transform=axes[0].transAxes,
-            )
-            axes[0].set_title("Saturation Function", fontsize=12, fontweight="bold")
+            panel_idx += 1
 
         # ============================================================================
-        # RIGHT PLOT: Adstock weights
+        # ADSTOCK PLOT (if present)
         # ============================================================================
         if est_adstock is not None:
+            ax = axes[panel_idx]
             est_adstock_params = est_adstock.get_params()
             l_max = est_adstock_params.get("l_max", 12)
             lags = np.arange(l_max + 1)
@@ -908,15 +920,15 @@ def plot_transforms(
                     true_weights = true_weights / true_weights.sum()
 
                 width = 0.35
-                axes[1].bar(
+                ax.bar(
                     lags - width / 2,
                     true_weights,
                     width,
                     alpha=0.8,
                     label="True",
                     color="gray",
                 )
-                axes[1].bar(
+                ax.bar(
                     lags + width / 2,
                     est_weights,
                     width,
@@ -925,15 +937,15 @@ def plot_transforms(
                     color="C0",
                 )
             else:
-                axes[1].bar(lags, est_weights, alpha=0.7, color="C0", label="Estimated")
+                ax.bar(lags, est_weights, alpha=0.7, color="C0", label="Estimated")
 
-            axes[1].set_xlabel("Lag (periods)", fontsize=11)
-            axes[1].set_ylabel("Adstock Weight", fontsize=11)
-            axes[1].set_title(
+            ax.set_xlabel("Lag (periods)", fontsize=11)
+            ax.set_ylabel("Adstock Weight", fontsize=11)
+            ax.set_title(
                 "Adstock Function (Carryover Effect)", fontsize=12, fontweight="bold"
             )
-            axes[1].legend(fontsize=LEGEND_FONT_SIZE, framealpha=0.9)
-            axes[1].grid(True, alpha=0.3, axis="y")
+            ax.legend(fontsize=LEGEND_FONT_SIZE, framealpha=0.9)
+            ax.grid(True, alpha=0.3, axis="y")
 
             # Add parameter text
             param_text = "Estimated:\n"
@@ -948,28 +960,16 @@ def plot_transforms(
                 param_text += f"  half_life={half_life_true:.2f}\n"
                 param_text += f"  alpha={true_alpha:.3f}\n"
 
-            axes[1].text(
+            ax.text(
                 0.95,
                 0.95,
                 param_text.strip(),
-                transform=axes[1].transAxes,
+                transform=ax.transAxes,
                 fontsize=9,
                 verticalalignment="top",
                 horizontalalignment="right",
                 bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
             )
-        else:
-            axes[1].text(
-                0.5,
-                0.5,
-                "No adstock transform",
-                ha="center",
-                va="center",
-                transform=axes[1].transAxes,
-            )
-            axes[1].set_title(
-                "Adstock Function (Carryover Effect)", fontsize=12, fontweight="bold"
-            )
 
         plt.tight_layout()
         return fig, axes

causalpy/skl_models.py

@@ -195,26 +195,32 @@ def __init__(
 
         # Validate estimation method and required parameters
         if estimation_method == "grid":
-            if saturation_grid is None:
+            # At least one transform must be specified
+            if saturation_grid is None and adstock_grid is None:
                 raise ValueError(
-                    "saturation_grid is required for grid search method. "
-                    "E.g., saturation_grid={'slope': [1.0, 2.0], 'kappa': [3, 5]}"
+                    "At least one of saturation_grid or adstock_grid must be provided for grid search. "
+                    "To use only adstock: set saturation_type=None and provide adstock_grid. "
+                    "To use only saturation: provide saturation_grid and set adstock_grid=None."
                 )
-            if adstock_grid is None:
+            # If saturation_type is specified, grid must be provided
+            if saturation_type is not None and saturation_grid is None:
                 raise ValueError(
-                    "adstock_grid is required for grid search method. "
-                    "E.g., adstock_grid={'half_life': [2, 3, 4]}"
+                    f"saturation_grid is required when saturation_type='{saturation_type}'. "
+                    "E.g., saturation_grid={'lam': [0.2, 0.5, 0.8]}"
                 )
         elif estimation_method == "optimize":
-            if saturation_bounds is None:
+            # At least one transform must be specified
+            if saturation_bounds is None and adstock_bounds is None:
                 raise ValueError(
-                    "saturation_bounds is required for optimize method. "
-                    "E.g., saturation_bounds={'slope': (0.5, 5.0), 'kappa': (2, 10)}"
+                    "At least one of saturation_bounds or adstock_bounds must be provided for optimize method. "
+                    "To use only adstock: set saturation_type=None and provide adstock_bounds. "
+                    "To use only saturation: provide saturation_bounds and set adstock_bounds=None."
                 )
-            if adstock_bounds is None:
+            # If saturation_type is specified, bounds must be provided
+            if saturation_type is not None and saturation_bounds is None:
                 raise ValueError(
-                    "adstock_bounds is required for optimize method. "
-                    "E.g., adstock_bounds={'half_life': (1, 10)}"
+                    f"saturation_bounds is required when saturation_type='{saturation_type}'. "
+                    "E.g., saturation_bounds={'lam': (0.1, 1.0)}"
                 )
         else:
             raise ValueError(