main_sample.py

import argparse
from utils import *


def parse_args():
    parser = argparse.ArgumentParser(description="Calculate the ATE under sample-level privacy.")
    parser.add_argument('--path', nargs='?', default='./dataset/',
                        help='Data path.')
    parser.add_argument('--dataset', nargs='?', default='IHDP',
                        help='Choose a dataset.')
    
    # parser.add_argument('--epsilon', type=float, default=5,
    #                     help='Privacy budget.')
    # parser.add_argument('--seed_n', type=int, default=1,
    #                     help='Choose a seed.')
    
    
    parser.add_argument('--N', type=int, default=5,
                        help='The number of neighbors.')
    parser.add_argument('--h', type=float, default=0.001,
                        help='Error coefficient.')
    
    # parser.add_argument('--model_r', type=float, default=0.5,
    #                     help='The privacy budget ratio for the regression model.')
    
    parser.add_argument('--eps_1', type=float, default=0.1,
                        help='The privacy budget ratio for the first phase.')
    parser.add_argument('--eps_2', type=float, default=0.7,
                        help='The privacy budget ratio for the second phase.')
    
    parser.add_argument('--save_csv', type=bool, default=False,
                        help='Whether to save the results.')
    
    return parser.parse_args()




if __name__ == '__main__':


    args = parse_args()
    
    root_path = args.path 
    dataset_name = args.dataset

    epsilon = args.epsilon
    
    seed = args.seed
    n_neighbors = args.N
    h = args.h
    

    e1_r = args.eps_1
    e2_r = args.eps_2
    save_csv = args.save_csv
    

    # epsilon = float(epsilon)
    # seed = int(seed)
    n_neighbors = int(n_neighbors)
    h = float(h)
    save_csv = bool(save_csv)

    # model_r = args.model_r
    # model_r = float(model_r)
    model_r = 0.5
    e1_r = float(e1_r)
    e2_r = float(e2_r)
    e3_r = 1 - e1_r - e2_r

    e11_r = e1_r * model_r
    e12_r = e1_r - e11_r

    # eps11 = epsilon * e1_r * e11_r
    # eps12 = epsilon * e1_r * e12_r
    # eps2 = epsilon * e2_r
    # eps3 = epsilon * e3_r 

    print("All params: %s " % (args))


    data, cov_ind, cov_len, B = load_dataset(root_path,dataset_name)

    covariates = data[[f'x{i}' for i in range(cov_ind, cov_ind+cov_len)]].values

    data_len = len(data)

    reg_C = 1.0

    model = LogisticRegression(penalty='l2', C=reg_C, solver='lbfgs')
    model.fit(covariates, data['treatment'])
    w_star = model.coef_.flatten()  
    sensitivity_w = 2 * len(w_star) / (data_len * reg_C)

    data['propensity_score'] = model.predict_proba(covariates)[:, 1]
    propensity_score = np.array(data['propensity_score']).reshape(data_len,1)

    distance_matrix_true = cal_dis_mat(covariates)

    treat_ind_true = data[data['treatment']==1].index
    control_ind_true = data[data['treatment']==0].index

    len_treat = len(treat_ind_true)
    len_control = len(control_ind_true)

    df_matched, control_match_count, treated_match_count = psm_match(data,n_neighbors)

    max_control_matches = np.max(control_match_count)
    max_treated_matches = np.max(treated_match_count)

    ate_true = (df_matched['Y1'] - df_matched['Y0']).mean()

    print(f'sensitivity_w:{sensitivity_w}')
    print(f'max_control_matches:{max_control_matches}, max_treated_matches:{max_treated_matches}')
    print('ATE true:{:.5f}'.format(ate_true))

    res_cols = ['eps_total', 'dataset', 'eps11_r', 'eps12_r', 'eps2_r', 'eps3_r', 'k1', 'k2', 'h', 'n_neighbors', 'seed', 'ate_dp', 'ate_re']

    all_res = pd.DataFrame(columns=res_cols)    

    all_epsilon = np.arange(8) + 1
    all_seeds = np.arange(10)

    for ei in range(len(all_epsilon)):
        epsilon = all_epsilon[ei]
        for si in range(len(all_seeds)):
            seed = all_seeds[si]
            np.random.seed(seed)
            random.seed(seed)
            t1 = time.time()

            eps11 = epsilon * e11_r
            eps12 = epsilon * e12_r
            eps2 = epsilon * e2_r
            eps3 = epsilon * e3_r

            # eps11: add noise to the true weight
            noise_w = np.random.laplace(0, sensitivity_w / eps11, size=w_star.shape)
            w_dp = w_star + noise_w 
            
            # eps12: add noise to the e(X)
            logits = np.dot(covariates, w_dp) 
            propensity_scores = 1 / (1 + np.exp(-logits))  # sigmoid 
            
            noise_p = np.random.laplace(0, 1 / eps12, size=propensity_scores.shape)
            propensity_scores_dp = np.clip(propensity_scores + noise_p, 0, 1)
            
            data['propensity_score_dp'] = propensity_scores_dp
            
            # eps2: add noise to the treated/control
            data['treatment_dp'] = data['treatment'].apply(lambda t: random_response(t, eps2))

            len_treat = len(data[data['treatment_dp'] == 1])
            len_control = len(data[data['treatment_dp'] == 0])
            
            treat_ind = data[data['treatment_dp']==1].index
            control_ind = data[data['treatment_dp']==0].index

            
            sel_cov_noisy = data['propensity_score_dp'].values.reshape(-1,1)

            sel_dis_mat = cal_dis_mat(sel_cov_noisy)

            M_noise = cal_M1(sel_dis_mat,treat_ind,control_ind,n_neighbors)
            k1, k2 = cal_k_sample(eps3,len_treat,len_control,M_noise,h)
            
            ate_dp = cal_ate(data,sel_dis_mat,eps3,B,n_neighbors,k1,k2)
            ate_re = np.abs((ate_dp-ate_true) / ate_true)
            
            print(f'[Time:{t2-t1:.3f}s] Sample-DP (epsilon:{epsilon}, e11_r:{e11_r}, e12_r:{e12_r}, e2_r:{e2_r}, e3_r:{e3_r}, seed:{seed}): ate_dp={ate_dp:4f}, ate_re={ate_re:.4f}')

            res1 = [epsilon, dataset_name, e11_r, e12_r, e2_r, e3_r, k1, k2, h, n_neighbors, seed, ate_dp, ate_re]

            all_res.loc[len(all_res)] = res1
            
    group_cols = ['eps_total', 'dataset', 'eps11_r', 'eps12_r', 'eps2_r', 'eps3_r', 'k1', 'k2', 'h', 'n_neighbors']

    grouped_res = all_res.groupby(group_cols).agg(
        ate_dp_mean=('ate_dp', 'mean'),
        ate_dp_std=('ate_dp', 'std'),
        ate_re_mean=('ate_re', 'mean'),
        ate_re_std=('ate_re', 'std')
        ).reset_index()
    
    print('*******************Final results (sample-level)*******************')
    print(grouped_res)

    save_root = './result/'

    if not os.path.exists(save_root):
        os.mkdir(save_root)
  
    # save_csv = True

    all_csv_name = save_root + 'sample_DP_dataset_%s_e11_%s_e12_%s_e2_%s_neigh_%s_h_%s_all_res.csv' %(dataset_name,e11_r,e12_r,e2_r,n_neighbors,h)
    group_csv_name = save_root + 'sample_DP_dataset_%s_e11_%s_e12_%s_e2_%s_neigh_%s_h_%s_group_res.csv' %(dataset_name,e11_r,e12_r,e2_r,n_neighbors,h)


    if save_csv:
        all_res.to_csv(all_csv_name,index=False)
        grouped_res.to_csv(group_csv_name,index=False)
        print('Save %s and %s done.'%(all_csv_name,group_csv_name))