Merge pull request #16 from anotherjoshsmith/generate_input

Generate input

.travis.yml

@@ -20,6 +20,7 @@ python:
 
 install:
 - travis_retry pip install $PIP_DEPS
+- travis_retry pip install numpy cython
 - travis_retry pip install -r requirements.txt
 - travis_retry pip install -e .
 
@@ -29,7 +30,7 @@ before_script: # configure a headless display to test plot generation
 - sleep 3 # give xvfb some time to start
 
 script:
-- flake8 --ignore N802,N806,W503 `find . -name \*.py | grep -v setup.py | grep -v version.py | grep -v __init__.py | grep -v /docs/`
+- flake8 --ignore N802,N806,W503,E126 `find . -name \*.py | grep -v setup.py | grep -v version.py | grep -v __init__.py | grep -v /docs/`
 - mkdir for_test
 - cd for_test
 - py.test --pyargs plumitas --cov-report term-missing --cov=plumitas

plumitas/__init__.py

@@ -1,3 +1,4 @@
 from __future__ import absolute_import, division, print_function
 from .version import __version__  # noqa
 from .core import *
+from .input import *

plumitas/input.py

@@ -0,0 +1,148 @@
+def header_to_string(header, topology):
+    string = []
+    if header['restart']:
+        string.append('RESTART\n')
+
+    if header['wholemolecules']:
+        string.append('WHOLEMOLECULES ')
+        for idx, query in enumerate(header['wholemolecules']):
+            atoms = topology.select(query)
+            a_str = ','.join(str(x) for x in atoms)
+            entity = f'ENTITY{str(idx)}={a_str} '
+            string.append(entity)
+        string.append('\n')
+
+    string.append('\n')
+    return ''.join(string)
+
+
+def groups_to_string(groups, topology):
+    string = []
+
+    for key, group in groups.items():
+        label = f'{key}: '
+        action = ''.join(action for action in group.keys())
+        query = group[action]
+        atoms = topology.select(query)
+        a_str = ','.join(str(x) for x in atoms)
+        group_string = f'{label}{action.upper()} ATOMS={a_str}\n'
+        string.append(group_string)
+
+    string.append('\n')
+    return ''.join(string)
+
+
+def cvs_to_string(cvs, topology):
+    string = []
+    for key, value in cvs.items():
+        cv_type = ''.join(key for key in value.keys())
+        label = f'{key}: {cv_type.upper()} '
+        string.append(label)
+
+        if cv_type == 'torsion':
+            if value['torsion']['atoms']:
+                atoms = value['torsion']['atoms']
+                string.append(f'{atoms}\n')
+                continue
+
+            if (not value['torsion']['resid']
+                    or not value['torsion']['resid']):
+                print('If "atoms" are not specified, you must supply '
+                      'both "resid" and "angle".')
+                continue
+
+            resid = value['torsion']['resid']
+            atom_lookup = {
+                'phi': '(residue {} and name C) '
+                       'or (residue {} and name N CA C)'
+                       ''.format((resid - 1), resid),
+                'psi': '(residue {} and name N CA C) '
+                       'or (residue {} and name N'
+                       ''.format(resid, (resid + 1))
+            }
+
+            query = atom_lookup[value['torsion']['angle']]
+            atoms = topology.select(query)
+            a_str = ','.join(str(x) for x in atoms)
+            string.append(f'{a_str}\n')
+
+    string.append('\n')
+    return ''.join(string)
+
+
+def bias_to_string(bias):
+    string = []
+    method = ''.join(key for key in bias.keys())
+    string.append(f'{method.upper()} ...\n')
+
+    for key, value in bias[method].items():
+        string.append(f'{key.upper()}={value}\n')
+
+    string.append(f'... {method.upper()}\n')
+    string.append('\n')
+    return ''.join(string)
+
+
+def footer_to_string(footer):
+    string = []
+    for action, arguments in footer.items():
+        string.append(f'{action.upper()} ')
+        for key, value in footer[action].items():
+            string.append(f'{key.upper()}={value} ')
+        string.append('\n')
+
+    string.append('\n')
+    return ''.join(string)
+
+
+def generate_input(mdtraj_top, out_file='plumed.dat', **kwargs):
+    """
+    Converts an input dictionary object into a plumed run file.
+
+    Parameters
+    ----------
+    mdtraj_top : mdtraj.traj.Topology
+        A mdtraj Topology object generated from an input configuration
+        for your system. This allows automated atom selection with
+        simple VMD-style atom queries. This method is especially useful
+        when dealing with many collective variables and/or atom groups.
+    out_file : string
+        Name of file to be used with enhanced sampling simulation
+        in PLUMED. Default plumed.dat is common choice.
+    **kwargs : dict
+        Dictionary containing plumed input stored in header, groups,
+        cvs, bias, and footer sections. Eventually, these dicts will
+        be conveniently created with an interactive GUI.
+
+    Returns
+    -------
+    None
+    """
+    plumed_dat = []
+    if 'header' in kwargs.keys():
+        plumed_dat.append(
+            header_to_string(kwargs['header'], mdtraj_top)
+        )
+    if 'groups' in kwargs.keys():
+        plumed_dat.append(
+            groups_to_string(kwargs['groups'], mdtraj_top)
+        )
+    if 'collective_variables' in kwargs.keys():
+        plumed_dat.append(
+            cvs_to_string(kwargs['collective_variables'],
+                          mdtraj_top)
+        )
+    if 'bias' in kwargs.keys():
+        plumed_dat.append(
+            bias_to_string(kwargs['bias'])
+        )
+    if 'footer' in kwargs.keys():
+        plumed_dat.append(
+            footer_to_string(kwargs['footer'])
+        )
+
+    plumed_input = ''.join(plumed_dat)
+
+    with open(out_file, 'w') as f:
+        f.write(plumed_input)
+    return

plumitas/tests/test_plumitas.py

@@ -2,6 +2,7 @@
 import os.path as op
 
 import plumitas as plm
+import mdtraj as md
 
 data_path = op.join(plm.__path__[0], 'data')
 
@@ -77,3 +78,51 @@ def test_no_bias():
     project.get_bias_params(bin_space_plumed, bias_type='pbmetad')
     project.get_bias_params(bin_plumed, bias_type='metad')
     project.get_bias_params(bin_space_plumed, bias_type='metad')
+
+
+def test_input():
+    conf = op.join(data_path, "topology/conf.gro")
+    traj = md.load(conf)
+    top = traj.top
+
+    plumed = {'header': {'restart': False,
+                         'wholemolecules': ['protein', 'resname ALA']},
+
+              'groups': {'sidechain_com': {'com': 'sidechain and resname ALA'}
+                         },
+
+              'collective_variables': {
+                  'phi': {'torsion': {'angle': 'phi',
+                                      'resid': 2,
+                                      'atoms': '',
+                                      }
+                          },
+                  'psi': {'torsion': {'atoms': '7,9,15,17',
+                                      'angle': 'psi',
+                                      'resid': 2,
+                                      },
+                          },
+              },
+
+              'bias': {'pbmetad': {'label': 'pbmetad',
+                                   'arg': 'phi,psi',
+                                   'temp': '300',
+                                   'pace': '500',
+                                   'biasfactor': '15',
+                                   'height': '1.2',
+                                   'sigma': '0.35,0.35',
+                                   'grid_min': '-pi,-pi',
+                                   'grid_max': 'pi,pi',
+                                   'grid_spacing': '0.1,0.1',
+                                   'file': 'HILLS_phi,HILLS_psi'
+                                   }
+                       },
+
+              'footer': {'print': {'stride': '500',
+                                   'arg': 'phi,psi,pbmetad.bias',
+                                   'file': 'COLVAR'
+                                   },
+                         }
+              }
+
+    plm.generate_input(top, **plumed)

plumitas/version.py

@@ -5,7 +5,7 @@
 _version_major = 0
 _version_minor = 0
 _version_micro = 1  # use '' for first of series, number for 1 and above
-_version_extra = 'dev2'
+_version_extra = 'dev3'
 # _version_extra = ''  # Uncomment this for full releases
 
 # Construct full version string from these.

requirements.txt

@@ -2,3 +2,5 @@ numpy
 scipy
 matplotlib
 pandas
+cython
+mdtraj