analyze_gromacs.py

#!/usr/bin/env python

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np


DTYPES = {
    'em': [('step', int),
        ('ljshort', float), ('ljlong', float),
        ('coulshort', float), ('coullong', float),
        ('Epot', float)],
    'nvt': [('t', int),
        ('Epot', float), ('Ekin', float), ('Etot', float),
        ('Ttot', float), ('Tprot', float), ('Tnprot', float)],
    'npt': [('t', int),
        ('Epot', float), ('Ekin', float), ('Etot', float),
        ('p', float), ('V', float), ('rho', float)]
    }


def get_data(datafile, dtype):
    return np.genfromtxt(datafile, dtype=DTYPES[dtype],
            converters={0: lambda x: int(float(x))})


def plot_em(data, title):
    f = plt.figure()
    f.suptitle(title)
    lj = f.add_subplot(311, xlabel='EM Step', ylabel=r'$E_{LJ}$ (kJ mol$^{-1}$)')
    coulomb = f.add_subplot(312, xlabel='EM Step', ylabel=r'$E_{Coulomb}$ (kJ mol$^{-1}$)')
    epot = f.add_subplot(313, xlabel='EM Step', ylabel=r'$E_{pot}$ (kJ mol$^{-1}$)')

    lj.plot(data['step'], data['ljshort'], label=r'$LJ_{short}$')
    lj.plot(data['step'], data['ljlong'], label=r'$LJ_{long}$')
    coulomb.plot(data['step'], data['coulshort'], label=r'$Coulomb_{short}$')
    coulomb.plot(data['step'], data['coullong'], label=r'$Coulomb_{reciprocal}$')
    epot.plot(data['step'], data['Epot'], label=r'$E_{pot}$')

    lj.legend(); coulomb.legend(); epot.legend()
    return f


def plot_nvt(data, title):
    f = plt.figure()
    f.suptitle(title)
    E = f.add_subplot(211, xlabel='Time (ps)', ylabel=r'E (kJ mol$^{-1}$)')
    T = f.add_subplot(212, xlabel='Time (ps)', ylabel=r'T (kJ mol$^{-1}$)')

    E.plot(data['t'], data['Ekin'], label=r'$E_{kin}$')
    E.plot(data['t'], data['Epot'], label=r'$E_{pot}$')
    E.plot(data['t'], data['Etot'], label=r'$E_{total}$')
    T.plot(data['t'], data['Tprot'], label=r'$T_{protein}$')
    T.plot(data['t'], data['Tnprot'], label=r'$T_{non-protein}$')
    T.plot(data['t'], data['Ttot'], label=r'$T_{total}$')

    E.legend(); T.legend()
    return f


def plot_npt(data, title):
    f = plt.figure()
    f.suptitle(title)
    E = f.add_subplot(411, xlabel='Time (ps)', ylabel=r'E (kJ mol$^{-1}$)')
    p = f.add_subplot(412, xlabel='Time (ps)', ylabel=r'p (bar)')
    V = f.add_subplot(413, xlabel='Time (ps)', ylabel=r'V (nm$^3$)')
    rho = f.add_subplot(414, xlabel='Time (ps)', ylabel=r'$\rho$ (kg m$^{3}$)')

    E.plot(data['t'], data['Ekin'], label=r'$E_{kin}$')
    E.plot(data['t'], data['Epot'], label=r'$E_{pot}$')
    E.plot(data['t'], data['Etot'], label=r'$E_{total}$')
    p.plot(data['t'], data['p'], label='p (bar)')
    V.plot(data['t'], data['V'], label=r'V')
    rho.plot(data['t'], data['rho'], label=r'$\rho$')

    E.legend(); p.legend(); V.legend(); rho.legend()
    return f


def plot_md(data, title):
    import sys
    print "Not yet implemented"
    sys.exit(1)


def parse_args():
    import argparse as ap
    p = ap.ArgumentParser()
    p.add_argument('datafile', help="File to interpret")
    p.add_argument('-t', '--type', required=True, choices=['em', 'nvt', 'npt',
        'md'], help="Type of file to interpret")
    p.add_argument('-T', '--title', help="Plot title")
    p.add_argument('output', help="Output file")
    return p.parse_args()


def main():
    args = parse_args()
    data = get_data(args.datafile, args.type)
    plotfuncs = {
            'em': plot_em,
            'nvt': plot_nvt,
            'npt': plot_npt,
            'md': plot_md
            }
    figure = plotfuncs[args.type](data, args.title)
    figure.savefig(args.output)


if __name__ == '__main__':
    main()