Repository: d2cml-ai/csdid (language: python)

Issue:
allow_unbalanced_panel = False is not working 
```
from csdid.att_gt import ATTgt 
import pandas as pd 
# Load the CSV file 
file_path = 'C:/Users/Alexander/Documents/GitHub/csdid/data/cohort_2023.csv' 
data = pd.read_csv(file_path)
# Prepare  the data
data['customer_id_numeric'] = data['customer_id'].astype('category').cat.codes 
# Create an instance of the ATTgt class and fit the model 
out = ATTgt(yname="ln_gms",gname="first.treat", idname="customer_id_numeric",        
     tname="time_period", data=data, allow_unbalanced_panel = False).fit(est_method='ipw', bstrap = True) 
# Display the results 
out.summ_attgt().summary2
```

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
c:\Users\Alexander\anaconda3\envs\tensorflow\lib\site-packages\pandas\core\indexing.py in _validate_tuple_indexer(self, key)
    903             try:
--> 904                 self._validate_key(k, i)
    905             except ValueError as err:

c:\Users\Alexander\anaconda3\envs\tensorflow\lib\site-packages\pandas\core\indexing.py in _validate_key(self, key, axis)
   1517         else:
-> 1518             raise ValueError(f"Can only index by location with a [{self._valid_types}]")
   1519 

ValueError: Can only index by location with a [integer, integer slice (START point is INCLUDED, END point is EXCLUDED), listlike of integers, boolean array]

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
C:\Users\ALEXAN~1\AppData\Local\Temp/ipykernel_24432/3877272001.py in <module>
     10 data['customer_id_numeric'] = data['customer_id'].astype('category').cat.codes
     11 # Create an instance of the ATTgt class and fit the model
---> 12 out = ATTgt(yname="ln_gms",gname="first.treat", idname="customer_id_numeric",        
     13      tname="time_period", data=data, allow_unbalanced_panel = False).fit(est_method='ipw', bstrap = True) 
     14 # Display the results

c:\Users\Alexander\anaconda3\envs\tensorflow\lib\site-packages\csdid\att_gt.py in __init__(self, yname, tname, idname, gname, data, control_group, xformla, panel, allow_unbalanced_panel, clustervar, weights_name, anticipation, cband, biters, alp)
...
--> 906                 raise ValueError(
    907                     "Location based indexing can only have "
    908                     f"[{self._valid_types}] types"

ValueError: Location based indexing can only have [integer, integer slice (START point is INCLUDED, END point is EXCLUDED), listlike of integers, boolean array] types
Output is truncated. View as a [scrollable element](command:cellOutput.enableScrolling?4242d70f-e8c8-4f7f-938c-97ed0d0531c4) or open in a [text editor](command:workbench.action.openLargeOutput?4242d70f-e8c8-4f7f-938c-97ed0d0531c4). Adjust cell output [settings](command:workbench.action.openSettings?%5B%22%40tag%3AnotebookOutputLayout%22%5D)...
```


IndexError while using ATTgt package 
Hi, 

I am trying to understand the updated package and running it against a simulated data as follows: 

```
out = ATTgt(yname='outcome',
           gname='treatment_month',
           idname='seller_id',
           tname='month',
           allow_unbalanced_panel=True,
           xformla=f'outcome~1',
           control_group='never_treated',
           data=panel_data.reset_index()
           ).fit(est_method='dr')
```

I am getting a following error: 
```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[8], line 9
      1 out = ATTgt(yname='outcome',
      2            gname='treatment_month',
      3            idname='seller_id',
      4            tname='month',
      5            allow_unbalanced_panel=True,
      6            xformla=f'outcome~1',
      7            control_group='never_treated',
      8            data=panel_data.reset_index()
----> 9            ).fit(est_method='dr')

File ~\AppData\Roaming\Python\Python311\site-packages\csdid\att_gt.py:39, in ATTgt.fit(self, est_method, base_period, bstrap)
     36 def fit(self, est_method = 'dr', base_period = 'varying', bstrap = True):
     37   # print(self.dp)
     38   dp = self.dp
---> 39   result, inffunc = compute_att_gt(dp)
     40   att = result['att']
     41   crit_val, se, V = np.zeros(len(att)), np.zeros(len(att)), np.zeros(len(att))

File ~\AppData\Roaming\Python\Python311\site-packages\csdid\attgt_fnc\compute_att_gt.py:83, in compute_att_gt(dp, est_method, base_period)
     81 n1 = data[gname] == 0
     82 n2 = (data[gname] > (tlist[np.max([t_i, pret]) + tfac]) + anticipation)
---> 83 n3 = np.where(data[gname] != glist[g], True, False)
     84 row_eval = n1 | n2 & n3
     85 data = data.assign(C = 1 * row_eval)

IndexError: index 15 is out of bounds for axis 0 with size 4
```

I tried changing control group definition and other parameters but still got the same error. 

I am attaching a code to generate simulated data: 
```
# Set seed for reproducibility
np.random.seed(51)

# Define parameters
num_months = 36
num_sellers = 100
mean_outcome = 10
noise_level = 0.5
treatment_effect = 2
treatment_noise_level = 0.2
treatment_months = [15, 18, 21, 24]
cohort_sizes = [15, 15, 15, 15]

# Initialize data structure
data = []

# Generate seller data
for seller_id in range(1, num_sellers + 1):
    # Randomly assign join month between 1 and 12
    join_month = np.random.randint(1, 13)
    
    # Randomly decide if seller leaves between month 24 and 36
    leave_month = np.random.choice([0] + list(range(24, 37)), p=[0.75] + [0.25 / 13] * 13)
    
    # Default outcome with noise
    outcome = mean_outcome + np.random.normal(0, noise_level, num_months)
    
    # Treatment assignment
    if seller_id <= 40:
        treatment_status = 0
        treatment_month = np.nan
    else:
        cumulative_sizes = np.cumsum(cohort_sizes)
        if seller_id <= 40 + cumulative_sizes[0]:
            cohort = 0
        elif seller_id <= 40 + cumulative_sizes[1]:
            cohort = 1
        elif seller_id <= 40 + cumulative_sizes[2]:
            cohort = 2
        else:
            cohort = 3
        
        treatment_month = treatment_months[cohort]
        treatment_status = 1
        
        # Apply treatment effect after treatment month
        outcome[treatment_month - 1:] += treatment_effect + np.random.uniform(-treatment_noise_level, treatment_noise_level, num_months - treatment_month + 1)
    
    # Create panel data for each month
    for month in range(1, num_months + 1):
        if month >= join_month and (leave_month == 0 or month < leave_month):
            data.append([seller_id, month, outcome[month - 1], treatment_status, treatment_month])

# Create DataFrame
columns = ['seller_id', 'month', 'outcome', 'treatment_status', 'treatment_month']
panel_data = pd.DataFrame(data, columns=columns)
panel_data['seller_id'] = panel_data['seller_id'].astype(str)
panel_data = panel_data.set_index(keys=['seller_id', 'month'])

panel_data = panel_data.reset_index()
panel_data = panel_data.sort_values(by=['seller_id', 'month'], ascending=False)
panel_data['outcome12'] = panel_data.groupby('seller_id')['outcome'].transform(lambda x: x.rolling(12).sum())
panel_data['outcome12'] = panel_data['outcome12'].fillna(panel_data.groupby('seller_id')['outcome12'].transform('mean'))
panel_data = panel_data.set_index(keys=['seller_id', 'month'])

#Specifically for CSDID
panel_data = panel_data.reset_index()
panel_data['seller_id'] = panel_data['seller_id'].astype('int64')
panel_data['treatment_month'].fillna(0, inplace=True)
panel_data['treatment_month'] = panel_data['treatment_month'].astype('int64')
```

I also tried the code with the data provided in the repository example, and it works perfectly fine. 
My simulated data looks exactly like the data you provided. I will really really appreciate your help to debug the issue here. 

Thank you. 





Discussion / hints:

Fix patch:
diff --git a/csdid/attgt_fnc/compute_att_gt.py b/csdid/attgt_fnc/compute_att_gt.py
index f300084..d4917fa 100644
--- a/csdid/attgt_fnc/compute_att_gt.py
+++ b/csdid/attgt_fnc/compute_att_gt.py
@@ -3,216 +3,505 @@
 from drdid import drdid, reg_did, ipwd_did
 
 from csdid.utils.bmisc import panel2cs2
+import warnings
 
-fml = patsy.dmatrices
-
-
-# data Struct
-# output = {
-#   "att" : []
-#   'group': []
-#   'year': []
-#   'post ': []
-# }
 
+fml = patsy.dmatrices
+# Initialize a list to store data for each iteration
+results_list = []
 
 def compute_att_gt(dp, est_method = "dr", base_period = 'varying'):
-  yname = dp['yname']
-  tname = dp['tname']
-  idname = dp['idname']
-  xformla = dp['xformla']
-  data = dp['data']
-  weights_name = dp['weights_name']
-  # base_period = dp['base_period']
-  panel = dp['panel']
-  true_rep_cross_section = dp['true_rep_cross_section']
-  control_group = dp['control_group']
-  anticipation = dp['anticipation']
-  gname = dp['gname']
-  n = dp['n']
-  nT = dp['nT']
-  nG = dp['nG']
-  tlist = dp['tlist']
-  glist = dp['glist']
-  tlist_len, tfac = len(tlist), 0
-  if base_period != 'universal':
-    tlist_len = tlist_len - 1
-    tfac = 1
-
-  inf_func = []
-
-  att_est, group, year, post_array = [], [], [], []
-
-  def add_att_data(att = 0, pst = 0, inf_f = []):
-    inf_func.append(inf_f)
-    att_est.append(att)
-    group.append(g)
-    year.append(tn)
-    post_array.append(pst)
-
-  never_treated = control_group == 'nevertreated'
-  if never_treated:
-    data = data.assign(C = 1 * (data[gname] == 0))
-  data = data.assign(y_main = data[yname])
-
-  # g, t = glist[0], tlist[0]
-
-  # for _, g, in enumerate(glist):
-  for g_index, g in enumerate(glist):
-
-    # g = glist[1]
-    G_main = (data[gname] == g)
-    data = data.assign(G_m = 1 * G_main)
-
-    for t_i in range(tlist_len):
-      pret = t_i
-      tn = tlist[t_i + tfac]
-      if base_period == 'universal' or g < tn:
-        try:
-          pret = np.where(tlist + anticipation < g)[0][-1]
-        except:
-          raise f"There are no pre-treatment periods for the group first treated at {g}\nUnits from this group are dropped"
-          # break
-
-
-      if base_period == 'universal':
-        if pret == tn:
-          add_att_data()
-
-      if not never_treated:
-        n1 = data[gname] == 0
-        n2 = (data[gname] > (tlist[np.max([t_i, pret]) + tfac]) + anticipation)
-        #n3 = np.where(data[gname] != glist[g], True, False)
-        n3 = np.where(data[gname] != glist[g_index], True, False)
-
-        row_eval = n1 | n2 & n3
-        data = data.assign(C = 1 * row_eval)
-
-      post_treat = 1 * (g <= tn)
-      disdat = data[(data[tname] == tn) | (data[tname] == tlist[pret])]
-
-      if panel: 
-        disdat = panel2cs2(disdat, yname, idname, tname)
-        disdat = disdat.dropna()
-        n = len(disdat)
-        dis_idx = np.array(disdat.G_m == 1) | np.array(disdat.C == 1)
-        disdat = disdat.loc[dis_idx, :]
-        n1 = len(disdat)
-        G = disdat.G_m
-        C = disdat.C
-        w = disdat.w
-
-        ypre = disdat.y0 if tn > pret else disdat.y1
-        ypost = disdat.y0 if tn < pret else disdat.y1
-        _, covariates = fml(xformla, data = disdat, return_type = 'dataframe')
-
-        G, C, w, ypre = map(np.array, [G, C, w, ypre])
-        ypost, covariates = map(np.array, [ypost, covariates])
-
-        if callable(est_method):
-          est_att_f = est_method
-        elif est_method == "reg":
-          est_att_f = reg_did.reg_did_panel
-        elif est_method == "ipw":
-          est_att_f = ipwd_did.std_ipw_did_panel
-        elif est_method == "dr":
-          est_att_f = drdid.drdid_panel
-
-        att_gt, att_inf_func = est_att_f(ypost, ypre, G, i_weights=w, covariates=covariates)
-
-        inf_zeros = np.zeros(n)
-        att_inf = n / n1 * att_inf_func
-        inf_zeros[dis_idx] = att_inf
-
-        add_att_data(att_gt, inf_f=inf_zeros)
-
-      if not panel:
-        right_ids = np.array(disdat.query('(G_m == 1) or (C == 1)').rowid.to_numpy())
-        dis_idx = (data['rowid'].isin(right_ids)) &\
-          ((data[tname] == tlist[t_i + tfac]) |\
-            (data[tname] == tlist[pret]))
-
-        disdat = data.loc[dis_idx]
-
-        G = disdat.G_m.to_numpy()
-        C = disdat.C.to_numpy()
-        Y = disdat[yname].to_numpy()
-        post = 1 * (disdat[tname] == tlist[t_i + tfac]).to_numpy()
-        w = disdat.w.to_numpy()
-
-        # G, C, Y, post, w = map(np.array, [G, C, Y, post, w])
-
-
-        n1 = sum(G + C)
-
-        skip_this_att_gt = False
-        if np.sum(G * post) == 0:
-          print(f"No units in group {g} in time period {tn}")
-          skip_this_att_gt = True 
-
-        if np.sum(G * (1 - post)) == 0:
-          print(f"No units in group {g} in time period {tn}")
-          skip_this_att_gt = True 
-
-        if np.sum(C * post) == 0:
-          print(f"No available control units for group {g} in time period {tn}")
-          skip_this_att_gt = True 
-
-        if np.sum(C * (1 - post)) == 0:
-          print(f"No available control units for group {g} in time period {tn}")
-          skip_this_att_gt = True 
-
-        if skip_this_att_gt:
-          add_att_data()
-
-        try:
-          _, covariates = fml(xformla, data = disdat, return_type = 'dataframe')
-          covariates = np.array(covariates)
-        except:
-          y_str, x_str = xformla.split("~")
-          xs1 = x_str.split('+')
-          xs1_col_names = [x.strip() for x in xs1 if x.strip() != '1']
-          n_dis = len(disdat)
-          ones = np.ones((n_dis, 1))
-          try:
-            covariates = disdat[xs1_col_names].to_numpy()
-            covariates = np.append(covariates, ones, axis=1)
-          except:
-            covariates = ones
-
-        
-
-        if callable(est_method):
-          est_att_f = est_method
-        elif est_method == "reg":
-          est_att_f = reg_did.reg_did_rc
-        elif est_method == "ipw":
-          est_att_f = ipwd_did.std_ipw_did_rc
-        elif est_method == "dr":
-          est_att_f = drdid.drdid_rc
-        att_gt, att_inf_func = est_att_f(y=Y, post=post, D = G, i_weights=w, covariates=covariates)
-
-        inf_func_df = pd.DataFrame(
-          {
-            "inf_func": att_inf_func,
-            "right_ids": right_ids
-          }
-        )
-        inf_zeros = np.zeros(n)
-        aggte_infffuc = inf_func_df.groupby('right_ids').inf_func.sum()
-        try:
-          dis_idx1 = np.isin(data['rowid'].unique(), aggte_infffuc.index.to_numpy())
-        except:
-          dis_idx1 = np.isin(data['rowid'].unique().to_numpy(), aggte_infffuc.index.to_numpy())
-        inf_zeros[dis_idx1] = np.array(aggte_infffuc)
-
-        add_att_data(att_gt, pst = post_treat, inf_f=inf_zeros)
-        # print(att_est)
-
-  output = {
+    yname = dp['yname']
+    tname = dp['tname']
+    idname = dp['idname']
+    xformla = dp['xformla']
+    data = dp['data']
+    weights_name = dp['weights_name']
+    # base_period = dp['base_period']
+    panel = dp['panel']
+    true_rep_cross_section = dp['true_rep_cross_section']
+    control_group = dp['control_group']
+    anticipation = dp['anticipation']
+    gname = dp['gname']
+    n = dp['n']
+    nT = dp['nT']
+    nG = dp['nG']
+    tlist = dp['tlist']
+    glist = dp['glist']
+
+    # Calculate time periods and adjustment factor
+    tlist_len = len(tlist) - 1 if base_period != "universal" else len(tlist)
+    tfac = 1 if base_period != "universal" else 0
+
+    inf_func = []
+
+    att_est, group, year, post_array = [], [], [], []
+
+    def add_att_data(att = 0, pst = 0, inf_f = []):
+        inf_func.append(inf_f)
+        att_est.append(att)
+        group.append(g)
+        year.append(tn)
+        post_array.append(pst)
+
+    # Handle never treated case
+    never_treated = control_group == 'nevertreated'
+    if never_treated:
+        data['C'] = (data[gname] == 0).astype(int)
+    data['y_main'] = data[yname]
+
+    # Loop over groups
+    for g_index, g in enumerate(glist):  
+        # Set up .G once
+        # Create a binary column 'G_m' to indicate if a row belongs to the current group 'g'
+        G_main = (data[gname] == glist[g_index])
+        data = data.assign(G_m=1 * G_main)
+
+        # Loop over time periods
+        for t_i in range(tlist_len):
+
+            # Set pretreatment period
+            pret = t_i  # Initialize pretreatment period as current time period index
+            tn = tlist[t_i + tfac]  # Current time period (adjusted for tfac)
+
+            # Universal base period
+            if base_period == 'universal':  # Check if using a universal base period
+                try:
+                    # Set pretreatment period as the last period before treatment
+                    pret = np.where((tlist + anticipation) < g)[0][-1]
+                except IndexError:
+                    # Handle cases where no pretreatment periods exist
+                    raise ValueError(
+                        f"There are no pre-treatment periods for the group first treated at {g}. Units from this group are dropped."
+                    )
+
+            # If we are in the universal base period
+            if base_period == 'universal' and pret == tn:
+                # Normalize results to zero and skip computation
+                add_att_data(att=0, pst=0, inf_f=np.zeros(len(data)))
+                continue
+
+            # For non-never treated groups, set up the control group indicator 'C'
+            if not never_treated:
+                # Units that are either never treated (gname == 0) or treated in the future
+                n1 = (data[gname] == 0)
+                n2 = (data[gname] > (tlist[max(t_i, pret) + tfac] + anticipation))
+                n3 = (data[gname] != glist[g_index])  # Not in the current group
+                row_eval = n1 | (n2 & n3)  # Combine conditions
+                data = data.assign(C=1 * row_eval)  # Assign the control indicator
+
+        # -----------------------------------------------------------------------------
+  
+            
+            # Check if in post-treatment period
+            if glist[g_index] <= tlist[t_i + tfac]:
+                # Use same base period as for post-treatment periods
+                # This matches R's: tail(which((tlist+anticipation) < glist[g]), 1)
+                pret_mask = (np.array(tlist) + anticipation) < glist[g_index]
+                if not any(pret_mask):
+                    warnings.warn(f"There are no pre-treatment periods for the group first treated at {glist[g_index]}\nUnits from this group are dropped")
+                    break
+                
+                pret = np.where(pret_mask)[0][-1]  # Gets last element like R's tail(..., 1)
+                # print("this is alex", pret)
+
+
+            # Print details for debugging
+            # print(f"Current period: {tlist[t_i + tfac]}")
+            # print(f"Current group: {glist[g_index]}")
+            # print(f"Set pretreatment period to be: {tlist[pret]}")
+
+            # -----------------------------------------------------------------------------
+            # Debugging and validation
+            if base_period == 'universal' and pret == tn:
+                # Normalize results to zero and break the loop
+                add_att_data(att=0, pst=post_treat, inf_f=np.zeros(len(data)))
+                continue
+            
+            # Post-treatment dummy variable
+            post_treat = 1 * (g <= tn)
+
+            # Subset the data for the current and pretreatment periods
+            disdat = data[(data[tname] == tn) | (data[tname] == tlist[pret])]
+            # print("Shape of disdat:", disdat.shape)
+
+        # results for the case with panel data
+        #-----------------------------------------------------------------------------
+
+            if panel: 
+                disdat = panel2cs2(disdat, yname, idname, tname)
+                disdat = disdat.dropna()
+                n = len(disdat)
+                dis_idx = np.array(disdat.G_m == 1) | np.array(disdat.C == 1)
+                disdat = disdat.loc[dis_idx, :]
+                n1 = len(disdat)
+                G = disdat.G_m
+                C = disdat.C
+                w = disdat.w
+
+                ypre = disdat.y0 if tn > pret else disdat.y1
+                ypost = disdat.y0 if tn < pret else disdat.y1
+                _, covariates = fml(xformla, data = disdat, return_type = 'dataframe')
+
+                G, C, w, ypre = map(np.array, [G, C, w, ypre])
+                ypost, covariates = map(np.array, [ypost, covariates])
+
+                if callable(est_method):
+                    est_att_f = est_method
+                elif est_method == "reg":
+                    est_att_f = reg_did.reg_did_panel
+                elif est_method == "ipw":
+                    est_att_f = ipwd_did.std_ipw_did_panel
+                elif est_method == "dr":
+                    est_att_f = drdid.drdid_panel
+
+                att_gt, att_inf_func = est_att_f(ypost, ypre, G, i_weights=w, covariates=covariates)
+
+                inf_zeros = np.zeros(n)
+                att_inf = n / n1 * att_inf_func
+                inf_zeros[dis_idx] = att_inf
+
+                add_att_data(att_gt, inf_f=inf_zeros)
+
+        #-----------------------------------------------------------------------------
+        # results for the case with no panel data
+        #-----------------------------------------------------------------------------
+            
+            if not panel:
+                # Fixed right_ids selection
+                right_ids = disdat.loc[disdat.G_m.eq(1) | disdat.C.eq(1), 'rowid'].to_numpy()
+                # print("Shape of ids:",right_ids.shape)  # Show dimensions
+                # Consistent time period selection
+                dis_idx = (data['rowid'].isin(right_ids)) & \
+                            (data[tname].isin([tlist[t_i + tfac], tlist[pret]]))
+
+                disdat = data.loc[dis_idx]
+
+                G = disdat.G_m.to_numpy()
+                C = disdat.C.to_numpy()
+                Y = disdat[yname].to_numpy()
+                post = 1 * (disdat[tname] == tlist[t_i + tfac]).to_numpy()
+                w = disdat.w.to_numpy()
+                n1 = sum(G + C)
+
+                # Store the current iteration's data
+                current_data = {
+                    'Y': Y,
+                    'post': post,
+                    'G': G,
+                    'group': g,
+                    'time_period': tn
+                }
+                
+                # print(f"Lengths: Y={len(Y)}, post={len(post)}, G={len(G)}, group={type(g)}, time_period={type(tn)}")
+                # results_list.append(pd.DataFrame(current_data))
+
+                #-----------------------------------------------------------------------------
+                # checks to make sure that we have enough observations
+                
+                skip_this_att_gt = False
+
+                if np.sum(G * post) == 0:
+                    print(f"No units in group {g} in time period {t_i + tfac + 1}, e1")
+                    skip_this_att_gt = True 
+
+                if np.sum(G * (1 - post)) == 0:
+                    print(f"No units in group {g} in time period {t_i + 1}, e2")
+                    skip_this_att_gt = True 
+
+                if np.sum(C * post) == 0:
+                    print(f"No available control units for group {g} in time period {t_i + tfac + 1}, e3")
+                    skip_this_att_gt = True 
+
+                if np.sum(C * (1 - post)) == 0:
+                    print(f"No available control units for group {g} in time period {t_i+1}, e4")
+                    skip_this_att_gt = True 
+
+                if skip_this_att_gt:
+                    # Append results with missing ATT and NA influence function
+                    add_att_data(att=np.nan, pst=post_treat, inf_f=np.full(n, np.nan))
+                    #add_att_data()
+                    continue
+
+                # return (inf_func)
+                try:
+                    _, covariates = fml(xformla, data = disdat, return_type = 'dataframe')
+                    covariates = np.array(covariates)
+                except Exception as e:
+                    print(f"Warning: Formula processing failed: {e}")
+                    y_str, x_str = xformla.split("~")
+                    xs1 = x_str.split('+')
+                    xs1_col_names = [x.strip() for x in xs1 if x.strip() != '1']
+                    n_dis = len(disdat)
+                    ones = np.ones((n_dis, 1))
+                    try:
+                        covariates = disdat[xs1_col_names].to_numpy()
+                        covariates = np.append(covariates, ones, axis=1)
+                    except:
+                        covariates = ones
+
+                #-----------------------------------------------------------------------------
+                # code for actually computing att(g,t)
+                #-----------------------------------------------------------------------------
+                # print(Y, post, G, w, covariates)
+                
+                if callable(est_method):
+                    est_att_f = est_method
+                elif est_method == "reg":
+                    est_att_f = reg_did.reg_did_rc
+                elif est_method == "ipw":
+                    est_att_f = ipwd_did.std_ipw_did_rc
+                elif est_method == "dr":
+                    est_att_f = drdid.drdid_rc
+                
+                att_gt, att_inf_func = est_att_f(y=Y, post=post, D = G, i_weights=w, covariates=covariates)
+                # print(att_inf_func)
+                att_inf_func = (n/n1)*att_inf_func
+                
+                inf_func_df = pd.DataFrame(
+                {
+                    "inf_func": att_inf_func,
+                    "right_ids": right_ids
+                }
+                ).fillna(0)
+
+                inf_zeros = np.zeros(n)
+                aggte_infffuc = inf_func_df.groupby('right_ids').inf_func.sum()
+                try:
+                    dis_idx1 = np.isin(data['rowid'].unique(), aggte_infffuc.index.to_numpy())
+                except:
+                    dis_idx1 = np.isin(data['rowid'].unique().to_numpy(), aggte_infffuc.index.to_numpy())
+                
+                inf_zeros[dis_idx1] = np.array(aggte_infffuc)
+
+                add_att_data(att_gt, pst = post_treat, inf_f=inf_zeros)
+                # print(att_est)
+
+    output = {
     'group': group ,
     'year': year,
     "att" : att_est,
     'post ': post_array
-  }
-  return (output, np.array(inf_func))
+    }
+
+    return (output, np.vstack(inf_func))
+
+
+# import numpy as np, pandas as pd
+# import patsy 
+# from drdid import drdid, reg_did, ipwd_did
+
+# from csdid.utils.bmisc import panel2cs2
+
+# fml = patsy.dmatrices
+
+
+# # data Struct
+# # output = {
+# #   "att" : []
+# #   'group': []
+# #   'year': []
+# #   'post ': []
+# # }
+
+
+# def compute_att_gt(dp, est_method = "dr", base_period = 'varying'):
+#   yname = dp['yname']
+#   tname = dp['tname']
+#   idname = dp['idname']
+#   xformla = dp['xformla']
+#   data = dp['data']
+#   weights_name = dp['weights_name']
+#   # base_period = dp['base_period']
+#   panel = dp['panel']
+#   true_rep_cross_section = dp['true_rep_cross_section']
+#   control_group = dp['control_group']
+#   anticipation = dp['anticipation']
+#   gname = dp['gname']
+#   n = dp['n']
+#   nT = dp['nT']
+#   nG = dp['nG']
+#   tlist = dp['tlist']
+#   glist = dp['glist']
+#   tlist_len, tfac = len(tlist), 0
+#   if base_period != 'universal':
+#     tlist_len = tlist_len - 1
+#     tfac = 1
+
+#   inf_func = []
+
+#   att_est, group, year, post_array = [], [], [], []
+
+#   def add_att_data(att = 0, pst = 0, inf_f = []):
+#     inf_func.append(inf_f)
+#     att_est.append(att)
+#     group.append(g)
+#     year.append(tn)
+#     post_array.append(pst)
+
+#   never_treated = control_group == 'nevertreated'
+#   if never_treated:
+#     data = data.assign(C = 1 * (data[gname] == 0))
+#   data = data.assign(y_main = data[yname])
+
+#   # g, t = glist[0], tlist[0]
+
+#   # for _, g, in enumerate(glist):
+#   for g_index, g in enumerate(glist):
+
+#     # g = glist[1]
+#     G_main = (data[gname] == g)
+#     data = data.assign(G_m = 1 * G_main)
+
+#     for t_i in range(tlist_len):
+#       pret = t_i
+#       tn = tlist[t_i + tfac]
+#       if base_period == 'universal' or g < tn:
+#         try:
+#           pret = np.where(tlist + anticipation < g)[0][-1]
+#         except:
+#           raise f"There are no pre-treatment periods for the group first treated at {g}\nUnits from this group are dropped"
+#           # break
+
+
+#       if base_period == 'universal':
+#         if pret == tn:
+#           add_att_data()
+
+#       if not never_treated:
+#         n1 = data[gname] == 0
+#         n2 = (data[gname] > (tlist[np.max([t_i, pret]) + tfac]) + anticipation)
+#         #n3 = np.where(data[gname] != glist[g], True, False)
+#         n3 = np.where(data[gname] != glist[g_index], True, False)
+
+#         row_eval = n1 | n2 & n3
+#         data = data.assign(C = 1 * row_eval)
+
+#       post_treat = 1 * (g <= tn)
+#       disdat = data[(data[tname] == tn) | (data[tname] == tlist[pret])]
+
+#       if panel: 
+#         disdat = panel2cs2(disdat, yname, idname, tname)
+#         disdat = disdat.dropna()
+#         n = len(disdat)
+#         dis_idx = np.array(disdat.G_m == 1) | np.array(disdat.C == 1)
+#         disdat = disdat.loc[dis_idx, :]
+#         n1 = len(disdat)
+#         G = disdat.G_m
+#         C = disdat.C
+#         w = disdat.w
+
+#         ypre = disdat.y0 if tn > pret else disdat.y1
+#         ypost = disdat.y0 if tn < pret else disdat.y1
+#         _, covariates = fml(xformla, data = disdat, return_type = 'dataframe')
+
+#         G, C, w, ypre = map(np.array, [G, C, w, ypre])
+#         ypost, covariates = map(np.array, [ypost, covariates])
+
+#         if callable(est_method):
+#           est_att_f = est_method
+#         elif est_method == "reg":
+#           est_att_f = reg_did.reg_did_panel
+#         elif est_method == "ipw":
+#           est_att_f = ipwd_did.std_ipw_did_panel
+#         elif est_method == "dr":
+#           est_att_f = drdid.drdid_panel
+
+#         att_gt, att_inf_func = est_att_f(ypost, ypre, G, i_weights=w, covariates=covariates)
+
+#         inf_zeros = np.zeros(n)
+#         att_inf = n / n1 * att_inf_func
+#         inf_zeros[dis_idx] = att_inf
+
+#         add_att_data(att_gt, inf_f=inf_zeros)
+
+#       if not panel:
+#         right_ids = np.array(disdat.query('(G_m == 1) or (C == 1)').rowid.to_numpy())
+#         dis_idx = (data['rowid'].isin(right_ids)) &\
+#           ((data[tname] == tlist[t_i + tfac]) |\
+#             (data[tname] == tlist[pret]))
+
+#         disdat = data.loc[dis_idx]
+
+#         G = disdat.G_m.to_numpy()
+#         C = disdat.C.to_numpy()
+#         Y = disdat[yname].to_numpy()
+#         post = 1 * (disdat[tname] == tlist[t_i + tfac]).to_numpy()
+#         w = disdat.w.to_numpy()
+
+#         # G, C, Y, post, w = map(np.array, [G, C, Y, post, w])
+
+
+#         n1 = sum(G + C)
+
+#         skip_this_att_gt = False
+#         if np.sum(G * post) == 0:
+#           print(f"No units in group {g} in time period {tn}")
+#           skip_this_att_gt = True 
+
+#         if np.sum(G * (1 - post)) == 0:
+#           print(f"No units in group {g} in time period {tn}")
+#           skip_this_att_gt = True 
+
+#         if np.sum(C * post) == 0:
+#           print(f"No available control units for group {g} in time period {tn}")
+#           skip_this_att_gt = True 
+
+#         if np.sum(C * (1 - post)) == 0:
+#           print(f"No available control units for group {g} in time period {tn}")
+#           skip_this_att_gt = True 
+
+#         if skip_this_att_gt:
+#           add_att_data()
+
+#         try:
+#           _, covariates = fml(xformla, data = disdat, return_type = 'dataframe')
+#           covariates = np.array(covariates)
+#         except:
+#           y_str, x_str = xformla.split("~")
+#           xs1 = x_str.split('+')
+#           xs1_col_names = [x.strip() for x in xs1 if x.strip() != '1']
+#           n_dis = len(disdat)
+#           ones = np.ones((n_dis, 1))
+#           try:
+#             covariates = disdat[xs1_col_names].to_numpy()
+#             covariates = np.append(covariates, ones, axis=1)
+#           except:
+#             covariates = ones
+
+        
+
+#         if callable(est_method):
+#           est_att_f = est_method
+#         elif est_method == "reg":
+#           est_att_f = reg_did.reg_did_rc
+#         elif est_method == "ipw":
+#           est_att_f = ipwd_did.std_ipw_did_rc
+#         elif est_method == "dr":
+#           est_att_f = drdid.drdid_rc
+#         att_gt, att_inf_func = est_att_f(y=Y, post=post, D = G, i_weights=w, covariates=covariates)
+
+#         inf_func_df = pd.DataFrame(
+#           {
+#             "inf_func": att_inf_func,
+#             "right_ids": right_ids
+#           }
+#         )
+#         inf_zeros = np.zeros(n)
+#         aggte_infffuc = inf_func_df.groupby('right_ids').inf_func.sum()
+#         try:
+#           dis_idx1 = np.isin(data['rowid'].unique(), aggte_infffuc.index.to_numpy())
+#         except:
+#           dis_idx1 = np.isin(data['rowid'].unique().to_numpy(), aggte_infffuc.index.to_numpy())
+#         inf_zeros[dis_idx1] = np.array(aggte_infffuc)
+
+#         add_att_data(att_gt, pst = post_treat, inf_f=inf_zeros)
+#         # print(att_est)
+
+#   output = {
+#     'group': group ,
+#     'year': year,
+#     "att" : att_est,
+#     'post ': post_array
+#   }
+#   return (output, np.array(inf_func))
\ No newline at end of file
diff --git a/csdid/attgt_fnc/preprocess_did.py b/csdid/attgt_fnc/preprocess_did.py
index 979b557..714e2c2 100644
--- a/csdid/attgt_fnc/preprocess_did.py
+++ b/csdid/attgt_fnc/preprocess_did.py
@@ -1,3 +1,248 @@
+# from typing import Union, List, Optional, Dict
+# import pandas as pd, numpy as np
+# import patsy 
+# from csdid.utils.bmisc import makeBalancedPanel
+
+# fml = patsy.dmatrices
+
+# def pre_process_did(yname, tname, idname, gname, data: pd.DataFrame, 
+#   control_group: Union[str, List[str]] = "nevertreated", 
+#   anticipation = 0, xformla : str = None,
+#   panel = True, allow_unbalanced_panel = True, cband = False,
+#   clustervar = None,  weights_name = None
+#   ) -> dict:
+
+#   n, t = data.shape
+#   # control_group = "nevertreated"
+#   columns = [idname, tname, yname, gname]
+#   # print(columns)
+#   # Columns
+#   if clustervar is not None:
+#     columns += [clustervar]
+#   if weights_name is not None:
+#     columns += [weights_name]
+#     w = data[weights_name]
+#   else:
+#     w = np.ones(n)
+
+
+#   if xformla is None:
+#     xformla = f'{yname} ~ 1'
+
+#   # if xformla is None:
+#   try:
+#     _, x_cov = fml(xformla, data = data, return_type='dataframe')
+#     _, n_cov = x_cov.shape
+#     data = pd.concat([data[columns], x_cov], axis=1)
+#     data = data.assign(w = w)
+#   except:
+#     data = data.assign(intercept = 1)
+#     clms = columns + ['intercept']
+#     n_cov = len(data.columns)
+#     # patsy dont work with pyspark
+#     data = data[clms]
+#     if weights_name is None:
+#       data = data.assign(w = 1)
+#     else:
+#       data = data.assign(w = lambda x: x[weights_name] * 1)
+
+
+#   data = data.dropna()
+#   ndiff = n - len(data) 
+#   if ndiff != 0: 
+#     print(f'dropped, {ndiff}, rows from original data due to missing data')
+#   try:
+
+#     tlist = np.sort(data[tname].unique())
+#     glist = np.sort(data[gname].unique())
+#   except:
+#     tlist = np.sort(data[tname].unique().to_numpy())
+#     glist = np.sort(data[gname].unique().to_numpy())
+
+#   asif_nev_treated = data[gname] > np.max(tlist)
+#   asif_nev_treated.fillna(False, inplace=True)
+#   data.loc[asif_nev_treated, gname] = 0
+
+#   if len(glist[glist == 0]) == 0:
+#     if control_group == "nevertreated":
+#       raise ValueError("There is no available never-treated group")
+#     else:
+#       value = np.max(glist) - anticipation
+#       data = data.query(f'{tname} < @value')
+#       tlist = np.sort(data[tname].unique())
+#       glist = np.sort(data[gname].unique())
+#       glist = glist[glist < np.max(glist)]
+
+#   glist = glist[glist > 0]
+#   # first prerios 
+#   fp = tlist[0]
+#   glist = glist[glist > fp + anticipation]
+
+#   treated_fp = (data[gname] <= fp) & ~(data[gname] == 0)
+#   treated_fp.fillna(False, inplace=True)
+
+#   try:
+
+#     nfirst_period = np.sum(treated_fp) if panel \
+#       else len(data.loc[treated_fp, idname].unique())
+#   except:
+#     nfirst_period = treated_fp.sum() if panel \
+#       else len(data.loc[treated_fp, idname].unique())
+
+#   if nfirst_period > 0:
+#     warning_message = f"Dropped {nfirst_period} units that were already treated in the first period."
+#     print(warning_message)
+#     glist_in = np.append(glist, [0])
+#     data = data.query(f'{gname} in @glist_in')
+#     tlist = np.sort(data[tname].unique())
+#     glist = np.sort(data[gname].unique())
+#     glist = glist[glist > 0]
+#     fp = tlist[0]
+#     glist = glist[glist > fp + anticipation]
+
+#   #todo: idname must be numeric
+#   true_rep_cross_section = False
+#   if not panel:
+#     true_rep_cross_section = True
+
+
+
+#   #-----------------------------------------------------------------------------
+#   # setup data in panel case
+#   #-----------------------------------------------------------------------------
+
+#   # # Check if data is a balanced panel if panel = True and allow_unbalanced_panel = True
+#   # if panel and allow_unbalanced_panel:
+#   #     # First, focus on complete cases
+#   #     keepers = data.dropna()
+#   #     data_comp = keepers.copy()
+
+#   #     # Make it a balanced dataset
+#   #     n_all = len(data_comp[idname].unique())
+#   #     data_bal = makeBalancedPanel(data_comp, idname, tname)
+#   #     n_bal = len(data_bal[idname].unique())
+
+#   #     if n_bal < n_all:
+#   #         # If fewer unique IDs in the balanced panel, it means the panel is unbalanced
+#   #         print("You have an unbalanced panel. Proceeding as such.")
+#   #         allow_unbalanced_panel = True
+#   #     else:
+#   #         # If the number of unique IDs remains the same, it means the panel is balanced
+#   #         print("You have a balanced panel. Setting allow_unbalanced_panel = False.")
+#   #         allow_unbalanced_panel = False
+
+
+#   if panel:
+#     if allow_unbalanced_panel: 
+#       panel = False
+#       true_rep_cross_section = False
+#     else:
+#       keepers = data.dropna().index
+#       n = len(data[idname].unique)
+#       print(n)
+#       n_keep = len(data.iloc[keepers, idname].unique())
+
+#       if len(data.loc[keepers] < len(data)):
+#         print(f"Dropped {n-n_keep} observations that had missing data.")
+#         data = data.loc[keepers]
+#       # make balanced data set
+#       n_old = len(data[idname].unique())
+#       data = makeBalancedPanel(data, idname=idname, tname=tname)
+#       n = len(data[idname].unique())
+#       if len(data) == 0:
+#         raise ValueError("All observations dropped to convert data to balanced panel. Consider setting `panel=False` and/or revisit 'idname'.")
+#       if n < n_old:
+#         warnings.warn(f"Dropped {n_old-n} observations while converting to balanced panel.")
+#       tn = tlist[0]
+#       n = len(data.query(f'{tname} == @tn'))
+
+#   # # Add rowid
+#   # if not panel:
+#   #     # Handle missing data
+#   #     keepers = data.dropna().index
+#   #     dropped_count = len(data) - len(keepers)
+#   #     if dropped_count > 0:
+#   #         print(f"Dropped {dropped_count} observations that had missing data.")
+      
+#   #     # Drop incomplete rows
+#   #     data = data.loc[keepers]
+
+#   #     # Check if all rows are dropped
+#   #     if len(data) == 0:
+#   #         raise ValueError("All observations dropped due to missing data problems.")
+
+#   #     # Add rowid based on true_repeated_cross_section
+#   #     if true_rep_cross_section:
+#   #         data = data.assign(rowid=range(len(data)))
+#   #         idname = 'rowid'
+#   #     else:
+#   #         if idname not in data.columns:
+#   #             raise ValueError(f"Column {idname} is not in the dataset.")
+#   #         data = data.assign(rowid=lambda x: x[idname] * 1)
+
+#   #     # Number of unique observations
+#   #     n = data[idname].nunique()
+
+#   if not panel:
+#     # Check for complete cases (equivalent to complete.cases in R)
+#     keepers = data.notna().all(axis=1)
+#     dropped_count = len(data) - keepers.sum()
+    
+#     if dropped_count > 0:
+#       import warnings
+#       warnings.warn(f"Dropped {dropped_count} observations that had missing data.")
+#       data = data[keepers].copy()
+    
+#     # If drop all data, raise error
+#     if len(data) == 0:
+#       raise ValueError("All observations dropped due to missing data problems.")
+    
+#     # Add rowid column
+#     if true_rep_cross_section:
+#       data.loc[:, 'rowid'] = range(1, len(data) + 1)  # 1-based indexing like R
+#       idname = 'rowid'
+#     else:
+#       # Set rowid to idname for repeated cross section/unbalanced
+#       if idname not in data.columns:
+#         raise ValueError(f"Column {idname} not found in dataset")
+#       data.loc[:, 'rowid'] = data[idname]
+    
+#     # Calculate n as unique number of cross section observations
+#     n = data[idname].nunique()
+
+#   data = data.sort_values([idname, tname])
+#   data = data.assign(w1 = lambda x: x['w'] * 1)
+#   # data.loc[:, ".w"] = data['w']
+#   if len(glist) == 0:
+#     raise f"No valid groups. The variable in '{gname}' should be expressed as the time a unit is first treated (0 if never-treated)."
+#   if len(tlist) == 2:
+#     cband = False
+#   gsize = data.groupby(data[gname]).size().reset_index(name="count")
+#   gsize["count"] /= len(tlist)
+
+#   reqsize = n_cov + 5
+#   gsize = gsize[gsize["count"] < reqsize]
+
+#   if len(gsize) > 0:
+#     gpaste = ",".join(map(str, gsize[gname]))
+#     warnings.warn(f"Be aware that there are some small groups in your dataset.\n  Check groups: {gpaste}.")
+
+#     if 0 in gsize[gname].to_numpy() and control_group == "nevertreated":
+#       raise "Never-treated group is too small, try setting control_group='notyettreated'."
+#   nT, nG = map(len, [tlist, glist])
+#   did_params = {
+#     'yname' : yname, 'tname': tname,
+#     'idname' : idname, 'gname': gname,
+#     'xformla': xformla, 'data': data,
+#     'tlist': tlist, 'glist': glist,
+#     'n': n, 'nG': nG, 'nT': nT,
+#     'control_group': control_group, 'anticipation': anticipation,
+#     'weights_name': weights_name, 'panel': panel,
+#     'true_rep_cross_section': true_rep_cross_section,
+#     'clustervars': clustervar
+#   }
+#   return did_params
+
 
 import pandas as pd, numpy as np
 import patsy 
@@ -16,6 +261,7 @@ def pre_process_did(yname, tname, idname, gname, data: pd.DataFrame,
   n, t = data.shape
   control_group = control_group[0]
   columns = [idname, tname, yname, gname]
+  control_group = "nevertreated"
   # print(columns)
   # Columns
   if clustervar is not None:
@@ -181,4 +427,3 @@ def pre_process_did(yname, tname, idname, gname, data: pd.DataFrame,
     'clustervars': clustervar
   }
   return did_params
-


Test patch:
diff --git a/test_csdid_aq.ipynb b/archive/test_csdid_aq.ipynb
similarity index 100%
rename from test_csdid_aq.ipynb
rename to archive/test_csdid_aq.ipynb