add round_to kwarg to plot methods

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -23,10 +23,13 @@
 from patsy import build_design_matrices, dmatrices
 from sklearn.linear_model import LinearRegression as sk_lin_reg
 
-from causalpy.custom_exceptions import BadIndexException  # NOQA
-from causalpy.custom_exceptions import DataException, FormulaException
+from causalpy.custom_exceptions import (
+    BadIndexException,  # NOQA
+    DataException,
+    FormulaException,
+)
 from causalpy.plot_utils import plot_xY
-from causalpy.utils import _is_variable_dummy_coded
+from causalpy.utils import _is_variable_dummy_coded, round_num
 
 LEGEND_FONT_SIZE = 12
 az.style.use("arviz-darkgrid")
@@ -228,7 +231,7 @@ def _input_validation(self, data, treatment_time):
                 "If data.index is not DatetimeIndex, treatment_time must be pd.Timestamp."  # noqa: E501
             )
 
-    def plot(self, counterfactual_label="Counterfactual", **kwargs):
+    def plot(self, counterfactual_label="Counterfactual", round_to=None, **kwargs):
         """
         Plot the results
         """
@@ -275,8 +278,8 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs):
 
         ax[0].set(
             title=f"""
-            Pre-intervention Bayesian $R^2$: {self.score.r2:.3f}
-            (std = {self.score.r2_std:.3f})
+            Pre-intervention Bayesian $R^2$: {round_num(self.score.r2, round_to)}
+            (std = {round_num(self.score.r2_std, round_to)})
             """
         )
 
@@ -580,7 +583,7 @@ def _input_validation(self):
                 coded. Consisting of 0's and 1's only."""
             )
 
-    def plot(self):
+    def plot(self, round_to=None):
         """Plot the results.
         Creating the combined mean + HDI legend entries is a bit involved.
         """
@@ -658,7 +661,7 @@ def plot(self):
         # formatting
         ax.set(
             xticks=self.x_pred_treatment[self.time_variable_name].values,
-            title=self._causal_impact_summary_stat(),
+            title=self._causal_impact_summary_stat(round_to),
         )
         ax.legend(
             handles=(h_tuple for h_tuple in handles),
@@ -711,11 +714,14 @@ def _plot_causal_impact_arrow(self, ax):
             va="center",
         )
 
-    def _causal_impact_summary_stat(self) -> str:
+    def _causal_impact_summary_stat(self, round_to=None) -> str:
         """Computes the mean and 94% credible interval bounds for the causal impact."""
         percentiles = self.causal_impact.quantile([0.03, 1 - 0.03]).values
-        ci = "$CI_{94\\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]"
-        causal_impact = f"{self.causal_impact.mean():.2f}, "
+        ci = (
+            "$CI_{94\\%}$"
+            + f"[{round_num(percentiles[0], round_to)}, {round_num(percentiles[1], round_to)}]"
+        )
+        causal_impact = f"{round_num(self.causal_impact.mean(), round_to)}, "
         return f"Causal impact = {causal_impact + ci}"
 
     def summary(self) -> None:
@@ -893,7 +899,7 @@ def _is_treated(self, x):
         """
         return np.greater_equal(x, self.treatment_threshold)
 
-    def plot(self):
+    def plot(self, round_to=None):
         """
         Plot the results
         """
@@ -918,12 +924,15 @@ def plot(self):
         labels = ["Posterior mean"]
 
         # create strings to compose title
-        title_info = f"{self.score.r2:.3f} (std = {self.score.r2_std:.3f})"
+        title_info = f"{round_num(self.score.r2, round_to)} (std = {round_num(self.score.r2_std, round_to)})"
         r2 = f"Bayesian $R^2$ on all data = {title_info}"
         percentiles = self.discontinuity_at_threshold.quantile([0.03, 1 - 0.03]).values
-        ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]"
+        ci = (
+            r"$CI_{94\%}$"
+            + f"[{round_num(percentiles[0], round_to)}, {round_num(percentiles[1], round_to)}]"
+        )
         discon = f"""
-            Discontinuity at threshold = {self.discontinuity_at_threshold.mean():.2f},
+            Discontinuity at threshold = {round_num(self.discontinuity_at_threshold.mean(), round_to)},
             """
         ax.set(title=r2 + "\n" + discon + ci)
         # Intervention line
@@ -1104,7 +1113,7 @@ def _is_treated(self, x):
         """Returns ``True`` if `x` is greater than or equal to the treatment threshold."""  # noqa: E501
         return np.greater_equal(x, self.kink_point)
 
-    def plot(self):
+    def plot(self, round_to=None):
         """
         Plot the results
         """
@@ -1129,12 +1138,15 @@ def plot(self):
         labels = ["Posterior mean"]
 
         # create strings to compose title
-        title_info = f"{self.score.r2:.3f} (std = {self.score.r2_std:.3f})"
+        title_info = f"{round_num(self.score.r2, round_to)} (std = {round_num(self.score.r2_std, round_to)})"
         r2 = f"Bayesian $R^2$ on all data = {title_info}"
         percentiles = self.gradient_change.quantile([0.03, 1 - 0.03]).values
-        ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]"
+        ci = (
+            r"$CI_{94\%}$"
+            + f"[{round_num(percentiles[0], round_to)}, {round_num(percentiles[1], round_to)}]"
+        )
         grad_change = f"""
-            Change in gradient = {self.gradient_change.mean():.2f},
+            Change in gradient = {round_num(self.gradient_change.mean(), round_to)},
             """
         ax.set(title=r2 + "\n" + grad_change + ci)
         # Intervention line
@@ -1292,7 +1304,7 @@ def _input_validation(self) -> None:
                 """
             )
 
-    def plot(self):
+    def plot(self, round_to=None):
         """Plot the results"""
         fig, ax = plt.subplots(
             2, 1, figsize=(7, 9), gridspec_kw={"height_ratios": [3, 1]}
@@ -1339,18 +1351,21 @@ def plot(self):
         )
 
         # Plot estimated caual impact / treatment effect
-        az.plot_posterior(self.causal_impact, ref_val=0, ax=ax[1])
+        az.plot_posterior(self.causal_impact, ref_val=0, ax=ax[1], round_to=round_to)
         ax[1].set(title="Estimated treatment effect")
         return fig, ax
 
-    def _causal_impact_summary_stat(self) -> str:
+    def _causal_impact_summary_stat(self, round_to) -> str:
         """Computes the mean and 94% credible interval bounds for the causal impact."""
         percentiles = self.causal_impact.quantile([0.03, 1 - 0.03]).values
-        ci = r"$CI_{94%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]"
+        ci = (
+            r"$CI_{94%}$"
+            + f"[{round_num(percentiles[0], round_to)}, {round_num(percentiles[1], round_to)}]"
+        )
         causal_impact = f"{self.causal_impact.mean():.2f}, "
         return f"Causal impact = {causal_impact + ci}"
 
-    def summary(self) -> None:
+    def summary(self, round_to=None) -> None:
         """
         Print text output summarising the results
         """
@@ -1359,7 +1374,7 @@ def summary(self) -> None:
         print(f"Formula: {self.formula}")
         print("\nResults:")
         # TODO: extra experiment specific outputs here
-        print(self._causal_impact_summary_stat())
+        print(self._causal_impact_summary_stat(round_to))
         self.print_coefficients()
 
     def _get_treatment_effect_coeff(self) -> str:

causalpy/skl_experiments.py

@@ -17,6 +17,8 @@
 import seaborn as sns
 from patsy import build_design_matrices, dmatrices
 
+from causalpy.utils import round_num
+
 LEGEND_FONT_SIZE = 12
 
 
@@ -113,7 +115,7 @@ def __init__(
         # cumulative impact post
         self.post_impact_cumulative = np.cumsum(self.post_impact)
 
-    def plot(self, counterfactual_label="Counterfactual", **kwargs):
+    def plot(self, counterfactual_label="Counterfactual", round_to=None, **kwargs):
         """Plot experiment results"""
         fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8))
 
@@ -128,7 +130,9 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs):
             ls=":",
             c="k",
         )
-        ax[0].set(title=f"$R^2$ on pre-intervention data = {self.score:.3f}")
+        ax[0].set(
+            title=f"$R^2$ on pre-intervention data = {round_num(self.score, round_to)}"
+        )
 
         ax[1].plot(self.datapre.index, self.pre_impact, "k.")
         ax[1].plot(
@@ -258,9 +262,11 @@ class SyntheticControl(PrePostFit):
     ... )
     """
 
-    def plot(self, plot_predictors=False, **kwargs):
+    def plot(self, plot_predictors=False, round_to=None, **kwargs):
         """Plot the results"""
-        fig, ax = super().plot(counterfactual_label="Synthetic control", **kwargs)
+        fig, ax = super().plot(
+            counterfactual_label="Synthetic control", round_to=round_to, **kwargs
+        )
         if plot_predictors:
             # plot control units as well
             ax[0].plot(self.datapre.index, self.pre_X, "-", c=[0.8, 0.8, 0.8], zorder=1)
@@ -397,7 +403,7 @@ def __init__(
         # TODO: THIS IS NOT YET CORRECT
         self.causal_impact = self.y_pred_treatment[1] - self.y_pred_counterfactual[0]
 
-    def plot(self):
+    def plot(self, round_to=None):
         """Plot results"""
         fig, ax = plt.subplots()
 
@@ -462,7 +468,7 @@ def plot(self):
             xlim=[-0.05, 1.1],
             xticks=[0, 1],
             xticklabels=["pre", "post"],
-            title=f"Causal impact = {self.causal_impact[0]:.2f}",
+            title=f"Causal impact = {round_num(self.causal_impact[0], round_to)}",
         )
         ax.legend(fontsize=LEGEND_FONT_SIZE)
         return (fig, ax)
@@ -607,7 +613,7 @@ def _is_treated(self, x):
         """
         return np.greater_equal(x, self.treatment_threshold)
 
-    def plot(self):
+    def plot(self, round_to=None):
         """Plot results"""
         fig, ax = plt.subplots()
         # Plot raw data
@@ -627,8 +633,8 @@ def plot(self):
             label="model fit",
         )
         # create strings to compose title
-        r2 = f"$R^2$ on all data = {self.score:.3f}"
-        discon = f"Discontinuity at threshold = {self.discontinuity_at_threshold:.2f}"
+        r2 = f"$R^2$ on all data = {round_num(self.score, round_to)}"
+        discon = f"Discontinuity at threshold = {round_num(self.discontinuity_at_threshold, round_to)}"
         ax.set(title=r2 + "\n" + discon)
         # Intervention line
         ax.axvline(

causalpy/utils.py

@@ -1,6 +1,7 @@
 """
 Utility functions
 """
+import numpy as np
 import pandas as pd
 
 
@@ -13,3 +14,37 @@ def _is_variable_dummy_coded(series: pd.Series) -> bool:
 def _series_has_2_levels(series: pd.Series) -> bool:
     """Check that the variable in the provided Series has 2 levels"""
     return len(pd.Categorical(series).categories) == 2
+
+
+def round_num(n, round_to):
+    """
+    Return a string representing a number with `round_to` significant figures.
+
+    Parameters
+    ----------
+    n : float
+        number to round
+    round_to : int
+        number of significant figures
+    """
+    sig_figs = _format_sig_figs(n, round_to)
+    return f"{n:.{sig_figs}g}"
+
+
+def _format_sig_figs(value, default=None):
+    """Get a default number of significant figures.
+
+    Gives the integer part or `default`, whichever is bigger.
+
+    Examples
+    --------
+    0.1234 --> 0.12
+    1.234  --> 1.2
+    12.34  --> 12
+    123.4  --> 123
+    """
+    if default is None:
+        default = 2
+    if value == 0:
+        return 1
+    return max(int(np.log10(np.abs(value))) + 1, default)