Create tsne.py

tsne.py

@@ -0,0 +1,259 @@
+import numpy as np
+import pandas as pd
+from openTSNE import TSNE
+import plotly.graph_objs as go
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+from sklearn.decomposition import PCA
+from scipy.optimize import linear_sum_assignment
+
+class TSNE_Plot():
+    def __init__(self, sentence, embed, label = None, n_clusters :int = 3, n_annotation_positions:int = 20):
+        assert n_clusters > 0, "N must be greater than 0" 
+        self.N = n_clusters
+        self.test_X = pd.DataFrame({"text": sentence, "embed": [np.array(i) for i in embed]})
+        self.test_y = pd.DataFrame({'label':label}) if label is not None else pd.DataFrame({"label": self.cluster()})
+        self.embed = self.calculate_tsne()
+        self.init_df()
+
+        self.n_annotation_positions = n_annotation_positions
+        self.show_sentence = []
+        self.random_sentence()
+
+        
+        self.annotation_positions = []
+        self.get_annotation_positions()
+        self.mapping = {}
+    
+    def cluster(self):
+        from sklearn.cluster import KMeans
+        n_components = min(50, len(self.test_X))
+        pca = PCA(n_components=n_components)
+        compact_embedding = pca.fit_transform(np.array(self.test_X["embed"].tolist()))
+        kmeans = KMeans(n_clusters=self.N)
+        kmeans.fit(compact_embedding)
+        labels = kmeans.labels_
+        return labels
+    
+    def generate_colormap(self, n_labels):
+        #创建一个均匀分布的颜色映射
+        color_norm = mcolors.Normalize(vmin=0, vmax=len(n_labels) - 1)
+        # 使用 plt.cm 中预先定义的colormap，你可以自由选择其他colormap如"hsv", "hot", "cool", "viridis"等
+        scalar_map = plt.cm.ScalarMappable(norm=color_norm, cmap='jet') 
+
+        colormap = {}
+        for label in range(len(n_labels)):
+            # 将颜色值转换为十六进制
+            color_hex = mcolors.to_hex(scalar_map.to_rgba(label))
+            colormap[n_labels[label]] = color_hex
+        return colormap
+
+    def divide_hex_color_by_half(self, hex_color):
+        if len(hex_color) > 0 and hex_color[0] == "#":
+            hex_color = hex_color[1:]
+
+        red_hex, green_hex, blue_hex = hex_color[0:2], hex_color[2:4], hex_color[4:6]
+        
+        red_half = int(red_hex, 16) // 10 + (255-25)
+        green_half = int(green_hex, 16) // 10 + (255-25)
+        blue_half = int(blue_hex, 16) // 10 + (255-25)
+        
+        half_hex_color = "#{:02x}{:02x}{:02x}".format(red_half, green_half, blue_half)
+        return half_hex_color
+
+
+    def get_annotation_positions(self):
+        min_x, max_x = self.df['x'].min()-1, self.df['x'].max()+2
+        n = self.n_annotation_positions
+
+        y_min, y_max = self.df['y'].min() * 3, self.df['y'].max() * 3
+
+        add = 0 if n % 2 == 0 else 1
+        y_values = np.linspace(y_min, y_max, n//2+add)
+
+        left_positions = [(min_x, y) for y in y_values]
+        right_positions = [(max_x, y) for y in y_values]
+
+
+        self.annotation_positions = left_positions + right_positions
+
+    
+    def euclidean_distance(self, p1, p2):
+        return np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
+
+    def map_points(self):
+        # Get points from the dataframe using the show_sentence indices
+        points1 = [(self.embed[i][0], self.embed[i][1]) for i in self.show_sentence]
+
+        # Create a distance matrix between the points
+        distance_matrix = np.zeros((len(points1), len(self.annotation_positions)))
+
+        for i, point1 in enumerate(points1):
+            for j, point2 in enumerate(self.annotation_positions):
+                distance_matrix[i, j] = self.euclidean_distance(point1, point2)
+
+        # Apply linear_sum_assignment to find the optimal mapping
+        row_ind, col_ind = linear_sum_assignment(distance_matrix)
+
+        for i, j in zip(row_ind, col_ind):
+            self.mapping[self.show_sentence[i]] = self.annotation_positions[j]
+
+
+    def show_text(self, show_sentence, text):
+        sentence = []
+        for i in range(len(text)):
+            if i in show_sentence:
+                s = text[i][:10] + "..." + text[i][-10:]
+                sentence.append(s)
+            else:
+                sentence.append("")
+        return sentence
+
+    def init_df(self):
+        X, Y = np.split(self.embed, 2, axis=1)
+        data = {
+            "x": X.flatten(),
+            "y": Y.flatten(),
+        }
+
+        self.df = pd.DataFrame(data)
+
+
+    def format_data(self):
+        sentence = self.show_text(self.show_sentence, self.test_X["text"])
+        X, Y = np.split(self.embed, 2, axis=1)
+        n = len(self.test_X)
+        data = {
+            "x": X.flatten(),
+            "y": Y.flatten(),
+            "label": self.test_y["label"],
+            "sentence" : sentence,
+            "size" : [20 if i in self.show_sentence else 10 for i in range(n)],
+            "pos" : [{"x_offset": self.mapping.get(i, (0, 0))[0], "y_offset": self.mapping.get(i, (0, 0))[1]} for i in range(n)],
+            "annotate" : [True if i in self.show_sentence else False for i in range(n)],
+        }
+        self.df = pd.DataFrame(data)
+
+    def calculate_tsne(self):   
+        embed = np.array(self.test_X["embed"].tolist())
+        n_components = min(50, len(self.test_X))
+        pca = PCA(n_components=n_components)
+        compact_embedding = pca.fit_transform(embed)
+        tsne = TSNE(
+            perplexity=30,
+            metric="cosine",
+            n_jobs=8,
+            random_state=42,
+            verbose=False,
+        )
+        embedding_train = tsne.fit(compact_embedding)
+        embedding_train = embedding_train.optimize(n_iter=1000, momentum=0.8)
+        return embedding_train
+
+    def random_sentence(self):
+        #多次随机可能会影响可视化结果
+        n_samples = len(self.test_y)
+
+        show_sentence = []
+        while len(show_sentence) < self.n_annotation_positions:
+            show_sentence.append(np.random.randint(0, n_samples))
+            show_sentence = list(set(show_sentence))
+
+        # 确保每个标签至少有一个句子，用在show_sentence中最多的标签的句子来补充
+        label_count = self.test_y["label"].value_counts()
+        max_label = label_count.index[0]
+        max_count = label_count[0]
+        for i in range(max_count):
+            for j in range(len(label_count)):
+                if label_count[j] == i:
+                    show_sentence.append(self.test_y[self.test_y["label"] == label_count.index[j]].index[0])
+        self.show_sentence = list(set(show_sentence))
+
+    def plot(self, return_fig=False):
+        min_x, max_x = self.df['x'].min()-1, self.df['x'].max()+2
+        fig = go.Figure()
+        fig = go.Figure(layout=go.Layout(
+            autosize=False,  # 禁止图像自动调整大小
+            height=800,  # 您可以根据需要调整这个值
+            width=1500,  # 您可以根据需要调整这个值
+            # plot_bgcolor="#262626",
+        ))
+        
+        label_colors = self.generate_colormap(self.df['label'].unique())
+
+        line_legend_group = "lines"
+
+        # 为每个类别的点创建散点图
+        for label, color in label_colors.items():
+            mask = self.df["label"] == label
+            fig.add_trace(go.Scatter(x=self.df[mask]['x'], y=self.df[mask]['y'], mode='markers', 
+                                    marker=dict(color=color, size=self.df[mask]['size']),  # 添加 size 参数
+                                    showlegend=True, legendgroup=line_legend_group,
+                                    name = str(label)) 
+                                    )
+
+
+        # 为每个句子创建注释
+        for x, y, label, sentence, pos, annotate in zip(self.df.x, self.df.y, self.df.label, self.df.sentence, self.df.pos, self.df.annotate):
+            if not sentence:
+                continue
+            if not annotate:
+                continue
+            # pos在左边
+            criteria = (pos["x_offset"] - min_x) < 1e-2
+
+            sentence_annotation = dict(
+                x=pos["x_offset"],
+                y=pos["y_offset"],
+                xref="x",
+                yref="y",
+                text=sentence,
+                showarrow=False,
+                xanchor="right" if criteria else 'left',
+                yanchor='middle',
+                font=dict(color="black"),
+                bordercolor=label_colors.get(label, "black"),
+                borderpad=2,
+                bgcolor=self.divide_hex_color_by_half(label_colors.get(label, "black"))
+            )
+            fig.add_annotation(sentence_annotation)
+
+            x_start = x - 1 if criteria else x + 1
+            x_turn = x - 0.5 if criteria else x + 0.5
+            y_turn = y
+
+            fig.add_trace(go.Scatter(x=[pos["x_offset"], x_start, x_turn, x], y=[pos["y_offset"], pos["y_offset"], y_turn, y], mode='lines', 
+                                    line=dict(color=label_colors.get(label, "black")), showlegend=False, legendgroup=line_legend_group))
+
+        # 取消坐标轴的数字
+        fig.update_xaxes(tickvals=[])
+        fig.update_yaxes(tickvals=[])
+
+        if not return_fig:
+            fig.show()
+        else:
+            return fig
+    
+    def tsne_plot(self, n_sentence = 20, return_fig=False):
+        # 计算t-SNE，返回降维后的数据，每个元素为一个二维向量
+        embedding_train = self.calculate_tsne()
+
+        # 随机抽取显示文本, n为抽取的数量，show_sentence为一个列表，每个元素为显示文本的索引
+        if self.n_annotation_positions != min(n_sentence, len(self.test_y)):
+            self.n_annotation_positions = min(n_sentence, len(self.test_y))
+            self.random_sentence()
+            self.get_annotation_positions()
+
+        # find the optimal sentence positions
+        self.map_points()
+
+        # 格式化数据，输出为一个pandas的DataFrame，包含x, y, label, sentence, sentence_pos, size
+        # x, y为降维后的坐标，label为类别，sentence为显示的文本，sentence_pos为文本的位置("left", "right")，size为被选中文本的大小
+        self.format_data()
+        # self.df = self.df.sort_values('y').reset_index(drop=True)
+
+        if not return_fig:
+            # 绘制图像
+            self.plot()
+        else:
+            return self.plot(return_fig=return_fig)