diffindiff 2.0.1__tar.gz → 2.0.3__tar.gz

{diffindiff-2.0.1 → diffindiff-2.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: diffindiff
-Version: 2.0.1
+Version: 2.0.3
 Summary: diffindiff: Python library for convenient Difference-in-Differences Analyses
 Author: Thomas Wieland
 Author-email: geowieland@googlemail.com
@@ -38,17 +38,18 @@ Thomas Wieland [ORCID](https://orcid.org/0000-0001-5168-9846) [EMail](mailto:geo
   - Create predictive counterfactuals
 - **DiD analysis**: 
   - Perfom standard DiD analysis  
-  - Model Extensions:
+  - Model extensions:
     - Staggered adoption
     - Multiple treatments
     - Two-way fixed effects models
     - Group- or individual-specific treatment effects
     - Group- or individual-specific time trends
     - Including covariates
-    - After-treatment period
+    - Including fter-treatment period
     - Triple Difference (DDD)
     - Own counterfactuals
-    - Bonferroni correction
+    - Bonferroni correction for treatment effects
+    - Placebo test
 - **Visualization**:
   - Plot observed and expected time course of treatment and control group
   - Plot expected time course of treatment group and counterfactual
@@ -60,7 +61,6 @@ Thomas Wieland [ORCID](https://orcid.org/0000-0001-5168-9846) [EMail](mailto:geo
   - Test for type of adoption
   - Test whether the panel dataset is balanced
   - Test for parallel trend assumption
-  - Placebo test
 
 
 ## Literature

{diffindiff-2.0.1 → diffindiff-2.0.3}/README.md

@@ -16,17 +16,18 @@ Thomas Wieland [ORCID](https://orcid.org/0000-0001-5168-9846) [EMail](mailto:geo
   - Create predictive counterfactuals
 - **DiD analysis**: 
   - Perfom standard DiD analysis  
-  - Model Extensions:
+  - Model extensions:
     - Staggered adoption
     - Multiple treatments
     - Two-way fixed effects models
     - Group- or individual-specific treatment effects
     - Group- or individual-specific time trends
     - Including covariates
-    - After-treatment period
+    - Including fter-treatment period
     - Triple Difference (DDD)
     - Own counterfactuals
-    - Bonferroni correction
+    - Bonferroni correction for treatment effects
+    - Placebo test
 - **Visualization**:
   - Plot observed and expected time course of treatment and control group
   - Plot expected time course of treatment group and counterfactual
@@ -38,7 +39,6 @@ Thomas Wieland [ORCID](https://orcid.org/0000-0001-5168-9846) [EMail](mailto:geo
   - Test for type of adoption
   - Test whether the panel dataset is balanced
   - Test for parallel trend assumption
-  - Placebo test
 
 
 ## Literature

{diffindiff-2.0.1 → diffindiff-2.0.3}/diffindiff/didanalysis.py

@@ -1,11 +1,13 @@
-#-------------------------------------------------------------------------------
-# Name:        didanalysis (diffindiff)
+#-----------------------------------------------------------------------
+# Name:        didanalysis (diffindiff package)
 # Purpose:     Analysis functions for difference-in-differences analyses
-# Author:      Thomas Wieland (mail: geowieland@googlemail.com, ORCID: 0000-0001-5168-9846)
-# Version:     2.0.1
-# Last update: 2025-04-15 18:43
+# Author:      Thomas Wieland 
+#              ORCID: 0000-0001-5168-9846
+#              mail: geowieland@googlemail.com              
+# Version:     2.0.3
+# Last update: 2025-04-18 10:24
 # Copyright (c) 2025 Thomas Wieland
-#-------------------------------------------------------------------------------
+#-----------------------------------------------------------------------
 
 
 import pandas as pd
@@ -25,7 +27,8 @@ class DiffModel:
         did_modeldata,
         did_modelpredictions,
         did_model_statistics,
-        did_olsmodel
+        did_olsmodel,
+        did_prediction_intervals
         ):
 
         self.data = [
@@ -34,7 +37,8 @@ class DiffModel:
             did_modeldata, 
             did_modelpredictions, 
             did_model_statistics, 
-            did_olsmodel
+            did_olsmodel,
+            did_prediction_intervals
             ]
 
     def treatment_statistics(
@@ -82,7 +86,7 @@ class DiffModel:
         after_treatment_period_start = None
         after_treatment_period_end = None
         after_treatment_period_N = None
-        if len(model_config["after_treatment_col"]) > 0:
+        if len(model_config["after_treatment_col"]) > 0 and after_treatment_col is not None:
             after_treatment_period_start = treatment_period_end+pd.Timedelta(days=1)
             after_treatment_period_start = pd.to_datetime(after_treatment_period_start)
             after_treatment_period_end = pd.to_datetime(study_period_end)
@@ -364,7 +368,7 @@ class DiffModel:
 
             for key, value in covariates_effects.items():
                 covariates_effects_rows.append({
-                    "Covariates": value["Coefficient"],
+                    "": value["Coefficient"],
                     "Estimate": value["Estimate"],
                     "SE": value["SE"],
                     "t": value["t"],
@@ -523,13 +527,15 @@ class DiffModel:
             covariates_effects_df["CI lower"] = covariates_effects_df["CI lower"].map(lambda x: f"{x:,.3f}")
             covariates_effects_df["CI upper"] = covariates_effects_df["CI upper"].map(lambda x: f"{x:,.3f}")
             covariates_effects_df.iloc[:, 0] = covariates_effects_df.iloc[:, 0].apply(lambda x: f"{x:<{max_width_column1}}")
+            print("Covariates") 
             print(covariates_effects_df.to_string(index=False))
-        if not show_covariates:
+        if not show_covariates or no_covariates == 0:
             if no_covariates > 0:
                 print ("Covariates                 YES")
             else:
                 print ("Covariates                 NO")
 
+        print("")
         print("Fixed effects")        
         if model_config["FE_unit"]:
             print (" Units                     YES")
@@ -566,7 +572,7 @@ class DiffModel:
         print(treatment_diagnostics_df_t)
 
         print("-" * total_width)
-        print ("Input data diagnostics")
+        print ("Input data diagnostixx") # TODO ?? AENDERN
         if modeldata_isbalanced:
             print ("Balanced panel data        YES")
         else:
@@ -756,16 +762,21 @@ class DiffModel:
         ols_model = self.data[5]
         return ols_model
     
+    def prediction_intervals(self):
+
+        prediction_intervals = self.data[6]
+        return prediction_intervals
+    
     def placebo(
-            self,
-            treatment: str = None,
-            after_treatment_col: str = None,
-            TG_col: str = None,
-            TT_col: str = None,
-            divide: float = 0.5,
-            resample: float = 1.0,
-            random_state = 71
-            ):
+        self,
+        treatment: str = None,
+        after_treatment_col: str = None,
+        TG_col: str = None,
+        TT_col: str = None,
+        divide: float = 0.5,
+        resample: float = 1.0,
+        random_state = 71
+        ):
 
         model_config = self.data[1]
         model_data = self.data[2]
@@ -796,9 +807,9 @@ class DiffModel:
         TT_col_ = "TT_" + treatment
         TGxTT_ = "Placebo_" + treatment
         if TG_col is None and TG_col_ not in model_config["TG_col"]:
-            raise ValueError("Model object does not include treatment group identification variable for treatment ", treatment)
+            raise ValueError("Cannot find treatment group identification variable for treatment " + treatment + ". Please state TG_col = [treatment_group_dummy].")
         if TT_col is None and TT_col_ not in model_config["TT_col"]:
-            raise ValueError("Model object does not include treatment time variable for treatment ", treatment)
+            raise ValueError("Cannot findt treatment time variable for treatment " + treatment + ". Please state TG_col = [treatment_time_dummy].")
 
         unit_col = model_config["unit_col"]
         time_col = model_config["time_col"]
@@ -1127,19 +1138,20 @@ class DiffModel:
         return model_data_TG_CG    
 
     def plot_counterfactual(
-            self,
-            treatment = None,
-            x_label: str = "Time",
-            y_label: str = "Outcome",
-            y_lim = None,
-            plot_title: str = "Treatment group Counterfactual",
-            lines_col: list = ["blue", "green"],
-            lines_style: list = ["solid", "dashed"],
-            lines_labels: list = ["TG", "TG counterfactual"],
-            plot_legend: bool = True,
-            plot_grid: bool = True,
-            plot_size: list = [12, 6]
-            ):
+        self,
+        treatment: str = None,
+        after_treatment_col: str = None,
+        x_label: str = "Time",
+        y_label: str = "Outcome",
+        y_lim = None,
+        plot_title: str = "Treatment group Counterfactual",
+        lines_col: list = ["blue", "green"],
+        lines_style: list = ["solid", "dashed"],
+        lines_labels: list = ["TG", "TG counterfactual"],
+        plot_legend: bool = True,
+        plot_grid: bool = True,
+        plot_size: list = [12, 6]
+        ):
 
         model_config = self.data[1]
         outcome_col = model_config["outcome_col"]
@@ -1158,7 +1170,10 @@ class DiffModel:
             else:
                 raise ValueError ("Model object has no column for treatment group with respect to ", str(no_treatments), " treatments. Choose one with parameter treatment.")
 
-        model_data_mod = self.counterfactual()
+        model_data_mod = self.counterfactual(
+            treatment = treatment,
+            after_treatment_col = after_treatment_col
+            )
 
         if treatment is not None:
 
@@ -1182,6 +1197,8 @@ class DiffModel:
 
             treatment = treatment_diagnostics[0]["treatment"]
 
+        treatment_group = [str(x) for x in treatment_group]
+
         TG_col = "TG_" + treatment
         
         model_data_mod[TG_col] = 0
@@ -1344,20 +1361,41 @@ def did_analysis(
         intercept = False
         TG_col = []
         print ("NOTE: Quasi-experiment includes more than one treatment. Unit fixed effects are used instead of control group baseline and treatment group deviation.")
-        
-    if ITE:
-        GTE = False    
-    if ITT:
-        GTT = False  
-    
+           
+    if ITE:        
+        FE_unit = True
+        print ("NOTE: Model includes individual treatment effects. Unit fixed effects are included.")
+        if GTE:
+            GTE = False
+            print ("NOTE: Both group and individual treatment effects were stated. Switching to individual treatment effects only.")
+    if ITT:        
+        FE_unit = True        
+        TT_col = []
+        print ("NOTE: Model includes individual time trends. Unit fixed effects are included. Treatment time variable is dropped.")
+        if FE_time:
+            FE_time = False
+            print ("NOTE: Time fixed effects are dropped.")
+        if GTT:
+            GTT = False
+            print ("NOTE: Both group and individual time trends were stated. Switching to individual time trends only.")
+            
     if staggered_adoption:
         FE_unit = True
         FE_time = True        
         print ("NOTE: Quasi-experiment includes one or more staggered treatments. Two-way fixed effects model is used.")
 
-    if FE_unit and FE_time:
+    FE_group = False
+    if group_by is not None and group_by != "":
+        FE_group = True        
+               
+    if FE_unit:
         TG_col = []
+    if FE_time:
         TT_col = []
+    if FE_group:    
+        TG_col = []        
+        intercept = False
+        print ("NOTE: Quasi-experiment includes group fixed effects. Control group baseline and treatment group deviation are dropped.")   
     
     if after_treatment_col is not None or (isinstance (after_treatment_col, list) and len(after_treatment_col) > 0):
         if isinstance (after_treatment_col, str):
@@ -1458,20 +1496,14 @@ def did_analysis(
         outcome_col = "log_"+f'{outcome_col}'
 
     did_formula = f'{outcome_col} ~ {" + ".join(treatment_col)}'
-
-    if TG_col is not None or len(TG_col) > 0:
+    
+    if TG_col is not None and len(TG_col) > 0:
         did_formula = did_formula + f' + {" + ".join(TG_col)}'        
-    if TT_col is not None or len(TT_col) > 0:
-        did_formula = did_formula + f' + {" + ".join(TT_col)}'
-   
-    if ITT:
-        FE_unit = True
-        FE_time = False
-    if ITE:
-        FE_unit = True
+    if TT_col is not None and len(TT_col) > 0:
+        did_formula = did_formula + f' + {" + ".join(TT_col)}'   
 
     if len(after_treatment_col) > 0:
-        did_formula = did_formula + f'+ {" + ".join(after_treatment_col)}'
+        did_formula = did_formula + f' + {" + ".join(after_treatment_col)}'
     
     if FE_unit:
         unit_col_todummies = diffindiff.didtools.to_dummies(
@@ -1481,7 +1513,7 @@ def did_analysis(
             drop_first = intercept
             )        
         data = unit_col_todummies[0]
-        did_formula = did_formula + f'+ {unit_col_todummies[1]}'
+        did_formula = did_formula + f' + {unit_col_todummies[1]}'
         dummy_unit_vars = list(unit_col_todummies[2]["UNIT_"+unit_col].values)
         dummy_unit_original = list(unit_col_todummies[2][unit_col].values)
 
@@ -1493,7 +1525,7 @@ def did_analysis(
             drop_first = intercept
             )
         data = time_col_todummies[0]
-        did_formula = did_formula + f'+ {time_col_todummies[1]}'        
+        did_formula = did_formula + f' + {time_col_todummies[1]}'        
         dummy_time_vars = list(time_col_todummies[2]["TIME_"+time_col].values)
         dummy_time_original = list(time_col_todummies[2][time_col].values)   
 
@@ -1526,8 +1558,8 @@ def did_analysis(
                 new_col_name = f"{col}_x_time"
                 group_x_time = group_x_time.rename(columns={col: new_col_name})
             data = pd.concat([data, group_x_time], axis = 1)
-            GTT_columns_groupxtime = '+'.join(group_x_time.columns)
-            did_formula = did_formula + f'+{GTE_columns_group}+{GTT_columns_groupxtime}'
+            GTT_columns_groupxtime = ' + '.join(group_x_time.columns)
+            did_formula = did_formula + f' + {GTE_columns_group} + {GTT_columns_groupxtime}'
     
     if ITT:        
         if "date_counter" not in data.columns:
@@ -1542,7 +1574,7 @@ def did_analysis(
             new_col_name = f"{col}_x_time"
             unit_x_time = unit_x_time.rename(columns={col: new_col_name})
         data = pd.concat([data, unit_x_time], axis = 1)
-        ITT_columns_unitxtime = '+'.join(unit_x_time.columns)
+        ITT_columns_unitxtime = ' + '.join(unit_x_time.columns)
         did_formula = did_formula + f' + {ITT_columns_unitxtime}'
     
     if GTE:
@@ -1556,8 +1588,8 @@ def did_analysis(
                     new_col_name = f"{treatment}_{col}_x_time"
                     group_x_treatment = group_x_treatment.rename(columns={col: new_col_name})
             data = pd.concat([data, group_x_treatment], axis = 1)
-            GTE_columns_groupxtreatment = '+'.join(group_x_treatment.columns)
-            did_formula = did_formula + f'+{GTE_columns_group}+{GTE_columns_groupxtreatment}'
+            GTE_columns_groupxtreatment = ' + '.join(group_x_treatment.columns)
+            did_formula = did_formula + f' + {GTE_columns_group} + {GTE_columns_groupxtreatment}'
     
     if ITE:
         unit_x_treatment = pd.DataFrame()
@@ -1574,7 +1606,7 @@ def did_analysis(
         if group_by in covariates:
             covariates.remove(group_by)
         covariates_join = ' + '.join(covariates)
-        did_formula = did_formula + f'+{covariates_join}'
+        did_formula = did_formula + f' +{covariates_join}'
 
     if len(group_benefit) > 0:        
         group_benefit = diffindiff.didtools.unique(group_benefit)
@@ -1597,10 +1629,11 @@ def did_analysis(
         group_benefit = []
         DDD = False
     
-    if GTE or GTT or ITE or ITT:
-        intercept = False   
+    did_formula = did_formula[:-1] if did_formula.endswith(" ") else did_formula
+    did_formula = did_formula[:-1] if did_formula.endswith("+") else did_formula
+    did_formula = did_formula[:-1] if did_formula.endswith(" ") else did_formula
     if not intercept:
-        did_formula = did_formula + f' -1'
+        did_formula = did_formula + f' -1'    
     
     analysis_description = "Difference in Differences (DiD) Analysis"
     if DDD:
@@ -1622,6 +1655,7 @@ def did_analysis(
         "pre_post": pre_post,
         "FE_unit": FE_unit,
         "FE_time": FE_time,
+        "FE_group": FE_group,
         "intercept": intercept,
         "ITT": ITT,
         "GTT": GTT,
@@ -1831,7 +1865,7 @@ def did_analysis(
         FE_group_coef = {}        
         for i, group_dummy in enumerate(FE_group_vars):
             FE_group_coef[i] = {
-                "Coefficient": group_dummy,
+                "Coefficient": dummy_group_original[i],
                 "Estimate": ols_coefficients[group_dummy],
                 "SE": float(coef_standarderrors[group_dummy]),
                 "t": float(coef_t[group_dummy]),
@@ -1955,7 +1989,10 @@ def did_analysis(
             model_results["covariates_effects"] = covariates_effects
 
     model_predictions = ols_model.predict()
-
+    
+    prediction_intervals = ols_model.get_prediction()
+    prediction_intervals = prediction_intervals.summary_frame(alpha = confint_alpha)
+    
     model_statistics = {
         "rsquared": ols_model.rsquared,
         "rsquared_adj": ols_model.rsquared_adj,
@@ -1968,7 +2005,8 @@ def did_analysis(
         data,
         model_predictions,
         model_statistics,
-        ols_model
+        ols_model,
+        prediction_intervals
         ) 
 
     return did_model_output

{diffindiff-2.0.1 → diffindiff-2.0.3}/diffindiff/diddata.py

@@ -1,11 +1,14 @@
-#-------------------------------------------------------------------------------
-# Name:        diddata (diffindiff)
+#-----------------------------------------------------------------------
+# Name:        diddata (diffindiff package)
 # Purpose:     Creating data for Difference-in-Differences Analysis
-# Author:      Thomas Wieland (mail: geowieland@googlemail.com, ORCID: 0000-0001-5168-9846)
-# Version:     2.0.1
-# Last update: 2025-04-15 18:43
+# Author:      Thomas Wieland 
+#              ORCID: 0000-0001-5168-9846
+#              mail: geowieland@googlemail.com              
+# Version:     2.0.3
+# Last update: 2025-04-18 10:24
 # Copyright (c) 2025 Thomas Wieland
-#-------------------------------------------------------------------------------
+#-----------------------------------------------------------------------
+
 
 import pandas as pd
 import numpy as np
@@ -950,7 +953,7 @@ def create_counterfactual(
         )
     control_group = isnotreatment[2]
     
-    units_tt = didtools.treatment_times(
+    units_tt = diffindiff.didtools.treatment_times(
         data = data,
         unit_col = unit_col,
         time_col = time_col,
@@ -959,7 +962,7 @@ def create_counterfactual(
     units = diffindiff.didtools.unique(units_tt[unit_col])  
     
     if not isnotreatment[0]:
-        print ("No no-treatment control group. Counterfactual will not cover full treatment time.")    
+        print ("NOTE: No no-treatment control group. Counterfactual will not cover full treatment time.")    
     
     data_TG = pd.DataFrame(columns = data.columns)
     for unit in units:
@@ -980,7 +983,9 @@ def create_counterfactual(
             [data_TG, data_CG],
             ignore_index=True
         )
-        
+    
+    data_cf[X] = data_cf[X].apply(pd.to_numeric, errors='coerce')
+    
     counterfactual_pred = diffindiff.didtools.model_wrapper(
         y = data_cf[y],
         X = data_cf[X],

{diffindiff-2.0.1 → diffindiff-2.0.3}/diffindiff/didtools.py

@@ -1,11 +1,13 @@
-# -------------------------------------------------------------------------------
-# Name:        didtools (diffindiff)
-# Purpose:     Creating data for Difference-in-Differences Analysis
-# Author:      Thomas Wieland (mail: geowieland@googlemail.com, ORCID: 0000-0001-5168-9846)
-# Version:     2.0.1
-# Last update: 2025-04-15 18:44
+#-----------------------------------------------------------------------
+# Name:        didtools (diffindiff package)
+# Purpose:     Additional tools for Difference-in-Differences Analysis
+# Author:      Thomas Wieland 
+#              ORCID: 0000-0001-5168-9846
+#              mail: geowieland@googlemail.com              
+# Version:     2.0.3
+# Last update: 2025-04-18 12:08
 # Copyright (c) 2025 Thomas Wieland
-#-------------------------------------------------------------------------------
+#-----------------------------------------------------------------------
 
 
 import pandas as pd
@@ -34,11 +36,11 @@ def check_columns(
         raise ValueError(f"Data do not contain column(s): {', '.join(missing_columns)}")
 
 def is_balanced(
-    data,
-    unit_col,
-    time_col,
-    outcome_col,
-    other_cols = None
+    data: pd.DataFrame,
+    unit_col: str,
+    time_col: str,
+    outcome_col: str,
+    other_cols: list = None
     ):
 
     unit_freq = data[unit_col].nunique()
@@ -58,8 +60,8 @@ def is_balanced(
         return True
 
 def is_binary(
-    data,
-    treatment_col
+    data: pd.DataFrame,
+    treatment_col: str
     ):
     
     unique_values = set(data[treatment_col].dropna().unique())
@@ -76,7 +78,7 @@ def is_binary(
         return [False, "Unknown"]             
 
 def is_missing(
-    data,
+    data: pd.DataFrame,
     drop_missing: bool = True,
     missing_replace_by_zero: bool = False
     ):
@@ -104,10 +106,10 @@ def is_missing(
         ]
 
 def is_simultaneous(
-    data,
-    unit_col,
-    time_col,
-    treatment_col,
+    data: pd.DataFrame,
+    unit_col: str,
+    time_col: str,
+    treatment_col: str,
     pre_post = False
     ):
 
@@ -125,9 +127,9 @@ def is_simultaneous(
     return col_identical
 
 def is_notreatment(
-    data,
-    unit_col,
-    treatment_col
+    data: pd.DataFrame,
+    unit_col: str,
+    treatment_col: str
     ):
 
     data_relevant = data[[unit_col, treatment_col]]
@@ -150,12 +152,52 @@ def is_notreatment(
         control_group
         ]
 
+def treatment_group_col(
+    data: pd.DataFrame,
+    unit_col: str,
+    treatment_col: str,
+    create_TG_col: str = "TG"   
+    ):
+    
+    isnotreatment = is_notreatment(
+        data = data,
+        unit_col = unit_col,
+        treatment_col = treatment_col
+        )
+    
+    if not isnotreatment[0]:
+        print ("Model data does not contain a no-treatment control group. Treatment group column is constant = 1.")
+    
+    if create_TG_col in data.columns:
+        create_TG_col = "TG_"+treatment_col
+        print ("Column " + create_TG_col + " already exists. Saving treatment group in column TG_" + treatment_col)
+
+    treatment_group = isnotreatment[1]
+
+    data[create_TG_col] = 0
+    data.loc[data[unit_col].astype(str).isin(treatment_group), create_TG_col] = 1
+
+    return [
+        data, 
+        isnotreatment[0], 
+        create_TG_col
+        ]
+
+def untreated_units(
+    data: pd.DataFrame,
+    unit_col: str,
+    time_col: str,
+    treatment_col: str   
+    ):
+    # TODO ??
+    pass
+
 def is_parallel(
-    data,
-    unit_col,
-    time_col,
-    treatment_col,
-    outcome_col,
+    data: pd.DataFrame,
+    unit_col: str,
+    time_col: str,
+    treatment_col: str,
+    outcome_col: str,
     pre_post = False,
     alpha = 0.05
     ):
@@ -206,10 +248,10 @@ def is_parallel(
         ]
 
 def date_counter(
-        df,
-        date_col, 
-        new_col = "date_counter"
-        ):
+    df: pd.DataFrame,
+    date_col: str, 
+    new_col: str = "date_counter"
+    ):
     
     dates = df[date_col].unique()
 
@@ -226,6 +268,7 @@ def date_counter(
     return df
 
 def unique(data):
+
     if data is None or (isinstance(data, (list, np.ndarray, pd.Series, pd.DataFrame)) and len(data) == 0):
         return []
     
@@ -269,8 +312,9 @@ def model_wrapper(
     lgbm_learning_rate = 0.1,
     random_state = 71
     ):
-    if model_type not in ["ols", "olsbg", "dtbg", "rf", "gb", "knn", "svr", "xgb", "lgbm", "catboost"]:
-        raise ValueError("Please enter a valid model type")
+
+    if model_type not in ["ols", "olsbg", "dtbg", "rf", "gb", "knn", "svr", "xgb", "lgbm"]:
+        raise ValueError("Please enter a valid model type ('ols', 'olsbg', 'dtbg', 'rf', 'gb', 'knn', 'svr', 'xgb', 'lgbm')")
     
     X_train, X_test, y_train, y_test = train_test_split(
         X, 
@@ -348,10 +392,10 @@ def model_wrapper(
         ]
     
 def treatment_times(
-    data,
-    unit_col,
-    time_col,
-    treatment_col
+    data: pd.DataFrame,
+    unit_col: str,
+    time_col: str,
+    treatment_col: str
     ):
     
     check_columns(
@@ -389,10 +433,10 @@ def clean_column_name(value):
     return value.strip('_')
     
 def to_dummies(
-    data, 
-    col,
-    drop_first = False,
-    prefix = "DUMMY"
+    data: pd.DataFrame, 
+    col: str,
+    drop_first: bool = False,
+    prefix: str = "DUMMY"
     ):     
     
     unique_values = data[col].astype(str).unique()

{diffindiff-2.0.1 → diffindiff-2.0.3}/diffindiff/tests/tests_diffindiff.py

@@ -1,17 +1,19 @@
-#------------------------------------------------------------------------------------------
-
-# Name:        tests_diffindiff
+#-----------------------------------------------------------------------
+# Name:        tests_diffindiff (diffindiff package)
 # Purpose:     Tests and examples for the diffindiff package
-# Author:      Thomas Wieland (mail: geowieland@googlemail.com, ORCID: 0000-0001-5168-9846)
-# Version:     2.0.1
-# Last update: 2025-04-15 18:43
+# Author:      Thomas Wieland 
+#              ORCID: 0000-0001-5168-9846
+#              mail: geowieland@googlemail.com              
+# Version:     2.0.3
+# Last update: 2025-04-18 10:24
 # Copyright (c) 2025 Thomas Wieland
-#------------------------------------------------------------------------------------------
+#-----------------------------------------------------------------------
 
 
 import pandas as pd
 from diffindiff.didanalysis import DiffModel, did_analysis
 from diffindiff.diddata import DiffGroups, create_groups, DiffTreatment, create_treatment, DiffData, merge_data, create_data
+from diffindiff.didtools import treatment_group_col
 
 # Example 1: Effect of a curfew in German counties in the first
 # wave of the COVID-19 pandemic (DiD pre-post analysis)

{diffindiff-2.0.1 → diffindiff-2.0.3}/diffindiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: diffindiff
-Version: 2.0.1
+Version: 2.0.3
 Summary: diffindiff: Python library for convenient Difference-in-Differences Analyses
 Author: Thomas Wieland
 Author-email: geowieland@googlemail.com
@@ -38,17 +38,18 @@ Thomas Wieland [ORCID](https://orcid.org/0000-0001-5168-9846) [EMail](mailto:geo
   - Create predictive counterfactuals
 - **DiD analysis**: 
   - Perfom standard DiD analysis  
-  - Model Extensions:
+  - Model extensions:
     - Staggered adoption
     - Multiple treatments
     - Two-way fixed effects models
     - Group- or individual-specific treatment effects
     - Group- or individual-specific time trends
     - Including covariates
-    - After-treatment period
+    - Including fter-treatment period
     - Triple Difference (DDD)
     - Own counterfactuals
-    - Bonferroni correction
+    - Bonferroni correction for treatment effects
+    - Placebo test
 - **Visualization**:
   - Plot observed and expected time course of treatment and control group
   - Plot expected time course of treatment group and counterfactual
@@ -60,7 +61,6 @@ Thomas Wieland [ORCID](https://orcid.org/0000-0001-5168-9846) [EMail](mailto:geo
   - Test for type of adoption
   - Test whether the panel dataset is balanced
   - Test for parallel trend assumption
-  - Placebo test
 
 
 ## Literature

{diffindiff-2.0.1 → diffindiff-2.0.3}/setup.py

@@ -7,7 +7,7 @@ def read_README():
     
 setup(
     name='diffindiff',
-    version='2.0.1',
+    version='2.0.3',
     description='diffindiff: Python library for convenient Difference-in-Differences Analyses',
     packages=find_packages(include=["diffindiff", "diffindiff.tests"]),
     include_package_data=True,