From 52cb1407463d195cef2550448532e877b28191df Mon Sep 17 00:00:00 2001 From: bastien ollier Date: Wed, 19 Jun 2024 10:03:48 +0200 Subject: [PATCH] add 3d to kmeans --- frontend/pages/clustering:_kmeans.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/frontend/pages/clustering:_kmeans.py b/frontend/pages/clustering:_kmeans.py index 69d9920..c61bf40 100644 --- a/frontend/pages/clustering:_kmeans.py +++ b/frontend/pages/clustering:_kmeans.py @@ -11,7 +11,7 @@ if "data" in st.session_state: with st.form("my_form"): row1 = st.columns([1,1,1]) n_clusters = row1[0].selectbox("Number of clusters", range(1,data.shape[0])) - data_name = row1[1].multiselect("Data Name",data.select_dtypes(include="number").columns, max_selections=2) + data_name = row1[1].multiselect("Data Name",data.select_dtypes(include="number").columns, max_selections=3) n_init = row1[2].number_input("n_init",step=1,min_value=1) row2 = st.columns([1,1]) @@ -20,16 +20,24 @@ if "data" in st.session_state: st.form_submit_button("launch") - if len(data_name) == 2: + if len(data_name) >= 2 and len(data_name) <=3: x = data[data_name].to_numpy() kmeans = KMeans(n_clusters=n_clusters, init="random", n_init=n_init, max_iter=max_iter, random_state=111) y_kmeans = kmeans.fit_predict(x) - fig, ax = plt.subplots(figsize=(12,8)) - plt.scatter(x[:, 0], x[:, 1], c=y_kmeans, s=50, cmap="viridis") - centers = kmeans.cluster_centers_ - plt.scatter(centers[:, 0], centers[:, 1], c="black", s=200, marker="X") + fig = plt.figure() + if len(data_name) == 2: + ax = fig.add_subplot(projection='rectilinear') + plt.scatter(x[:, 0], x[:, 1], c=y_kmeans, s=50, cmap="viridis") + centers = kmeans.cluster_centers_ + plt.scatter(centers[:, 0], centers[:, 1], c="black", s=200, marker="X") + else: + ax = fig.add_subplot(projection='3d') + + ax.scatter(x[:, 0], x[:, 1],x[:, 2], c=y_kmeans, s=50, cmap="viridis") + centers = kmeans.cluster_centers_ + ax.scatter(centers[:, 0], centers[:, 1],centers[:, 2], c="black", s=200, marker="X") st.pyplot(fig) else: