1 2
import logging
2 2
import os as os
3 2
import shutil
4 2
from datetime import datetime
5 2
from functools import lru_cache
6

7 2
import parmed as pmd
8 2
import pytraj as pt
9 2
from parmed.structure import Structure as ParmedStructureClass
10

11 2
logger = logging.getLogger(__name__)
12

13

14 2
def get_key(dct, value):
15
    """
16
    Get dictionary key given the value.
17

18
    NOTE: this function will return a list of keys if there are more than one key with the same value
19
          in the order they are found.
20
    """
21 0
    key = [key for key in dct if (dct[key] == value)]
22

23 0
    if len(key) > 1:
24 0
        logger.warning(
25
            "There more than one key with the same value. Please check if this is the desired output."
26
        )
27

28 0
    return key
29

30

31 2
def override_dict(dct, custom):
32
    """Overrides dictionary values from that of a custom dictionary."""
33 2
    for key, value in custom.items():
34 2
        if value is not None:
35 2
            logger.debug("Overriding {} = {}".format(key, value))
36 2
            dct[key] = value
37

38

39 2
def return_parmed_structure(filename):
40
    """
41
    Return a structure object from a filename.
42

43
    Parameters
44
    ----------
45
    filename : str
46
        The name of the file to load
47

48
    Returns
49
    -------
50
    structure : :class:`parmed.structure.Structure`
51

52
    """
53
    # `parmed` can read both PDBs and
54
    # .inpcrd/.prmtop files with the same function call.
55 2
    try:
56 2
        structure = pmd.load_file(filename)
57 2
        logger.info("Loaded {}...".format(filename))
58 0
    except IOError:
59 0
        logger.error("Unable to load file: {}".format(filename))
60 2
    return structure
61

62

63 2
@lru_cache(maxsize=32)
64 2
def index_from_mask(structure, mask, amber_index=False):
65
    """
66
    Return the atom indicies for a given selection mask.
67

68
    Parameters
69
    ----------
70
    structure : `class`:`parmed.structure.Structure`
71
        The structure that contains the atoms
72
    mask : str
73
        The atom mask
74
    amber_index : bool
75
        If true, 1 will be added to the returned indices
76

77
    Returns
78
    -------
79
    indices : int
80
        Atom index or indices corresponding to the mask
81

82
    """
83 2
    if amber_index:
84 2
        index_offset = 1
85
    else:
86 2
        index_offset = 0
87 2
    if isinstance(structure, str):
88 2
        structure = return_parmed_structure(structure)
89 2
    elif isinstance(structure, ParmedStructureClass):
90 2
        pass
91
    else:
92 0
        raise Exception(
93
            "index_from_mask does not support the type associated with structure:"
94
            + type(structure)
95
        )
96
    # http://parmed.github.io/ParmEd/html/api/parmed/parmed.amber.mask.html?highlight=mask#module-parmed.amber.mask
97 2
    indices = [
98
        i + index_offset for i in pmd.amber.mask.AmberMask(structure, mask).Selected()
99
    ]
100 2
    logger.debug("There are {} atoms in the mask {}  ...".format(len(indices), mask))
101 2
    return indices
102

103

104 2
def make_window_dirs(
105
    window_list, stash_existing=False, path="./", window_dir_name="windows"
106
):
107
    """
108
    Make a series of windows to hold simulation data.
109

110
    Parameters
111
    ----------
112
    window_list : list
113
        List of simulation windows. The names in this list will be used for the folders.
114
    stash_existing : bool
115
        Whether to move an existing windows directory to a backup with the current time
116
    path :
117
        Root path for the directories
118
    window_dir_name :
119
        Name for the top level directory
120

121
    Returns
122
    -------
123

124
    """
125

126 2
    win_dir = os.path.join(path, window_dir_name)
127

128 2
    if stash_existing and os.path.isdir(win_dir):
129 0
        stash_dir = os.path.join(
130
            path, window_dir_name + "_{:%Y.%m.%d_%H.%M.%S}".format(datetime.now())
131
        )
132 0
        shutil.move(win_dir, stash_dir)
133

134 2
    for window in window_list:
135 2
        window_path = os.path.join(win_dir, window)
136 2
        if not os.path.exists(window_path):
137 2
            os.makedirs(window_path)
138

139

140 2
def strip_prmtop(prmtop, mask=":WAT,:Na+,:Cl-"):
141
    """Strip residues from a structure and write a new parameter file. This could probably also be done with ParmEd.
142

143
    Parameters:
144
    ----------
145
    prmtop : {str}
146
        Existing parameter file
147
    mask : {str}, optional
148
        The list of atom masks to strip (the default is [':WAT,:Na+,:Cl-'])
149

150
    Returns:
151
    -------
152
    stripped.topology
153
        The stripped topology that can be used to read stripped trajectories with `pytraj`
154

155
    """
156

157 2
    structure = pt.load_topology(os.path.normpath(prmtop))
158 2
    stripped = pt.strip(mask, structure)
159
    # stripped_name = os.path.join(os.path.splitext(prmtop)[0], '-stripped', os.path.splitext(prmtop)[1])
160
    # stripped.save(filename=stripped_name)
161
    # logger.debug('Stripping {} from parameter file and writing {}...'.format(mask, stripped_name))
162 2
    return stripped
163

164

165 2
def parse_mden(file):
166
    """
167
    Return energies from an AMBER `mden` file.
168

169
    Parameters
170
    ----------
171
    file : str
172
        Name of output file
173

174
    Returns
175
    -------
176
    energies : dict
177
        A dictionary containing VDW, electrostatic, bond, angle, dihedral, V14, E14, and total energy.
178

179
    """
180

181 2
    vdw, ele, bnd, ang, dih, v14, e14 = [], [], [], [], [], [], []
182

183 2
    with open(file, "r") as f:
184 2
        for line in f.readlines()[10:]:
185 2
            words = line.rstrip().split()
186 2
            if words[0] == "L6":
187 2
                vdw.append(float(words[3]))
188 2
                ele.append(float(words[4]))
189 2
            elif words[0] == "L7":
190 2
                bnd.append(float(words[2]))
191 2
                ang.append(float(words[3]))
192 2
                dih.append(float(words[4]))
193 2
            elif words[0] == "L8":
194 2
                v14.append(float(words[1]))
195 2
                e14.append(float(words[2]))
196

197 2
    energies = {
198
        "Bond": bnd,
199
        "Angle": ang,
200
        "Dihedral": dih,
201
        "V14": v14,
202
        "E14": e14,
203
        "VDW": vdw,
204
        "Ele": ele,
205
        "Total": [sum(x) for x in zip(bnd, ang, dih, v14, e14, vdw, ele)],
206
    }
207

208 2
    return energies
209

210

211 2
def parse_mdout(file):
212
    """
213
    Return energies from an AMBER `mdout` file.
214

215
    Parameters
216
    ----------
217
    file : str
218
        Name of output file
219

220
    Returns
221
    -------
222
    energies : dict
223
        A dictionary containing VDW, electrostatic, bond, angle, dihedral, V14, E14, and total energy.
224

225
    """
226

227 2
    vdw, ele, bnd, ang, dih, v14, e14 = [], [], [], [], [], [], []
228 2
    restraint = []
229

230 2
    with open(file, "r") as f:
231 2
        for line in f.readlines():
232 2
            words = line.rstrip().split()
233 2
            if len(words) > 1:
234 2
                if "BOND" in words[0]:
235 2
                    bnd.append(float(words[2]))
236 2
                    ang.append(float(words[5]))
237 2
                    dih.append(float(words[8]))
238 2
                if "VDWAALS" in words[0]:
239 2
                    vdw.append(float(words[2]))
240 2
                    ele.append(float(words[5]))
241 2
                if "1-4" in words[0]:
242 2
                    v14.append(float(words[3]))
243 2
                    e14.append(float(words[7]))
244 2
                    restraint.append(float(words[10]))
245

246 2
    energies = {
247
        "Bond": bnd,
248
        "Angle": ang,
249
        "Dihedral": dih,
250
        "V14": v14,
251
        "E14": e14,
252
        "VDW": vdw,
253
        "Ele": ele,
254
        "Restraint": restraint,
255
        "Total": [sum(x) for x in zip(bnd, ang, dih, v14, e14, vdw, ele)],
256
    }
257

258 2
    return energies

Read our documentation on viewing source code .

Loading