correlation.py

def calculate_correlation(s):
    s_temp = s[['A1','A3']]
    correlation = s_temp.corr().iloc[1,0]
    print("******************************")
    print(f'Correlation between A1 & A3: {correlation}')
    print("******************************")
    # if correlation > 0.6 or correlation < 0.6 remove A3
    if correlation > 0.6 or correlation < -0.6:
        s = s.drop(['A3'], axis=1)
        print(f'A3 was removed {correlation} > 0.6 or {correlation} < -0.6')
        print("******************************")
    
    return 

def df_to_array(df):
    A1 = df[['A1']].to_numpy()
    A2 = df[['A2']].to_numpy()
    if 'A3' in df:
        A3 = df[['A3']].to_numpy()

    df_as_matrix = np.array(A1,A2)
    if 'A3' in s:
        df_as_matrix = np.concatenate(df_as_matrix,A3)


    return df_as_matrix