Linting and documentation fix

Author: semihakbayrak

docs/examples/azcausal.ipynb

@@ -17,15 +17,16 @@ trend varies across each of the 10 control units.
 from causal_validation import Config, simulate
 from causal_validation.effects import StaticEffect
 from causal_validation.plotters import plot
-from causal_validation.transforms import Trend, Periodic
+from causal_validation.transforms import Trend
 from causal_validation.transforms.parameter import UnitVaryingParameter
+import numpy as np
 from scipy.stats import norm
 
-cfg = Config(
-    n_control_units=10,
-    n_pre_intervention_timepoints=60,
-    n_post_intervention_timepoints=30,
-)
+# Treatment assignment matrix
+D = np.zeros((90, 11))  # 90 time points, 11 units
+D[60:, -1] = 1  # Last unit treated after 60 time points
+
+cfg = Config(treatment_assignments=D)
 
 # Simulate the base observation
 base_data = simulate(cfg)
@@ -38,6 +39,8 @@ trended_data = trend_component(base_data)
 # Simulate a 5% lift in the treated unit's post-intervention data
 effect = StaticEffect(0.05)
 inflated_data = effect(trended_data)
+
+plot(inflated_data)
 ```
 
 ![Gaussian process posterior.](static/imgs/readme_fig.png)
@@ -50,6 +53,7 @@ combination with AZCausal by the following.
 
 ```python
 from azcausal.estimators.panel.sdid import SDID
+from causal_validation.estimator.utils import AZCausalWrapper
 from causal_validation.validation.placebo import PlaceboTest
 
 model = AZCausalWrapper(model=SDID())

docs/examples/basic.ipynb

@@ -1,6 +1,5 @@
 from dataclasses import (
     dataclass,
-    field,
 )
 import datetime as dt
 import typing as tp

docs/examples/placebo_test.ipynb

@@ -23,8 +23,8 @@
 class Noise(AdditiveOutputTransform):
     """
     Transform the treated units by adding TimeAndUnitVaryingParameter noise terms
-    sampled from a specified sampling distribution. By default, the sampling distribution
-    is Normal with 0 loc and 0.1 scale.
+    sampled from a specified sampling distribution. By default, the sampling
+    distribution is Normal with 0 loc and 0.1 scale.
     """
 
     noise_dist: TimeAndUnitVaryingParameter = field(

docs/index.md

@@ -1,24 +1,18 @@
-from copy import deepcopy
-import datetime as dt
 import string
 import typing as tp
 
-from azcausal.estimators.panel.did import DID
 from hypothesis import (
     given,
     settings,
     strategies as st,
 )
 import numpy as np
 import pandas as pd
-from pandas.core.indexes.datetimes import DatetimeIndex
-import pytest
 
 from causal_validation.data import (
     Dataset,
     DatasetContainer,
 )
-from causal_validation.types import InterventionTypes
 
 
 @given(
@@ -58,8 +52,9 @@ def test_dataset(T: int, N: int, K: int, seed: int):
         == (T * np.ones(N) - data1.n_post_intervention).tolist()
     )
 
-    assert data1.n_treated_units == 2
-    assert data1.n_control_units == N - 2
+    N_TREATED = 2
+    assert data1.n_treated_units == N_TREATED
+    assert data1.n_control_units == N - N_TREATED
     assert data1.treated_unit_indices == [2, 3]
     assert data1.control_unit_indices == [0, 1] + list(range(4, N))
 
@@ -92,7 +87,8 @@ def test_dataset_to_df(T: int, N: int, K: int, seed: int):
     df1 = data1.to_df()
     assert isinstance(df1, pd.DataFrame)
     assert df1.shape == (T, N * (K + 2))
-    assert df1.columns.nlevels == 2
+    EXPECTED_COLUMN_LEVELS = 2
+    assert df1.columns.nlevels == EXPECTED_COLUMN_LEVELS
     assert df1.columns[0] == ("U0", "Y")
     assert df1.columns[1] == ("U0", "D")
     assert df1.columns[2] == ("U0", "X0")
@@ -114,7 +110,7 @@ def test_dataset_to_df(T: int, N: int, K: int, seed: int):
     df2 = data2.to_df()
     assert isinstance(df2, pd.DataFrame)
     assert df2.shape == (T, N * 2)
-    assert df2.columns.nlevels == 2
+    assert df2.columns.nlevels == EXPECTED_COLUMN_LEVELS
     assert df2.columns[0] == ("U0", "Y")
     assert df2.columns[1] == ("U0", "D")
     assert df2.columns[2] == ("U1", "Y")
@@ -198,13 +194,11 @@ def test_dataset_container(seeds: tp.List[int], to_name: bool, T: int, N: int):
     K=st.integers(min_value=1, max_value=10),
     seed=st.integers(min_value=1, max_value=100),
     use_bernoulli=st.booleans(),
-    include_X=st.booleans(),
 )
 @settings(max_examples=10)
-def test_inflate(
-    T: int, N: int, K: int, seed: int, use_bernoulli: bool, include_X: bool
-):
+def test_inflate(T: int, N: int, K: int, seed: int, use_bernoulli: bool):
     rng = np.random.RandomState(seed)
+    include_X = bool(rng.binomial(1, 0.5))
     Y = rng.randn(T, N)
     D = rng.binomial(1, 0.3, (T, N)) if use_bernoulli else np.abs(rng.randn(T, N))
     X = rng.randn(T, N, K) if include_X else None

src/causal_validation/config.py

@@ -20,20 +20,12 @@
         min_value=-0.5, max_value=0.5, allow_nan=False, allow_infinity=False
     ),
     seed=st.integers(min_value=1, max_value=100),
-    use_bernoulli=st.booleans(),
-    include_X=st.booleans(),
 )
 @settings(max_examples=10)
-def test_static_effect(
-    T: int,
-    N: int,
-    K: int,
-    effect: float,
-    seed: int,
-    use_bernoulli: bool,
-    include_X: bool,
-):
+def test_static_effect(T: int, N: int, K: int, effect: float, seed: int):
     rng = np.random.RandomState(seed)
+    use_bernoulli = bool(rng.binomial(1, 0.5))
+    include_X = bool(rng.binomial(1, 0.5))
     Y = rng.randn(T, N)
     D = rng.binomial(1, 0.3, (T, N)) if use_bernoulli else np.abs(rng.randn(T, N))
     X = rng.randn(T, N, K) if include_X else None
@@ -53,7 +45,6 @@ def test_static_effect(
 
 
 @given(
-    T=st.integers(min_value=2, max_value=50),
     N=st.integers(min_value=2, max_value=50),
     K=st.integers(min_value=1, max_value=10),
     mean_effect=st.floats(
@@ -63,21 +54,19 @@ def test_static_effect(
         min_value=0.01, max_value=0.2, allow_nan=False, allow_infinity=False
     ),
     seed=st.integers(min_value=1, max_value=100),
-    use_bernoulli=st.booleans(),
-    include_X=st.booleans(),
 )
 @settings(max_examples=10)
 def test_random_effect(
-    T: int,
     N: int,
     K: int,
     mean_effect: float,
     stddev_effect: float,
     seed: int,
-    use_bernoulli: bool,
-    include_X: bool,
 ):
     rng = np.random.RandomState(seed)
+    use_bernoulli = bool(rng.binomial(1, 0.5))
+    include_X = bool(rng.binomial(1, 0.5))
+    T = 50
     Y = rng.randn(T, N)
     D = rng.binomial(1, 0.3, (T, N)) if use_bernoulli else np.abs(rng.randn(T, N))
     X = rng.randn(T, N, K) if include_X else None

src/causal_validation/transforms/noise.py

@@ -163,9 +163,10 @@ def test_simulate_control_weighted(n_units, n_timepoints, seed):
     )
     data = simulate(cfg)
 
+    N_TREATED = 2
     assert data.Y.shape == (n_timepoints, n_units)
-    assert data.n_control_units == n_units - 2
-    assert data.n_treated_units == 2
+    assert data.n_control_units == n_units - N_TREATED
+    assert data.n_treated_units == N_TREATED
     assert np.all(data.Y[:, :-2] @ cfg.weights[0] == data.Y[:, -2])
     assert np.all(data.Y[:, :-2] @ cfg.weights[1] == data.Y[:, -1])
 
@@ -208,9 +209,10 @@ def test_simulate_control_weighted_with_covariates(
     )
     data = simulate(cfg)
 
+    N_TREATED = 2
     assert data.Y.shape == (n_timepoints, n_units)
-    assert data.n_control_units == n_units - 2
-    assert data.n_treated_units == 2
+    assert data.n_control_units == n_units - N_TREATED
+    assert data.n_treated_units == N_TREATED
     assert np.all(data.Y[:, :-2] @ cfg.weights[0] == data.Y[:, -2])
     assert np.all(data.Y[:, :-2] @ cfg.weights[1] == data.Y[:, -1])
     X_treated1 = np.einsum("ijk,j->ik", data.X[:, :-2, :], cfg.weights[0])

tests/test_causal_validation/test_data.py

@@ -161,7 +161,7 @@ def test_varying_parameters():
     data_slots = constants.DATA_SLOTS
     base_data = simulate_data(GLOBAL_MEAN, DEFAULT_SEED, constants=constants)
     base_data_transform = periodic_transform(base_data)
-    for i, slot in enumerate(param_slots):
+    for slot in param_slots:
         setattr(
             periodic_transform,
             slot,