#!/usr/bin/env python

import sys

import argparse

import csv

import re

import yaml

import numpy as np

import matplotlib.pyplot as plt

import cycler

import pandas as pd

import MDAnalysis as mda

from MDAnalysis.analysis.distances import distance_array

from ffparam_charges_vs_qm import load_prm, COULOMB

def calc_mm_interaction_energy(u, ligand_spec, prm, solute_spec, vdw_shift=0.0, patches=None):

return np.array(electro), np.array(lj)

def scale_factors_for_lambda(lambda_val, alchElecLambdaStart=0.5, alchVdwLambdaEnd=0.5, annihilate=True):

return electro_scale, lj_scale

# TODO: Since this is just scaling, we don't have to sum across ligand atoms...do we?

def do_fep_scaling(electro, lj, lambda0, lambda1, alchElecLambdaStart, alchVdwLambdaEnd):

return electro0, lj0, electro1, lj1

# TODO: If weights are extremely nonuniform, complain that the energies1 do not cover a similar enough

# TODO: part of the phase space represented by energies0, and therefore the reweighted energy is probably

# TODO: far off.

def reweight_energies(e0, e1, beta=1.0/0.6):

return e1 * weights

# TODO: This should just take a CSV filename or something

def generate_plot(datatype, data, ligand_resname, dont_plot_atomtypes=''):

print >> sys.stderr, 'Wrote a pretty picture to %s.pdf' % prefix

if __name__ == '__main__':

ap = argparse.ArgumentParser(description='Calculates single point energies with MM, with respect to a ligand.')

ap.add_argument('ligand_resname', help='Ligand residue name (e.g. SKE, APM)')

ap.add_argument('namdconf', help='NAMD configuration containing "parameters" and "structure" keywords from which to glean PSF and force field filenames')

ap.add_argument('dcd', nargs='+', help='DCD trajectories to go with the PSF file specified in namdconf')

ap.add_argument('--solute-spec', default='protein or nucleic', help='Specify which part of the system we care about ligand interactions with')

ap.add_argument('--dont-plot-atomtypes', default='HA1,HA3,HP,HGA2,HGAAM1', help='Atom types not to plot, separated by commas')

ap.add_argument('--patches', help='YAML file containing patches to ligand parameters, such as charges')

ap.add_argument('--lambda0', type=float, help='FEP only: Lambda0 for the provided trajectory')

ap.add_argument('--lambda1', type=float, help='FEP only: Lambda1 for the provided trajectory')

ap.add_argument('--alchElecLambdaStart', type=float, default=0.5, help='FEP only: Governs lambda scaling, as it does in NAMD')

ap.add_argument('--alchVdwLambdaEnd', type=float, default=0.5, help='FEP only: Governs lambda scaling, as it does in NAMD')

ap.add_argument('--alchVdwShiftCoeff', type=float, default=6.0, help='FEP only: shift distance by this much A for L-J as in NAMD')

args = ap.parse_args()

psf_filename, prm_filenames = None, []

# Glean PSF and force field parameters from NAMD simulation input file

print >>sys.stderr, 'Reading NAMD simulation config', args.namdconf

with open(args.namdconf) as f:

psf_re = re.compile(r'structure\s+([\w\-/.]+)')

prm_re = re.compile(r'parameters\s+([\w\-/.]+)')

for line in f:

m = psf_re.match(line.strip())

if m: psf_filename = m.group(1)

m = prm_re.match(line.strip())

if m: prm_filenames.append(m.group(1))

print >>sys.stderr, 'Guessed PSF filename:', psf_filename

print >>sys.stderr, 'Guessed force field parameter files:', ', '.join(prm_filenames)

print >>sys.stderr, 'Solute spec:', args.solute_spec

u = mda.Universe(psf_filename, args.dcd)

prm = {}

for prm_fname in prm_filenames:

prm.update(load_prm(prm_fname))

print >>sys.stderr, 'Loaded %d MM parameters.' % len(prm)

patches = None

if args.patches is not None:

patches = yaml.load(open(args.patches))

print >>sys.stderr, 'Patches:', patches

ligand_spec = 'resname %s' % args.ligand_resname

# Calculate the energies assuming no FEP and no patches

print >>sys.stderr, 'Calculating energies with no patches.'

electro_orig, lj_orig = calc_mm_interaction_energy(u, ligand_spec, prm, solute_spec=args.solute_spec)

all_data = [('electro', electro_orig), ('lj', lj_orig)]

# If we have patches, recalculate the energies

if patches is not None:

print >> sys.stderr, 'Calculating energies WITH patches.'

vdw_shift = 0.0

if args.lambda0 is not None and args.lambda1 is not None:

vdw_shift = args.alchVdwShiftCoeff

electro_patched, lj_patched = calc_mm_interaction_energy(u, ligand_spec, prm, solute_spec=args.solute_spec,

patches=patches, vdw_shift=vdw_shift)

# Now we can reweight the patched energies.

# We are reweighting the sum of the ligand interaction energies, so that there is one energy per frame.

electro_reweighted = reweight_energies(np.sum(electro_orig, axis=1), np.sum(electro_patched, axis=1))

print >>sys.stderr, 'Mean electro_orig = %.3f' % np.mean(np.sum(electro_orig, axis=1))

print >>sys.stderr, 'Reweighted electro = %.3f' % np.mean(electro_reweighted)

# We actually don't need to do this, yet, because we don't change any L-J parameters

print >>sys.stderr, 'Not (directly) reweighting L-J energies because we have not implemented patching L-J terms'

lj_reweighted = np.sum(lj_orig, axis=1) / np.size(lj_orig, axis=0)

# For now we'll only redo the FEP if there were patches

if args.lambda0 is not None and args.lambda1 is not None:

electro_reweighted_lambda0, lj_reweighted_lambda0, \

electro_reweighted_lambda1, lj_reweighted_lambda1 = do_fep_scaling(electro_reweighted, lj_reweighted,

args.lambda0, args.lambda1, args.alchElecLambdaStart,

args.alchVdwLambdaEnd)

dG = electro_reweighted_lambda1 - electro_reweighted_lambda0

print >>sys.stderr, 'dG(electro) = %.3f kcal/mol' % np.sum(dG)

# Generate CSV files and accompanying plots

for datatype, data in all_data:

prefix = '%s_%s' % (datatype, args.ligand_resname)

ligand_atoms = u.select_atoms("resname %s" % args.ligand_resname)

labels = ['%s %d %.3f' % (a.type, a.index - ligand_atoms[0].index + 1, a.charge) for a in ligand_atoms]

with open('%s.csv' % prefix, 'w') as f:

writer = csv.writer(f)

writer.writerow(labels)

for row in data:

writer.writerow(row)

print >>sys.stderr, 'Wrote energy data to %s.csv' % prefix

generate_plot(datatype, data, args.ligand_resname, args.dont_plot_atomtypes)