drawfigures_recursively_leeshixin.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import scipy.ndimage
# import matplotlib.mlab as mlab

def plot_fun(filename, strech):
    # x, y, z = np.loadtxt(filename,
    #                      dtype={'names': ('Type', 'Num1', 'Atom', 'Num2', 'x', 'y', 'z'),
    #                             'formats': ('|S4', np.float, '|S1', np.float, np.float, np.float, np.float)},
    #                      delimiter=' ',
    #                      skiprows=1,
    #                      usecols=(4, 5, 6),
    #                      unpack=True)

    x, y = np.genfromtxt(filename, autostrip=True, skip_header=1, dtype=None,
                         skip_footer=1, usecols=(4, 5), unpack=True)
    x = x/strech
    y = y/strech
    x_max = (np.max(x)).astype(int)
    y_max = (np.max(y)).astype(int)

    num_x = np.size(x)
    num_y = np.size(y)

    length = np.min([num_x, num_y])
    C = np.zeros((x_max, y_max))
    # x_floor = (np.rint(x)).astype(int)
    # y_floor = (np.rint(y)).astype(int)

    for i in range(1, length):
        x_floor = (np.floor(x[i])).astype(int)
        y_floor = (np.floor(y[i])).astype(int)
        if x_floor >= x_max or y_floor >= y_max:
            continue

        C[x_floor][y_floor] = C[x_floor][y_floor] + 1

    C = scipy.ndimage.zoom(C, strech)
    # print(np.size(C[-1]))
    C = np.transpose(C)
    # print(np.size(C[-1]))

    # levels = np.linspace(0, np.max(C), 290)

    plt.figure(figsize=(7.5, 3.5))

    im = plt.imshow(C, interpolation='quadric', origin='lower', cmap=cm.jet)
    cbar = plt.colorbar(im, orientation='horizontal', ticks=[np.min(C), np.max(C)])
    cbar.ax.set_xticklabels(['Low lipid density', 'High lipid density'])

    # plt.figure()
    # CONT = plt.contour(C, cmap=cm.jet, extent=(0, x_max, 0, y_max))
    # plt.colorbar(CONT, orientation='horizontal')
    savefilename = 'images/' + filename[5:-4] + '_strech=' + str(strech) + '.png'
    # plt.show()
    plt.tick_params(axis='x', which='both', bottom='off', top='off', labelbottom='off')
    plt.tick_params(axis='y', which='both', left='off', right='off', labelleft='off')

    plt.savefig(savefilename, dpi=400)
    print(savefilename, ' has been saved.')

# start in 2999, stop in 290001, step is 2899;
for i in range(290000, 290001, 2899):
    filename = 'data/Config%07d.Pdb' % i
    print(filename)
    for strech in range(3, 4):
        plot_fun(filename, strech)