add warning when bandwidth parameter leads to only a few remaining datapoints

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Optional, Union
 
 import arviz as az
@@ -586,6 +587,11 @@ def __init__(
             fmin = self.treatment_threshold - self.bandwidth
             fmax = self.treatment_threshold + self.bandwidth
             filtered_data = self.data.query(f"{fmin} <= x <= {fmax}")
+            if len(filtered_data) <= 10:
+                warnings.warn(
+                    f"Choice of bandwidth parameter has lead to only {len(filtered_data)} remaining datapoints. Consider increasing the bandwidth parameter.",  # noqa: E501
+                    UserWarning,
+                )
             y, X = dmatrices(formula, filtered_data)
         else:
             y, X = dmatrices(formula, self.data)

causalpy/skl_experiments.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Optional
 
 import matplotlib.pyplot as plt
@@ -391,6 +392,11 @@ def __init__(
             fmin = self.treatment_threshold - self.bandwidth
             fmax = self.treatment_threshold + self.bandwidth
             filtered_data = self.data.query(f"{fmin} <= x <= {fmax}")
+            if len(filtered_data) <= 10:
+                warnings.warn(
+                    f"Choice of bandwidth parameter has lead to only {len(filtered_data)} remaining datapoints. Consider increasing the bandwidth parameter.",  # noqa: E501
+                    UserWarning,
+                )
             y, X = dmatrices(formula, filtered_data)
         else:
             y, X = dmatrices(formula, self.data)