diff --git a/pyproject.toml b/pyproject.toml index accaaef..2cf98ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,7 +171,7 @@ select = [ "TID", "ISC", ] -ignore = ["F722"] +ignore = ["F722", "PLW1641"] [tool.ruff.format] quote-style = "double" diff --git a/src/causal_validation/data.py b/src/causal_validation/data.py index 057acfa..a81f13a 100644 --- a/src/causal_validation/data.py +++ b/src/causal_validation/data.py @@ -21,15 +21,45 @@ @dataclass class Dataset: + """A causal inference dataset containing pre/post intervention observations + and optional associated covariates. + + Attributes: + Xtr: Pre-intervention control unit observations (N x D) + Xte: Post-intervention control unit observations (M x D) + ytr: Pre-intervention treated unit observations (N x 1) + yte: Post-intervention treated unit observations (M x 1) + _start_date: Start date for time indexing + Ptr: Pre-intervention control unit covariates (N x D x F) + Pte: Post-intervention control unit covariates (M x D x F) + Rtr: Pre-intervention treated unit covariates (N x 1 x F) + Rte: Post-intervention treated unit covariates (M x 1 x F) + counterfactual: Optional counterfactual outcomes (M x 1) + synthetic: Optional synthetic control outcomes (M x 1). + This is weighted combination of control units + minimizing a distance-based error w.r.t. the + treated in pre-intervention period. + _name: Optional name identifier for the dataset + """ Xtr: Float[np.ndarray, "N D"] Xte: Float[np.ndarray, "M D"] ytr: Float[np.ndarray, "N 1"] yte: Float[np.ndarray, "M 1"] _start_date: dt.date + Ptr: tp.Optional[Float[np.ndarray, "N D F"]] = None + Pte: tp.Optional[Float[np.ndarray, "M D F"]] = None + Rtr: tp.Optional[Float[np.ndarray, "N 1 F"]] = None + Rte: tp.Optional[Float[np.ndarray, "M 1 F"]] = None counterfactual: tp.Optional[Float[np.ndarray, "M 1"]] = None synthetic: tp.Optional[Float[np.ndarray, "M 1"]] = None _name: str = None + def __post_init__(self): + covariates = [self.Ptr, self.Pte, self.Rtr, self.Rte] + self.has_covariates = all(cov is not None for cov in covariates) + if not self.has_covariates: + assert all(cov is None for cov in covariates) + def to_df( self, index_start: str = dt.date(year=2023, month=1, day=1) ) -> pd.DataFrame: @@ -59,6 +89,13 @@ def n_units(self) -> int: def n_timepoints(self) -> int: return self.n_post_intervention + self.n_pre_intervention + @property + def n_covariates(self) -> int: + if self.has_covariates: + return self.Ptr.shape[2] + else: + return 0 + @property def control_units(self) -> Float[np.ndarray, "{self.n_timepoints} {self.n_units}"]: return np.vstack([self.Xtr, self.Xte]) @@ -67,6 +104,26 @@ def control_units(self) -> Float[np.ndarray, "{self.n_timepoints} {self.n_units} def treated_units(self) -> Float[np.ndarray, "{self.n_timepoints} 1"]: return np.vstack([self.ytr, self.yte]) + @property + def control_covariates( + self, + ) -> tp.Optional[ + Float[np.ndarray, "{self.n_timepoints} {self.n_units} {self.n_covariates}"] + ]: + if self.has_covariates: + return np.vstack([self.Ptr, self.Pte]) + else: + return None + + @property + def treated_covariates( + self, + ) -> tp.Optional[Float[np.ndarray, "{self.n_timepoints} 1 {self.n_covariates}"]]: + if self.has_covariates: + return np.vstack([self.Rtr, self.Rte]) + else: + return None + @property def pre_intervention_obs( self, @@ -79,6 +136,32 @@ def post_intervention_obs( ) -> tp.Tuple[Float[np.ndarray, "M D"], Float[np.ndarray, "M 1"]]: return self.Xte, self.yte + @property + def pre_intervention_covariates( + self, + ) -> tp.Optional[ + tp.Tuple[ + Float[np.ndarray, "N D F"], Float[np.ndarray, "N 1 F"], + ] + ]: + if self.has_covariates: + return self.Ptr, self.Rtr + else: + return None + + @property + def post_intervention_covariates( + self, + ) -> tp.Optional[ + tp.Tuple[ + Float[np.ndarray, "M D F"], Float[np.ndarray, "M 1 F"], + ] + ]: + if self.has_covariates: + return self.Pte, self.Rte + else: + return None + @property def full_index(self) -> DatetimeIndex: return self._get_index(self._start_date) @@ -97,7 +180,12 @@ def get_index(self, period: InterventionTypes) -> DatetimeIndex: return self.full_index def _get_columns(self) -> tp.List[str]: - colnames = ["T"] + [f"C{i}" for i in range(self.n_units)] + if self.has_covariates: + colnames = ["T"] + [f"C{i}" for i in range(self.n_units)] + [ + f"F{i}" for i in range(self.n_covariates) + ] + else: + colnames = ["T"] + [f"C{i}" for i in range(self.n_units)] return colnames def _get_index(self, start_date: dt.date) -> DatetimeIndex: @@ -116,7 +204,10 @@ def inflate(self, inflation_vals: Float[np.ndarray, "M 1"]) -> Dataset: Xtr, ytr = [deepcopy(i) for i in self.pre_intervention_obs] Xte, yte = [deepcopy(i) for i in self.post_intervention_obs] inflated_yte = yte * inflation_vals - return Dataset(Xtr, Xte, ytr, inflated_yte, self._start_date, yte) + return Dataset( + Xtr, Xte, ytr, inflated_yte, self._start_date, + self.Ptr, self.Pte, self.Rtr, self.Rte, yte, self.synthetic, self._name + ) def __eq__(self, other: Dataset) -> bool: ytr = np.allclose(self.ytr, other.ytr) @@ -151,14 +242,21 @@ def _slots(self) -> tp.Dict[str, int]: def drop_unit(self, idx: int) -> Dataset: Xtr = np.delete(self.Xtr, [idx], axis=1) Xte = np.delete(self.Xte, [idx], axis=1) + Ptr = np.delete(self.Ptr, [idx], axis=1) if self.Ptr is not None else None + Pte = np.delete(self.Pte, [idx], axis=1) if self.Pte is not None else None return Dataset( Xtr, Xte, self.ytr, self.yte, self._start_date, + Ptr, + Pte, + self.Rtr, + self.Rte, self.counterfactual, self.synthetic, + self._name, ) def to_placebo_data(self, to_treat_idx: int) -> Dataset: @@ -212,5 +310,7 @@ def reassign_treatment( Xtr = data.Xtr Xte = data.Xte return Dataset( - Xtr, Xte, ytr, yte, data._start_date, data.counterfactual, data.synthetic + Xtr, Xte, ytr, yte, data._start_date, + data.Ptr, data.Pte, data.Rtr, data.Rte, + data.counterfactual, data.synthetic, data._name ) diff --git a/src/causal_validation/validation/placebo.py b/src/causal_validation/validation/placebo.py index b8f7c36..b8143fb 100644 --- a/src/causal_validation/validation/placebo.py +++ b/src/causal_validation/validation/placebo.py @@ -1,7 +1,6 @@ from dataclasses import dataclass import typing as tp -from azcausal.core.effect import Effect import numpy as np import pandas as pd from pandera import ( @@ -11,14 +10,8 @@ ) from rich.progress import ( Progress, - ProgressBar, - track, ) from scipy.stats import ttest_1samp -from tqdm import ( - tqdm, - trange, -) from causal_validation.data import ( Dataset, @@ -108,7 +101,7 @@ def execute(self, verbose: bool = True) -> PlaceboTestResult: "[blue]Datasets", total=n_datasets, visible=verbose ) unit_task = progress.add_task( - f"[green]Control Units", + "[green]Control Units", total=n_control, visible=verbose, ) diff --git a/src/causal_validation/validation/rmspe.py b/src/causal_validation/validation/rmspe.py index 6b541ff..b606722 100644 --- a/src/causal_validation/validation/rmspe.py +++ b/src/causal_validation/validation/rmspe.py @@ -2,18 +2,14 @@ import typing as tp from jaxtyping import Float -import numpy as np import pandas as pd from pandera import ( Check, Column, DataFrameSchema, ) -from rich import box from rich.progress import ( Progress, - ProgressBar, - track, ) from causal_validation.validation.placebo import PlaceboTest @@ -87,7 +83,7 @@ def execute(self, verbose: bool = True) -> RMSPETestResult: "[blue]Datasets", total=n_datasets, visible=verbose ) unit_task = progress.add_task( - f"[green]Treatment and Control Units", + "[green]Treatment and Control Units", total=n_control + 1, visible=verbose, ) diff --git a/tests/test_causal_validation/test_data.py b/tests/test_causal_validation/test_data.py index 07554a4..1b7705c 100644 --- a/tests/test_causal_validation/test_data.py +++ b/tests/test_causal_validation/test_data.py @@ -23,6 +23,7 @@ simulate_data, ) from causal_validation.types import InterventionTypes +import datetime as dt MIN_STRING_LENGTH = 1 MAX_STRING_LENGTH = 20 @@ -199,6 +200,10 @@ def test_drop_unit(n_pre: int, n_post: int, n_control: int): assert reduced_data.ytr.shape == desired_shape_ytr assert reduced_data.yte.shape == desired_shape_yte + assert reduced_data.counterfactual == data.counterfactual + assert reduced_data.synthetic == data.synthetic + assert reduced_data._name == data._name + @pytest.mark.parametrize("n_pre, n_post, n_control", [(60, 30, 10), (60, 30, 20)]) def test_to_placebo(n_pre: int, n_post: int, n_control: int): @@ -287,6 +292,88 @@ def test_naming_setter(name: str, extra_chars: str): assert data.name == new_name +@given( + n_pre=st.integers(min_value=10, max_value=100), + n_post=st.integers(min_value=10, max_value=100), + n_control=st.integers(min_value=2, max_value=20), +) +@settings(max_examples=5) +def test_counterfactual_synthetic_attributes(n_pre: int, n_post: int, n_control: int): + constants = TestConstants( + N_POST_TREATMENT=n_post, + N_PRE_TREATMENT=n_pre, + N_CONTROL=n_control, + ) + data = simulate_data(0.0, DEFAULT_SEED, constants=constants) + + assert data.counterfactual is None + assert data.synthetic is None + + counterfactual_vals = np.random.randn(n_post, 1) + synthetic_vals = np.random.randn(n_post, 1) + + data_with_attrs = Dataset( + data.Xtr, data.Xte, data.ytr, data.yte, data._start_date, + data.Ptr, data.Pte, data.Rtr, data.Rte, + counterfactual_vals, synthetic_vals, "test_dataset" + ) + + np.testing.assert_array_equal(data_with_attrs.counterfactual, counterfactual_vals) + np.testing.assert_array_equal(data_with_attrs.synthetic, synthetic_vals) + assert data_with_attrs.name == "test_dataset" + + +@given( + n_pre=st.integers(min_value=10, max_value=100), + n_post=st.integers(min_value=10, max_value=100), + n_control=st.integers(min_value=2, max_value=20), +) +@settings(max_examples=5) +def test_inflate_method(n_pre: int, n_post: int, n_control: int): + constants = TestConstants( + N_POST_TREATMENT=n_post, + N_PRE_TREATMENT=n_pre, + N_CONTROL=n_control, + ) + data = simulate_data(0.0, DEFAULT_SEED, constants=constants) + + inflation_vals = np.ones((n_post, 1)) * 1.1 + inflated_data = data.inflate(inflation_vals) + + np.testing.assert_array_equal(inflated_data.Xtr, data.Xtr) + np.testing.assert_array_equal(inflated_data.ytr, data.ytr) + np.testing.assert_array_equal(inflated_data.Xte, data.Xte) + + expected_yte = data.yte * inflation_vals + np.testing.assert_array_equal(inflated_data.yte, expected_yte) + + np.testing.assert_array_equal(inflated_data.counterfactual, data.yte) + + +@given( + n_pre=st.integers(min_value=10, max_value=100), + n_post=st.integers(min_value=10, max_value=100), + n_control=st.integers(min_value=2, max_value=20), +) +@settings(max_examples=5) +def test_control_treated_properties(n_pre: int, n_post: int, n_control: int): + constants = TestConstants( + N_POST_TREATMENT=n_post, + N_PRE_TREATMENT=n_pre, + N_CONTROL=n_control, + ) + data = simulate_data(0.0, DEFAULT_SEED, constants=constants) + + control_units = data.control_units + expected_control = np.vstack([data.Xtr, data.Xte]) + np.testing.assert_array_equal(control_units, expected_control) + assert control_units.shape == (n_pre + n_post, n_control) + + treated_units = data.treated_units + expected_treated = np.vstack([data.ytr, data.yte]) + np.testing.assert_array_equal(treated_units, expected_treated) + assert treated_units.shape == (n_pre + n_post, 1) + @given( seeds=st.lists( elements=st.integers(min_value=1, max_value=1000), min_size=1, max_size=10 @@ -322,3 +409,75 @@ def test_dataset_container(seeds: tp.List[int], to_name: bool): else: assert k == f"Dataset {idx}" assert v == datasets[idx] + +@given( + n_pre=st.integers(min_value=10, max_value=100), + n_post=st.integers(min_value=10, max_value=100), + n_control=st.integers(min_value=2, max_value=20), +) +@settings(max_examples=5) +def test_covariate_properties_without_covariates( + n_pre: int, n_post: int, n_control: int +): + constants = TestConstants( + N_POST_TREATMENT=n_post, + N_PRE_TREATMENT=n_pre, + N_CONTROL=n_control, + ) + data = simulate_data(0.0, DEFAULT_SEED, constants=constants) + + assert data.has_covariates is False + assert data.control_covariates is None + assert data.treated_covariates is None + assert data.pre_intervention_covariates is None + assert data.post_intervention_covariates is None + assert data.n_covariates == 0 + + +@given( + n_pre=st.integers(min_value=10, max_value=50), + n_post=st.integers(min_value=10, max_value=50), + n_control=st.integers(min_value=2, max_value=10), + n_covariates=st.integers(min_value=1, max_value=5), + seed=st.integers(min_value=1, max_value=10000), +) +@settings(max_examples=5) +def test_covariate_properties_with_covariates( + n_pre: int, + n_post: int, + n_control: int, + n_covariates: int, + seed: int, +): + rng = np.random.RandomState(seed) + + Xtr = rng.uniform(-10, 10, (n_pre, n_control)) + Xte = rng.uniform(-10, 10, (n_post, n_control)) + ytr = rng.uniform(-10, 10, (n_pre, 1)) + yte = rng.uniform(-10, 10, (n_post, 1)) + Ptr = rng.uniform(-10, 10, (n_pre, n_control, n_covariates)) + Pte = rng.uniform(-10, 10, (n_post, n_control, n_covariates)) + Rtr = rng.uniform(-10, 10, (n_pre, 1, n_covariates)) + Rte = rng.uniform(-10, 10, (n_post, 1, n_covariates)) + + data = Dataset(Xtr, Xte, ytr, yte, dt.date(2023, 1, 1), Ptr, Pte, Rtr, Rte) + + assert data.n_covariates == n_covariates + assert data.has_covariates is True + + control_covariates = data.control_covariates + expected_control_cov = np.vstack([Ptr, Pte]) + np.testing.assert_array_equal(control_covariates, expected_control_cov) + assert control_covariates.shape == (n_pre + n_post, n_control, n_covariates) + + treated_covariates = data.treated_covariates + expected_treated_cov = np.vstack([Rtr, Rte]) + np.testing.assert_array_equal(treated_covariates, expected_treated_cov) + assert treated_covariates.shape == (n_pre + n_post, 1, n_covariates) + + pre_cov = data.pre_intervention_covariates + assert pre_cov == (Ptr, Rtr) + + post_cov = data.post_intervention_covariates + assert post_cov == (Pte, Rte) + diff --git a/tests/test_causal_validation/test_validation/test_placebo.py b/tests/test_causal_validation/test_validation/test_placebo.py index 858c5f3..172de83 100644 --- a/tests/test_causal_validation/test_validation/test_placebo.py +++ b/tests/test_causal_validation/test_validation/test_placebo.py @@ -30,7 +30,7 @@ def test_schema_coerce(): df = PlaceboSchema.example() cols = df.columns for col in cols: - if not col in ["Model", "Dataset"]: + if col not in ["Model", "Dataset"]: df[col] = np.ceil((df[col])) PlaceboSchema.validate(df) diff --git a/tests/test_causal_validation/test_validation/test_rmspe.py b/tests/test_causal_validation/test_validation/test_rmspe.py index 1bc6b37..ead1dfa 100644 --- a/tests/test_causal_validation/test_validation/test_rmspe.py +++ b/tests/test_causal_validation/test_validation/test_rmspe.py @@ -35,7 +35,7 @@ def test_schema_coerce(): df = RMSPESchema.example() cols = df.columns for col in cols: - if not col in ["Model", "Dataset"]: + if col not in ["Model", "Dataset"]: df[col] = np.ceil((df[col])) RMSPESchema.validate(df)