rankd.py 文件源码

def draw(self, **kwargs):
        """
        Draws the heatmap of the ranking matrix of variables.
        """
        # Set the axes aspect to be equal
        self.ax.set_aspect("equal")

        # Generate a mask for the upper triangle
        mask = np.zeros_like(self.ranks_, dtype=np.bool)
        mask[np.triu_indices_from(mask)] = True

        # Draw the heatmap
        # TODO: Move mesh to a property so the colorbar can be finalized
        data = np.ma.masked_where(mask, self.ranks_)
        mesh = self.ax.pcolormesh(data, cmap=self.colormap, vmin=-1, vmax=1)

        # Set the Axis limits
        self.ax.set(
            xlim=(0, data.shape[1]), ylim=(0, data.shape[0])
        )

        # Add the colorbar
        cb = self.ax.figure.colorbar(mesh, None, self.ax)
        cb.outline.set_linewidth(0)

        # Reverse the rows to get the lower left triangle
        self.ax.invert_yaxis()

        # Add ticks and tick labels
        self.ax.set_xticks(np.arange(len(self.ranks_)) + 0.5)
        self.ax.set_yticks(np.arange(len(self.ranks_)) + 0.5)
        if self.show_feature_names_:
            self.ax.set_xticklabels(self.features_, rotation=90)
            self.ax.set_yticklabels(self.features_)
        else:
            self.ax.set_xticklabels([])
            self.ax.set_yticklabels([])