Add rich formating and multi-model support

examples/basic.pct.py

@@ -51,7 +51,7 @@
 scales = [0.1, 0.5]
 
 fig, axes = plt.subplots(ncols=2, nrows=2, figsize=(10, 6), tight_layout=True)
-for (m, s), ax in zip(product(means, scales), axes.ravel()):
+for (m, s), ax in zip(product(means, scales), axes.ravel(), strict=False):
     cfg = Config(
         n_control_units=10,
         n_pre_intervention_timepoints=60,

examples/placebo_test.pct.py

@@ -44,10 +44,6 @@
 from causal_validation.effects import StaticEffect
 from causal_validation.models import AZCausalWrapper
 from causal_validation.plotters import plot
-from causal_validation.transforms import (
-    Periodic,
-    Trend,
-)
 from causal_validation.validation.placebo import PlaceboTest
 
 # %% [markdown]
@@ -99,3 +95,19 @@
 # %%
 result = PlaceboTest(model, data).execute()
 result.summary()
+
+# %% [markdown]
+# ## Model Comparison
+#
+# We can also use the results of a placebo test to compare two or more models. Using
+# `causal-validation`, this is as simple as supplying a series of models to the placebo
+# test and comparing their outputs. To demonstrate this, we will compare the previously
+# used synthetic difference-in-differences model with regular difference-in-differences.
+
+# %%
+did_model = AZCausalWrapper(model=DID())
+PlaceboTest([model, did_model], data).execute().summary()
+
+# %%
+
+# %%

pyproject.toml

@@ -31,7 +31,8 @@ dependencies = [
   "matplotlib",
   "numpy",
   "pandas",
-  "pandera"
+  "pandera",
+  "rich"
 ]
 
 [tool.hatch.build]

src/causal_validation/__about__.py

@@ -1,3 +1,3 @@
-__version__ = "0.0.4"
+__version__ = "0.0.5"
 
 __all__ = ["__version__"]

src/causal_validation/models.py

@@ -13,6 +13,9 @@ class AZCausalWrapper:
     model: Estimator
     error_estimator: tp.Optional[Error] = None
 
+    def __post_init__(self):
+        self._model_name = self.model.__class__.__name__
+
     def __call__(self, data: Dataset, **kwargs) -> Result:
         panel = data.to_azcausal()
         result = self.model.fit(panel, **kwargs)

src/causal_validation/validation/placebo.py

@@ -9,6 +9,8 @@
     Column,
     DataFrameSchema,
 )
+from rich import box
+from rich.table import Table
 from scipy.stats import ttest_1samp
 from tqdm import trange
 
@@ -29,37 +31,67 @@
 
 @dataclass
 class PlaceboTestResult:
-    effects: tp.List[Effect]
+    effects: tp.Dict[str, tp.List[Effect]]
 
-    def summary(self) -> pd.DataFrame:
-        _effects = [effect.value for effect in self.effects]
+    def _model_to_df(self, model_name: str, effects: tp.List[Effect]) -> pd.DataFrame:
+        _effects = [effect.value for effect in effects]
         _n_effects = len(_effects)
         expected_effect = np.mean(_effects)
         stddev_effect = np.std(_effects)
         std_error = stddev_effect / np.sqrt(_n_effects)
         p_value = ttest_1samp(_effects, 0, alternative="two-sided").pvalue
         result = {
+            "Model": model_name,
             "Effect": expected_effect,
             "Standard Deviation": stddev_effect,
             "Standard Error": std_error,
             "p-value": p_value,
         }
         result_df = pd.DataFrame([result])
-        PlaceboSchema.validate(result_df)
         return result_df
 
+    def to_df(self) -> pd.DataFrame:
+        df = pd.concat(
+            [
+                self._model_to_df(model, effects)
+                for model, effects in self.effects.items()
+            ]
+        )
+        PlaceboSchema.validate(df)
+        return df
+
+    def summary(self) -> Table:
+        table = Table(show_header=True, box=box.MARKDOWN)
+        df = self.to_df()
+
+        for column in df.columns:
+            table.add_column(str(column), style="magenta")
+
+        for _, value_list in enumerate(df.values.tolist()):
+            row = [str(x) for x in value_list]
+            table.add_row(*row)
+
+        return table
+
 
 @dataclass
 class PlaceboTest:
-    model: AZCausalWrapper
+    models: tp.Union[AZCausalWrapper, tp.List[AZCausalWrapper]]
     dataset: Dataset
 
+    def __post_init__(self):
+        if isinstance(self.models, AZCausalWrapper):
+            self.models: tp.List[AZCausalWrapper] = [self.models]
+
     def execute(self) -> PlaceboTestResult:
         n_control_units = self.dataset.n_units
-        results = []
-        for i in trange(n_control_units):
-            placebo_data = self.dataset.to_placebo_data(i)
-            result = self.model(placebo_data)
-            result = result.effect.percentage()
-            results.append(result)
+        results = {}
+        for model in self.models:
+            model_result = []
+            for i in trange(n_control_units):
+                placebo_data = self.dataset.to_placebo_data(i)
+                result = model(placebo_data)
+                result = result.effect.percentage()
+                model_result.append(result)
+            results[model._model_name] = model_result
         return PlaceboTestResult(effects=results)

tests/test_causal_validation/test_validation/test_placebo.py

@@ -10,6 +10,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from rich.table import Table
 
 from causal_validation.models import AZCausalWrapper
 from causal_validation.testing import (
@@ -54,11 +55,37 @@ def test_placebo_test(
 
     # Check that the structure of result
     assert isinstance(result, PlaceboTestResult)
-    assert len(result.effects) == n_control
+    for _, v in result.effects.items():
+        assert len(v) == n_control
 
     # Check the results are close to the true effect
-    summary = result.summary()
+    summary = result.to_df()
     PlaceboSchema.validate(summary)
     assert isinstance(summary, pd.DataFrame)
-    assert summary.shape == (1, 4)
+    assert summary.shape == (1, 5)
     assert summary["Effect"].iloc[0] == pytest.approx(0.0, abs=0.1)
+
+    rich_summary = result.summary()
+    assert isinstance(rich_summary, Table)
+    n_rows = result.summary().row_count
+    assert n_rows == summary.shape[0]
+
+
+@pytest.mark.parametrize("n_control", [9, 10])
+def test_multiple_models(n_control: int):
+    constants = TestConstants(N_CONTROL=n_control, GLOBAL_SCALE=0.001)
+    data = simulate_data(global_mean=20.0, seed=123, constants=constants)
+    trend_term = Trend(degree=1, coefficient=0.1)
+    data = trend_term(data)
+
+    model1 = AZCausalWrapper(DID())
+    model2 = AZCausalWrapper(SDID())
+    result = PlaceboTest([model1, model2], data).execute()
+
+    result_df = result.to_df()
+    result_rich = result.summary()
+    assert result_df.shape == (2, 5)
+    assert result_df.shape[0] == result_rich.row_count
+    assert result_df["Model"].tolist() == ["DID", "SDID"]
+    for _, v in result.effects.items():
+        assert len(v) == n_control