Format changes

Author: semihakbayrak

src/causal_validation/config.py

@@ -67,14 +67,10 @@ class Config:
     def __post_init__(self):
         self.rng = np.random.RandomState(self.seed)
         if self.covariate_means is not None:
-            assert self.covariate_means.shape == (
-                self.n_covariates,
-            )
+            assert self.covariate_means.shape == (self.n_covariates,)
 
         if self.covariate_stds is not None:
-            assert self.covariate_stds.shape == (
-                self.n_covariates,
-            )
+            assert self.covariate_stds.shape == (self.n_covariates,)
 
         if (self.n_covariates is not None) & (self.covariate_means is None):
             self.covariate_means = self.rng.normal(
@@ -94,7 +90,9 @@ def __post_init__(self):
             )
 
         n_units = self.treatment_assignments.shape[1]
-        treated_units = [i for i in range(n_units) if any(self.treatment_assignments[:, i] != 0)]
+        treated_units = [
+            i for i in range(n_units) if any(self.treatment_assignments[:, i] != 0)
+        ]
         n_treated_units = len(treated_units)
         n_control_units = n_units - n_treated_units
 
@@ -108,9 +106,4 @@ def __post_init__(self):
                 ]
             else:
                 assert len(self.weights) == n_treated_units
-                assert all(
-                    [
-                        len(w) == n_control_units
-                        for w in self.weights
-                    ]
-                )
+                assert all([len(w) == n_control_units for w in self.weights])

src/causal_validation/data.py

@@ -266,7 +266,7 @@ def _get_columns(self) -> tp.List[str]:
 
     def _get_index(self, start_date: dt.date) -> DatetimeIndex:
         return pd.date_range(start=start_date, freq="D", periods=self.n_timepoints)
-    
+
     def inflate(self, inflation_vals: Float[np.ndarray, "T N"]) -> Dataset:
         """
         Inflate the outputs Y by inflation_vals that are multiplicative factors

src/causal_validation/effects.py

@@ -45,11 +45,12 @@ class StaticEffect(AbstractEffect, _StaticEffect):
         effect (float): Rate effect to be applied, i.e., 0.3 = 30% lift.
         name (str): Name for the effect. 'Static Effect' by default.
     """
+
     effect: float
     name: str = "Static Effect"
 
     def get_effect(self, data: Dataset, **kwargs) -> Float[np.ndarray, "T N"]:
-        return np.ones(data.D.shape) + data.D*self.effect
+        return np.ones(data.D.shape) + data.D * self.effect
 
 
 @dataclass
@@ -64,6 +65,7 @@ class RandomEffect(AbstractEffect, _RandomEffect):
         stddev_effect (float): Rate effect std. dev. to be applied.
         name (str): Name for the effect. 'Random Effect' by default.
     """
+
     mean_effect: float
     stddev_effect: float
     name: str = "Random Effect"
@@ -72,11 +74,9 @@ def get_effect(
         self, data: Dataset, key: np.random.RandomState
     ) -> Float[np.ndarray, "T N"]:
         effect_sample = key.normal(
-            loc=self.mean_effect,
-            scale=self.stddev_effect,
-            size=data.D.shape
+            loc=self.mean_effect, scale=self.stddev_effect, size=data.D.shape
         )
-        return np.ones(data.D.shape) + data.D*effect_sample
+        return np.ones(data.D.shape) + data.D * effect_sample
 
 
 # Placeholder for now.

src/causal_validation/estimator/result.py

@@ -11,4 +11,4 @@ class Result:
     effect: Effect
     counterfactual: Float[NPArray, "N 1"]
     synthetic: Float[NPArray, "N 1"]
-    observed: Float[NPArray, "N 1"]
+    observed: Float[NPArray, "N 1"]

src/causal_validation/estimator/utils.py

@@ -3,37 +3,41 @@
 
 from azcausal.core.error import Error
 from azcausal.core.estimator import Estimator
-from azcausal.core.panel import Panel, CausalPanel
+from azcausal.core.panel import (
+    CausalPanel,
+    Panel,
+)
 from azcausal.core.result import Result as _Result
+from azcausal.util import to_panels
 from jaxtyping import Float
 import pandas as pd
 
 from causal_validation.data import Dataset
 from causal_validation.estimator import Result
 from causal_validation.types import NPArray
-from azcausal.util import to_panels
+
 
 def to_azcausal(dataset: Dataset) -> Panel:
     if dataset.n_treated_units != 1:
         raise ValueError("Only one treated unit is supported.")
     time_index = dataset.full_index
     unit_cols = dataset._get_columns()
-    
+
     data = []
     for time_idx in range(dataset.n_timepoints):
         for unit_idx, unit in enumerate(unit_cols):
-            data.append({
-                'variable': unit,
-                'time': time_index[time_idx],
-                'value': dataset.Y[time_idx, unit_idx],
-                'treated': int(dataset.D[time_idx, unit_idx])
-            })
-    
-    df_data =  pd.DataFrame(data)
+            data.append(
+                {
+                    "variable": unit,
+                    "time": time_index[time_idx],
+                    "value": dataset.Y[time_idx, unit_idx],
+                    "treated": int(dataset.D[time_idx, unit_idx]),
+                }
+            )
+
+    df_data = pd.DataFrame(data)
     panels = to_panels(df_data, "time", "variable", ["value", "treated"])
-    ctypes = dict(
-        outcome="value", time="time", unit="variable", intervention="treated"
-    )
+    ctypes = dict(outcome="value", time="time", unit="variable", intervention="treated")
     panel = CausalPanel(panels).setup(**ctypes)
     return panel
 

src/causal_validation/plotters.py

@@ -32,23 +32,39 @@ def plot(
     Y_control = data.control_unit_outputs
     Y_treated = data.treated_unit_outputs
     idx = data.full_index
-    
+
     if ax is None:
         _, ax = plt.subplots(figsize=(6, 3), tight_layout=True)
-    
+
     ax.plot(idx, Y_control, color=cols[0], label="Control", alpha=0.5)
-    
+
     for i, unit_idx in enumerate(data.treated_unit_indices):
-        unit_color = cols[1] if len(data.treated_unit_indices) == 1 else cols[1 + i % (len(cols) - 2)]
-        unit_label = "Treated" if len(data.treated_unit_indices) == 1 else f"Treated {unit_idx}"
+        unit_color = (
+            cols[1]
+            if len(data.treated_unit_indices) == 1
+            else cols[1 + i % (len(cols) - 2)]
+        )
+        unit_label = (
+            "Treated" if len(data.treated_unit_indices) == 1 else f"Treated {unit_idx}"
+        )
         ax.plot(idx, Y_treated[:, i], color=unit_color, label=unit_label)
-        
+
         treatment_date = data.treatment_date(unit_idx)
         if treatment_date is not None:
             line_color = cols[2] if len(data.treated_unit_indices) == 1 else unit_color
-            line_label = "Intervention" if len(data.treated_unit_indices) == 1 else f"Intervention {unit_idx}"
-            ax.axvline(x=treatment_date, color=line_color, label=line_label, linestyle="--", alpha=0.7)
-    
+            line_label = (
+                "Intervention"
+                if len(data.treated_unit_indices) == 1
+                else f"Intervention {unit_idx}"
+            )
+            ax.axvline(
+                x=treatment_date,
+                color=line_color,
+                label=line_label,
+                linestyle="--",
+                alpha=0.7,
+            )
+
     ax.xaxis.set_major_formatter(
         mdates.ConciseDateFormatter(ax.xaxis.get_major_locator())
     )

src/causal_validation/simulate.py

@@ -20,11 +20,12 @@ def simulate(config: Config, key: tp.Optional[np.random.RandomState] = None) ->
         )
     return base_data
 
+
 def _simulate_with_independent_treated_units(
     config: Config, key: np.random.RandomState
 ) -> Dataset:
     n_timepoints, n_units = config.treatment_assignments.shape
-    
+
     Y = key.normal(
         loc=config.global_mean, scale=config.global_scale, size=(n_timepoints, n_units)
     )
@@ -49,11 +50,12 @@ def _simulate_with_independent_treated_units(
 
     return data
 
+
 def _simulate_with_control_weighted_treated_units(
     config: Config, key: np.random.RandomState
 ) -> Dataset:
     n_timepoints, n_units = config.treatment_assignments.shape
-    
+
     Y = np.zeros((n_timepoints, n_units))
     if config.n_covariates is not None:
         X = np.zeros((n_timepoints, n_units, config.n_covariates))
@@ -71,9 +73,9 @@ def _simulate_with_control_weighted_treated_units(
     treated_unit_indices = data_void.treated_unit_indices
 
     Y_control = key.normal(
-        loc=config.global_mean, 
-        scale=config.global_scale, 
-        size=(n_timepoints, n_control_units)
+        loc=config.global_mean,
+        scale=config.global_scale,
+        size=(n_timepoints, n_control_units),
     )
 
     if config.n_covariates is not None:
@@ -100,4 +102,4 @@ def _simulate_with_control_weighted_treated_units(
         _start_date=config.start_date,
     )
 
-    return data
+    return data

src/causal_validation/testing.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass
-from jaxtyping import Float
 import typing as tp
+
+from jaxtyping import Float
 import numpy as np
 
 from causal_validation.config import Config
 from causal_validation.data import Dataset
 from causal_validation.simulate import simulate
-
 from causal_validation.types import (
     Number,
     TreatedSimulationTypes,
@@ -26,7 +26,7 @@ class TestConstants:
 
     def __post_init__(self):
         if self.TREATMENT_ASSIGNMENTS is None:
-            D = np.zeros((10,5))
+            D = np.zeros((10, 5))
             D[6:, 2] = 1
             D[8:, 3] = 1
             self.TREATMENT_ASSIGNMENTS = D

src/causal_validation/transforms/base.py

@@ -64,13 +64,7 @@ def apply_values(
     ) -> Dataset:
         Y = deepcopy(data.Y)
         Y = Y + transform_vals
-        return Dataset(
-            Y,
-            data.D,
-            data.X,
-            data._start_date,
-            data._name
-        )
+        return Dataset(Y, data.D, data.X, data._start_date, data._name)
 
 
 @dataclass(kw_only=True)
@@ -82,13 +76,7 @@ def apply_values(
     ) -> Dataset:
         Y = deepcopy(data.Y)
         Y = Y * transform_vals
-        return Dataset(
-            Y,
-            data.D,
-            data.X,
-            data._start_date,
-            data._name
-        )
+        return Dataset(Y, data.D, data.X, data._start_date, data._name)
 
 
 @dataclass(kw_only=True)
@@ -100,10 +88,4 @@ def apply_values(
     ) -> Dataset:
         X = deepcopy(data.X)
         X = X + transform_vals
-        return Dataset(
-            data.Y,
-            data.D,
-            X,
-            data._start_date,
-            data._name
-        )
+        return Dataset(data.Y, data.D, X, data._start_date, data._name)

src/causal_validation/transforms/noise.py

@@ -37,7 +37,7 @@ def get_values(self, data: Dataset) -> Float[np.ndarray, "T N"]:
         noise_treatment = self.noise_dist.get_value(
             n_units=data.n_treated_units, n_timepoints=data.n_timepoints
         )
-        noise[:,data.treated_unit_indices] = noise_treatment
+        noise[:, data.treated_unit_indices] = noise_treatment
         return noise