Apply kmeans on in each group in pandas DataFrame and save the clusters in a new column in the same DataFrame

user_5711976 0
Viewed 555

I have a dataframe containing some embeddings in column D. I would like to first groupby the data by column A and then apply kmeans on each group. Each group might contain nan values, so in the apply function I consider number of clusters as the number of non-nan values in column D devided by 2 (n_clusters = int(not_na_mask.sum()/2)). In the apply function I return df['cluster'].values.tolist(). I printed this values and it's correct for each group, but after running the whole script df_test['clusters'] only contains nan in all the rows.

Sample DataFrame:

df_test = pd.DataFrame({'A' : ['aa', 'bb', 'aa', 'bb','aa', 'bb', 'aa', 'cc', 'aa', 'aa', 'bb', 'bb', 'bb','cc', 'bb', 'aa', 'cc', 'aa'],
                       'B' : [1, 2, np.nan, 4, 6, np.nan, 7, 8, np.nan, 1, 4, 3, 4, 7, 5, 7, 9, np.nan],
                       'D' : [[2, 0, 1, 5, 4, 0], np.nan, [4, 7, 0, 1, 0, 2], [1., 1, 1, 2, 0, 5], np.nan , [1, 6, 3, 2, 1, 9], [4, 2, 1, 0, 0, 0], [3, 5, 6, 8, 8, 0], np.nan,
                             np.nan, [2, 5, 1, 7, 4, 0] , [4, 2, 0, 4, 0, 0], [1., 0, 1, 8, 0, 9], [1, 0, 7, 2, 1, 0], np.nan , [1, 1, 5, 0, 8, 0], [4, 1, 6, 1, 1, 0], np.nan]})

df_test:
    A   B   D
0   aa  1.0 [2, 0, 1, 5, 4, 0]
1   bb  2.0 NaN
2   aa  NaN [4, 7, 0, 1, 0, 2]
3   bb  4.0 [1.0, 1, 1, 2, 0, 5]
4   aa  6.0 NaN
5   bb  NaN [1, 6, 3, 2, 1, 9]
6   aa  7.0 [4, 2, 1, 0, 0, 0]
7   cc  8.0 [3, 5, 6, 8, 8, 0]
8   aa  NaN NaN
9   aa  1.0 NaN
10  bb  4.0 [2, 5, 1, 7, 4, 0]
11  bb  3.0 [4, 2, 0, 4, 0, 0]
12  bb  4.0 [1.0, 0, 1, 8, 0, 9]
13  cc  7.0 [1, 0, 7, 2, 1, 0]
14  bb  5.0 NaN
15  aa  7.0 [1, 1, 5, 0, 8, 0]
16  cc  9.0 [4, 1, 6, 1, 1, 0]
17  aa  NaN NaN

My approach for calculating kmeans:

def apply_kmeans_on_each_category(df):
    
    not_na_mask = df['D'].notna()
    
    embedding = df[not_na_mask]['D']
    n_clusters = int(not_na_mask.sum()/2)
    
    if n_clusters > 1:
        df['cluster'] = np.nan
        kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(embedding.tolist())
        df.loc[not_na_mask, 'cluster'] = kmeans.labels_
        return df['cluster'].values.tolist()
    else:
        return [np.nan] * len(df)

df_test['clusters'] = df_test.groupby('A').apply(apply_kmeans_on_each_category)

result:

df_test['clusters']:
0     NaN
1     NaN
2     NaN
3     NaN
4     NaN
5     NaN
6     NaN
7     NaN
8     NaN
9     NaN
10    NaN
11    NaN
12    NaN
13    NaN
14    NaN
15    NaN
16    NaN
17    NaN
Name: clusters, dtype: object
pandas k-means cluster-analysis embedding group-by
1 Answers

Made some slight changes. Meat of the change is to use transform instead of apply. Also, no need to pass the entire Grouper df, you can just pass column D directly as that's the only column you are using -

def apply_kmeans_on_each_category(df):
    not_na_mask = df.notna()
    
    embedding = df.loc[not_na_mask]
    n_clusters = int(not_na_mask.sum()/2)

    op = pd.Series([np.nan] * len(df), index=df.index)
    if n_clusters > 1:
        df['cluster'] = np.nan
        kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(embedding.tolist())
        op.loc[not_na_mask] = kmeans.labels_.tolist()
    return op
df_test['clusters'] = df_test.groupby('A')['D'].transform(apply_kmeans_on_each_category)

Output

    A   B   D   clusters
0   aa  1.0 [2, 0, 1, 5, 4, 0]  0.0
1   bb  2.0 NaN NaN
2   aa  NaN [4, 7, 0, 1, 0, 2]  1.0
3   bb  4.0 [1.0, 1, 1, 2, 0, 5]    0.0
4   aa  6.0 NaN NaN
5   bb  NaN [1, 6, 3, 2, 1, 9]  0.0
6   aa  7.0 [4, 2, 1, 0, 0, 0]  1.0
7   cc  8.0 [3, 5, 6, 8, 8, 0]  NaN
8   aa  NaN NaN NaN
9   aa  1.0 NaN NaN
10  bb  4.0 [2, 5, 1, 7, 4, 0]  1.0
11  bb  3.0 [4, 2, 0, 4, 0, 0]  1.0
12  bb  4.0 [1.0, 0, 1, 8, 0, 9]    0.0
13  cc  7.0 [1, 0, 7, 2, 1, 0]  NaN
14  bb  5.0 NaN NaN
15  aa  7.0 [1, 1, 5, 0, 8, 0]  0.0
16  cc  9.0 [4, 1, 6, 1, 1, 0]  NaN
17  aa  NaN NaN NaN
user_40980130
Related