diff --git a/causalpy/pymc_experiments.py b/causalpy/pymc_experiments.py index d1f2592c..9f385db2 100644 --- a/causalpy/pymc_experiments.py +++ b/causalpy/pymc_experiments.py @@ -59,10 +59,13 @@ def idata(self): return self.model.idata - def print_coefficients(self) -> None: + def print_coefficients(self, round_to=None) -> None: """ Prints the model coefficients + :param round_to: + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. + Example -------- >>> import causalpy as cp @@ -80,13 +83,13 @@ def print_coefficients(self) -> None: ... "progressbar": False ... }), ... ) - >>> result.print_coefficients() # doctest: +NUMBER + >>> result.print_coefficients(round_to=1) # doctest: +NUMBER Model coefficients: - Intercept 1.0, 94% HDI [1.0, 1.1] - post_treatment[T.True] 0.9, 94% HDI [0.9, 1.0] - group 0.1, 94% HDI [0.0, 0.2] + Intercept 1, 94% HDI [1, 1] + post_treatment[T.True] 1, 94% HDI [0.9, 1] + group 0.2, 94% HDI [0.09, 0.2] group:post_treatment[T.True] 0.5, 94% HDI [0.4, 0.6] - sigma 0.0, 94% HDI [0.0, 0.1] + sigma 0.08, 94% HDI [0.07, 0.1] """ print("Model coefficients:") coeffs = az.extract(self.idata.posterior, var_names="beta") @@ -95,13 +98,13 @@ def print_coefficients(self) -> None: for name in self.labels: coeff_samples = coeffs.sel(coeffs=name) print( - f"{name: <30}{coeff_samples.mean().data:.2f}, 94% HDI [{coeff_samples.quantile(0.03).data:.2f}, {coeff_samples.quantile(1-0.03).data:.2f}]" # noqa: E501 + f"{name: <30}{round_num(coeff_samples.mean().data, round_to)}, 94% HDI [{round_num(coeff_samples.quantile(0.03).data, round_to)}, {round_num(coeff_samples.quantile(1-0.03).data, round_to)}]" # noqa: E501 ) # add coeff for measurement std coeff_samples = az.extract(self.model.idata.posterior, var_names="sigma") name = "sigma" print( - f"{name: <30}{coeff_samples.mean().data:.2f}, 94% HDI [{coeff_samples.quantile(0.03).data:.2f}, {coeff_samples.quantile(1-0.03).data:.2f}]" # noqa: E501 + f"{name: <30}{round_num(coeff_samples.mean().data, round_to)}, 94% HDI [{round_num(coeff_samples.quantile(0.03).data, round_to)}, {round_num(coeff_samples.quantile(1-0.03).data, round_to)}]" # noqa: E501 ) @@ -138,7 +141,7 @@ class PrePostFit(ExperimentalDesign): ... } ... ), ... ) - >>> result.summary() # doctest: +NUMBER + >>> result.summary(round_to=1) # doctest: +NUMBER ==================================Pre-Post Fit================================== Formula: actual ~ 0 + a + g Model coefficients: @@ -231,7 +234,7 @@ def plot(self, counterfactual_label="Counterfactual", round_to=None, **kwargs): Plot the results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8)) @@ -331,15 +334,18 @@ def plot(self, counterfactual_label="Counterfactual", round_to=None, **kwargs): return fig, ax - def summary(self) -> None: + def summary(self, round_to=None) -> None: """ Print text output summarising the results + + :param round_to: + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ print(f"{self.expt_type:=^80}") print(f"Formula: {self.formula}") # TODO: extra experiment specific outputs here - self.print_coefficients() + self.print_coefficients(round_to) class InterruptedTimeSeries(PrePostFit): @@ -420,7 +426,7 @@ def plot(self, plot_predictors=False, **kwargs): """Plot the results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = super().plot(counterfactual_label="Synthetic control", **kwargs) if plot_predictors: @@ -589,7 +595,7 @@ def plot(self, round_to=None): """Plot the results. :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = plt.subplots() @@ -728,17 +734,19 @@ def _causal_impact_summary_stat(self, round_to=None) -> str: causal_impact = f"{round_num(self.causal_impact.mean(), round_to)}, " return f"Causal impact = {causal_impact + ci}" - def summary(self) -> None: + def summary(self, round_to=None) -> None: """ - Print text output summarising the results + Print text output summarising the results. + + :param round_to: + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ print(f"{self.expt_type:=^80}") print(f"Formula: {self.formula}") print("\nResults:") - # TODO: extra experiment specific outputs here - print(self._causal_impact_summary_stat()) - self.print_coefficients() + print(round_num(self._causal_impact_summary_stat(), round_to)) + self.print_coefficients(round_to) class RegressionDiscontinuity(ExperimentalDesign): @@ -894,7 +902,7 @@ def plot(self, round_to=None): Plot the results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = plt.subplots() # Plot raw data @@ -943,9 +951,12 @@ def plot(self, round_to=None): ) return fig, ax - def summary(self) -> None: + def summary(self, round_to: None) -> None: """ Print text output summarising the results + + :param round_to: + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ print(f"{self.expt_type:=^80}") @@ -954,9 +965,9 @@ def summary(self) -> None: print(f"Threshold on running variable: {self.treatment_threshold}") print("\nResults:") print( - f"Discontinuity at threshold = {self.discontinuity_at_threshold.mean():.2f}" + f"Discontinuity at threshold = {round_num(self.discontinuity_at_threshold.mean(), round_to)}" ) - self.print_coefficients() + self.print_coefficients(round_to) class RegressionKink(ExperimentalDesign): @@ -1111,7 +1122,7 @@ def plot(self, round_to=None): Plot the results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = plt.subplots() # Plot raw data @@ -1160,9 +1171,12 @@ def plot(self, round_to=None): ) return fig, ax - def summary(self) -> None: + def summary(self, round_to=None) -> None: """ Print text output summarising the results + + :param round_to: + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ print( @@ -1173,10 +1187,10 @@ def summary(self) -> None: Kink point on running variable: {self.kink_point} Results: - Change in slope at kink point = {self.gradient_change.mean():.2f} + Change in slope at kink point = {round_num(self.gradient_change.mean(), round_to)} """ ) - self.print_coefficients() + self.print_coefficients(round_to) class PrePostNEGD(ExperimentalDesign): @@ -1213,17 +1227,17 @@ class PrePostNEGD(ExperimentalDesign): ... } ... ) ... ) - >>> result.summary() # doctest: +NUMBER + >>> result.summary(round_to=1) # doctest: +NUMBER ==================Pretest/posttest Nonequivalent Group Design=================== Formula: post ~ 1 + C(group) + pre Results: - Causal impact = 1.8, $CI_{94%}$[1.7, 2.1] + Causal impact = 2, $CI_{94%}$[2, 2] Model coefficients: - Intercept -0.4, 94% HDI [-1.1, 0.2] - C(group)[T.1] 1.8, 94% HDI [1.6, 2.0] - pre 1.0, 94% HDI [0.9, 1.1] - sigma 0.5, 94% HDI [0.4, 0.5] + Intercept -0.5, 94% HDI [-1, 0.2] + C(group)[T.1] 2, 94% HDI [2, 2] + pre 1, 94% HDI [1, 1] + sigma 0.5, 94% HDI [0.5, 0.6] """ def __init__( @@ -1304,7 +1318,7 @@ def plot(self, round_to=None): """Plot the results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = plt.subplots( 2, 1, figsize=(7, 9), gridspec_kw={"height_ratios": [3, 1]} @@ -1362,12 +1376,15 @@ def _causal_impact_summary_stat(self, round_to) -> str: r"$CI_{94%}$" + f"[{round_num(percentiles[0], round_to)}, {round_num(percentiles[1], round_to)}]" ) - causal_impact = f"{self.causal_impact.mean():.2f}, " + causal_impact = f"{round_num(self.causal_impact.mean(), round_to)}, " return f"Causal impact = {causal_impact + ci}" def summary(self, round_to=None) -> None: """ Print text output summarising the results + + :param round_to: + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ print(f"{self.expt_type:=^80}") @@ -1375,7 +1392,7 @@ def summary(self, round_to=None) -> None: print("\nResults:") # TODO: extra experiment specific outputs here print(self._causal_impact_summary_stat(round_to)) - self.print_coefficients() + self.print_coefficients(round_to) def _get_treatment_effect_coeff(self) -> str: """Find the beta regression coefficient corresponding to the @@ -1452,7 +1469,6 @@ class InstrumentalVariable(ExperimentalDesign): ... formula=formula, ... model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs), ... ) - """ def __init__( diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py index 1b97b029..5c51d489 100644 --- a/causalpy/skl_experiments.py +++ b/causalpy/skl_experiments.py @@ -119,7 +119,7 @@ def plot(self, counterfactual_label="Counterfactual", round_to=None, **kwargs): """Plot experiment results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8)) @@ -270,7 +270,7 @@ def plot(self, plot_predictors=False, round_to=None, **kwargs): """Plot the results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = super().plot( counterfactual_label="Synthetic control", round_to=round_to, **kwargs @@ -415,7 +415,7 @@ def plot(self, round_to=None): """Plot results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = plt.subplots() @@ -629,7 +629,7 @@ def plot(self, round_to=None): """Plot results :param round_to: - Number of decimals used to round results. Defaults to 2. Use "none" to return raw numbers. + Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers. """ fig, ax = plt.subplots() # Plot raw data