Drawing.py 文件源码

def plotBasisFunctions(self, eigenvalues, eigenvectors):
        '''3d plot of the basis function. Right now I am plotting eigenvectors,
           so each coordinate of the eigenvector correspond to the value to be
           plotted for the correspondent state.''' 
        for i in xrange(len(eigenvalues)):  
            fig, ax = plt.subplots(subplot_kw = dict(projection = '3d'))
            X, Y = np.meshgrid(np.arange(self.numRows), np.arange(self.numCols))
            Z = eigenvectors[:,i].reshape(self.numCols, self.numRows)

            for ii in xrange(len(X)):
                for j in xrange(len(X[ii])/2):
                    tmp = X[ii][j]
                    X[ii][j] = X[ii][len(X[ii]) - j - 1]
                    X[ii][len(X[ii]) - j - 1] = tmp

            my_col = cm.jet(np.random.rand(Z.shape[0],Z.shape[1]))

            ax.plot_surface(X, Y, Z, rstride = 1, cstride = 1,
                cmap = plt.get_cmap('jet'))
            plt.gca().view_init(elev=30, azim=30)
            plt.savefig(self.outputPath + str(i) + '_eig' + '.png')
            plt.close()


        plt.plot(eigenvalues, 'o')
        plt.savefig(self.outputPath + 'eigenvalues.png')