Created a method to downsample for unbalanced data (for binary classification)

I modified the code in the undersampling chapter of this page to make it work.

python

#Method to ensure that the number of positive and negative data in train data is the same
#Reference: https://qiita.com/ryouta0506/items/619d9ac0d80f8c0aed92
#Cluster with kmeans and sample at a constant rate for each cluster

# X :pandas DataFrame
# target_column_name :name of the class. "Onset flag" etc.
# minority_label :Decimal label value. For example, "1".

def under_sampling(X, target_column_name, minority_label):
    
    #Hide it because it appears every time
    import warnings
    warnings.simplefilter('ignore', pd.core.common.SettingWithCopyWarning)
    
    #Divide into majority and minority
    X_majority = X.query(f'{target_column_name} != {minority_label}')
    X_minority = X.query(f'{target_column_name} == {minority_label}')

    #Clustering with KMeans
    from sklearn.cluster import KMeans
    km = KMeans(random_state=43)
    km.fit(X_majority)
    X_majority['Cluster'] = km.predict(X_majority)

    #Calculate how many samples to extract for each cluster
    ratio = X_majority['Cluster'].value_counts() / X_majority.shape[0] 
    n_sample_ary = (ratio * X_minority.shape[0]).astype('int64').sort_index()
    
    #Extract samples for each cluster
    dfs = []
    for i, n_sample in enumerate(n_sample_ary):
        dfs.append(X_majority.query(f'Cluster == {i}').sample(n_sample))
    
    #Make sure to combine minority data as well
    dfs.append(X_minority)
    
    #Create data after undersampling
    X_new = pd.concat(dfs, sort=True)
    
    #Deleted because it is unnecessary
    X_new = X_new.drop('Cluster', axis=1)
    
    return X_new