tsne.py

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)