lda_visualizer.py 文件源码

def visualize_document_topics_heatmap(self, outfile, set_topics=False):
        self.sort_doctopics_groups()
        doctopics_raw_hm = numpy.rot90(self.document_topics_raw)
        rows, columns = doctopics_raw_hm.shape
        rownames = self.topic_labels
        columnnames = self.document_names
        pyplot.pcolor(doctopics_raw_hm, norm=None, cmap='Blues')
        pyplot.gca().invert_yaxis()
        if self.group_names:
            ticks_groups = []
            bounds = []
            current_group = False
            start = 0
            for i,doc in enumerate(self.document_names):
                group = self.document_group_dict[doc]
                if group != current_group:
                    if i != 0:
                        bounds.append(i-1)
                        ticks_groups[start+int((i-start)/2)] = current_group
                    current_group = group
                    start=i
                ticks_groups.append('')
            ticks_groups[start+int((i-start)/2)] = current_group
            pyplot.xticks(numpy.arange(columns)+0.5,ticks_groups, fontsize=11)
            if set_topics:
                for index in set_topics:
                    pyplot.axhline(y=index)
                topic_names = self.return_topic_names(set_topics)
                pyplot.yticks(set_topics,topic_names,fontsize=8)
            else:
                pyplot.yticks(numpy.arange(rows)+0.5, rownames, fontsize=8)
            for bound in bounds:
                pyplot.axvline(x=bound)
        pyplot.colorbar(cmap='Blues')
        pyplot.savefig(outfile)
        pyplot.clf()