/
trajectory.py
474 lines (415 loc) · 21.7 KB
/
trajectory.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.
"""
This module provides classes used to define a MD trajectory.
"""
import itertools
import os
import warnings
from fnmatch import fnmatch
from typing import List, Union, Sequence
import numpy as np
from monty.io import zopen
from monty.json import MSONable
from pymatgen.core.structure import Structure, Lattice, Element, Specie, DummySpecie, Composition
from pymatgen.io.vasp.outputs import Xdatcar, Vasprun
__author__ = "Eric Sivonxay, Shyam Dwaraknath"
__version__ = "0.0"
__date__ = "Jan 25, 2019"
class Trajectory(MSONable):
"""
Trajectory object that stores structural information related to a MD simulation.
Provides basic functions such as slicing trajectory or obtaining displacements.
"""
def __init__(self, lattice: Union[List, np.ndarray, Lattice],
species: List[Union[str, Element, Specie, DummySpecie, Composition]],
frac_coords: List[Sequence[Sequence[float]]],
time_step: float = 2,
site_properties: dict = None,
frame_properties: dict = None,
constant_lattice: bool = True,
coords_are_displacement: bool = False,
base_positions: Sequence[Sequence[float]] = None):
"""
Create a trajectory object
Args:
lattice: The lattice as any 2D array. Each row should correspond to a lattice
vector. E.g., [[10,0,0], [20,10,0], [0,0,30]] specifies a
lattice with lattice vectors [10,0,0], [20,10,0] and [0,0,30].
species: List of species on each site. Can take in flexible input,
including:
i. A sequence of element / specie specified either as string
symbols, e.g. ["Li", "Fe2+", "P", ...] or atomic numbers,
e.g., (3, 56, ...) or actual Element or Specie objects.
ii. List of dict of elements/species and occupancies, e.g.,
[{"Fe" : 0.5, "Mn":0.5}, ...]. This allows the setup of
disordered structures.
frac_coords (MxNx3 array): list of fractional coordinates of
each species
time_step (int, float): Timestep of simulation in femtoseconds. Defaults to 2fs.
site_properties (list): Properties associated with the sites as a list of
dicts of sequences, e.g., [{"magmom":[5,5,5,5]}, {"magmom":[5,5,5,5]}]. The sequences
have to be the same length as the atomic species and fractional_coords. Number of supplied
dicts should match number of frames in trajectory
Defaults to None for no properties.
frame_properties (dict): Properties of the trajectory such as energy, pressure, etc. each property should
have a length equal to the trajectory length. eg: {'energy': [#, #, #, #], 'pressure': [0, 0.1, 0 0.02]}
constant_lattice (bool): Whether the lattice changes during the simulation, such as in an NPT MD simulation.
coords_are_displacement (bool): Whether supplied coordinates are given in displacements (True) or
positions (False)
base_positions (Nx3 array): The starting positions of all atoms in trajectory. Used to reconstruct positions
when converting from displacements to positions. Only needs to be specified if
coords_are_displacement=True. Defaults to first index of frac_coords if coords_are_displacement=False.
"""
# To support from_dict and as_dict
if isinstance(frac_coords, list):
frac_coords = np.array(frac_coords)
if isinstance(lattice, Lattice):
lattice = lattice.matrix
if isinstance(lattice, list):
lattice = np.array(lattice)
self.frac_coords = frac_coords
if coords_are_displacement:
if base_positions is None:
warnings.warn("Without providing an array of starting positions, \
the positions for each time step will not be available")
self.base_positions = base_positions
else:
self.base_positions = frac_coords[0]
self.coords_are_displacement = coords_are_displacement
if not constant_lattice and np.shape(lattice) == (3, 3):
self.lattice = [lattice for i in range(np.shape(self.frac_coords)[0])]
else:
self.lattice = lattice
self.constant_lattice = constant_lattice
self.species = species
self.site_properties = site_properties
self.frame_properties = frame_properties
self.time_step = time_step
def get_structure(self, i):
"""
Returns structure at specified index
Args:
i (int): Index of structure
Returns:
(Structure) pymatgen structure object
"""
return self[i]
def to_positions(self):
"""
Converts fractional coordinates of trajectory into positions
"""
if self.coords_are_displacement:
cumulative_displacements = np.cumsum(self.frac_coords, axis=0)
positions = self.base_positions + cumulative_displacements
self.frac_coords = positions
self.coords_are_displacement = False
def to_displacements(self):
"""
Converts position coordinates of trajectory into displacements between consecutive frames
"""
if not self.coords_are_displacement:
displacements = np.subtract(self.frac_coords, np.roll(self.frac_coords, 1, axis=0))
displacements[0] = np.zeros(np.shape(self.frac_coords[0]))
# Deal with PBC
displacements = [np.subtract(item, np.round(item)) for item in displacements]
self.frac_coords = displacements
self.coords_are_displacement = True
def extend(self, trajectory):
"""
Concatenate another trajectory
Args:
trajectory (Trajectory): Trajectory to add
"""
if self.time_step != trajectory.time_step:
raise ValueError('Trajectory not extended: Time steps of trajectories is incompatible')
if len(self.species) != len(trajectory.species) and self.species != trajectory.species:
raise ValueError('Trajectory not extended: species in trajectory do not match')
# Ensure both trajectories are in positions before combining
self.to_positions()
trajectory.to_positions()
self.site_properties = self._combine_site_props(self.site_properties, trajectory.site_properties,
np.shape(self.frac_coords)[0],
np.shape(trajectory.frac_coords)[0])
self.frame_properties = self._combine_frame_props(self.frame_properties, trajectory.frame_properties,
np.shape(self.frac_coords)[0],
np.shape(trajectory.frac_coords)[0])
self.frac_coords = np.concatenate((self.frac_coords, trajectory.frac_coords), axis=0)
self.lattice, self.constant_lattice = self._combine_lattice(self.lattice, trajectory.lattice,
np.shape(self.frac_coords)[0],
np.shape(trajectory.frac_coords)[0])
def __iter__(self):
for i in range(np.shape(self.frac_coords)[0]):
yield self[i]
def __len__(self):
return np.shape(self.frac_coords)[0]
def __getitem__(self, frames):
"""
Gets a subset of the trajectory if a slice is given, if an int is given, return a structure
Args:
frames (int, slice): int or slice of trajectory to return
Return:
(Trajectory, Structure) Subset of trajectory
"""
# If trajectory is in displacement mode, return the displacements at that frame
if self.coords_are_displacement:
if isinstance(frames, int):
if frames >= np.shape(self.frac_coords)[0]:
raise ValueError('Selected frame exceeds trajectory length')
# For integer input, return the displacements at that timestep
return self.frac_coords[frames]
if isinstance(frames, slice):
# For slice input, return a list of the displacements
start, stop, step = frames.indices(len(self))
return [self.frac_coords[i] for i in range(start, stop, step)]
if isinstance(frames, (list, np.ndarray)):
# For list input, return a list of the displacements
pruned_frames = [i for i in frames if i < len(self)] # Get rid of frames that exceed trajectory length
if len(pruned_frames) < len(frames):
warnings.warn('Some or all selected frames exceed trajectory length')
return [self.frac_coords[i] for i in pruned_frames]
raise Exception('Given accessor is not of type int, slice, list, or array')
# If trajectory is in positions mode, return a structure for the given frame or trajectory for the given frames
if isinstance(frames, int):
if frames >= np.shape(self.frac_coords)[0]:
raise ValueError('Selected frame exceeds trajectory length')
# For integer input, return the structure at that timestep
lattice = self.lattice if self.constant_lattice else self.lattice[frames]
site_properties = self.site_properties[frames] if self.site_properties else None
site_properties = self.site_properties[frames] if self.site_properties else None
return Structure(Lattice(lattice), self.species, self.frac_coords[frames],
site_properties=site_properties,
to_unit_cell=True)
if isinstance(frames, slice):
# For slice input, return a trajectory of the sliced time
start, stop, step = frames.indices(len(self))
pruned_frames = range(start, stop, step)
lattice = self.lattice if self.constant_lattice else [self.lattice[i] for i in pruned_frames]
frac_coords = [self.frac_coords[i] for i in pruned_frames]
if self.site_properties is not None:
site_properties = [self.site_properties[i] for i in pruned_frames]
else:
site_properties = None
if self.frame_properties is not None:
frame_properties = {}
for key, item in self.frame_properties.items():
frame_properties[key] = [item[i] for i in pruned_frames]
else:
frame_properties = None
return Trajectory(lattice, self.species, frac_coords, time_step=self.time_step,
site_properties=site_properties, frame_properties=frame_properties,
constant_lattice=self.constant_lattice, coords_are_displacement=False,
base_positions=self.base_positions)
if isinstance(frames, (list, np.ndarray)):
# For list input, return a trajectory of the specified times
pruned_frames = [i for i in frames if i < len(self)] # Get rid of frames that exceed trajectory length
if len(pruned_frames) < len(frames):
warnings.warn('Some or all selected frames exceed trajectory length')
lattice = self.lattice if self.constant_lattice else [self.lattice[i] for i in pruned_frames]
frac_coords = [self.frac_coords[i] for i in pruned_frames]
if self.site_properties is not None:
site_properties = [self.site_properties[i] for i in pruned_frames]
else:
site_properties = None
if self.frame_properties is not None:
frame_properties = {}
for key, item in self.frame_properties.items():
frame_properties[key] = [item[i] for i in pruned_frames]
else:
frame_properties = None
return Trajectory(lattice, self.species, frac_coords, time_step=self.time_step,
site_properties=site_properties, frame_properties=frame_properties,
constant_lattice=self.constant_lattice, coords_are_displacement=False,
base_positions=self.base_positions)
raise Exception('Given accessor is not of type int, slice, tuple, list, or array')
def copy(self):
"""
:return: Copy of Trajectory.
"""
return Trajectory(self.lattice, self.species, self.frac_coords, time_step=self.time_step,
site_properties=self.site_properties, frame_properties=self.frame_properties,
constant_lattice=self.constant_lattice, coords_are_displacement=False,
base_positions=self.base_positions)
@classmethod
def from_structures(cls, structures, constant_lattice=True, **kwargs):
"""
Convenience constructor to obtain trajectory from a list of structures.
Note: Assumes no atoms removed during simulation
Args:
structures (list): list of pymatgen Structure objects.
constant_lattice (bool): Whether the lattice changes during the simulation, such as in an NPT MD
simulation. True results in
Returns:
(Trajectory)
"""
frac_coords = [structure.frac_coords for structure in structures]
if constant_lattice:
lattice = structures[0].lattice.matrix
else:
lattice = [structure.lattice.matrix for structure in structures]
site_properties = {}
site_properties = [structure.site_properties for structure in structures]
return cls(lattice, structures[0].species, frac_coords, site_properties=site_properties,
constant_lattice=constant_lattice, **kwargs)
@classmethod
def from_file(cls, filename, constant_lattice=True, **kwargs):
"""
Convenience constructor to obtain trajectory from XDATCAR or vasprun.xml file
Args:
filename (str): The filename to read from.
constant_lattice (bool): Whether the lattice changes during the simulation, such as in an NPT MD
simulation. True results in
Returns:
(Trajectory)
"""
# TODO: Support other filetypes
fname = os.path.basename(filename)
if fnmatch(fname, "*XDATCAR*"):
structures = Xdatcar(filename).structures
elif fnmatch(fname, "vasprun*.xml*"):
structures = Vasprun(filename).structures
else:
raise ValueError("Unsupported file")
return cls.from_structures(structures, constant_lattice=constant_lattice, **kwargs)
def as_dict(self):
"""
:return: MSONAble dict.
"""
d = {"@module": self.__class__.__module__,
"@class": self.__class__.__name__,
"species": self.species, "time_step": self.time_step,
"site_properties": self.site_properties,
"frame_properties": self.frame_properties,
"constant_lattice": self.constant_lattice,
"coords_are_displacement": self.coords_are_displacement,
"base_positions": self.base_positions}
d["lattice"] = self.lattice.tolist()
d["frac_coords"] = self.frac_coords.tolist()
return d
@staticmethod
def _combine_lattice(attr_1, attr_2, len_1, len_2):
"""
Helper function to combine trajectory properties such as site_properties or lattice
"""
if np.shape(attr_1) == (3, 3) and np.shape(attr_2) == (3, 3):
attribute = attr_1
attribute_constant = True
elif np.shape(attr_1) == 3 and np.shape(attr_2) == 3:
attribute = np.concatenate((attr_1, attr_2), axis=0)
attribute_constant = False
else:
attribute = [attr_1.copy()] * len_1 if isinstance(attr_1, list) else attr_1.copy()
attribute.extend([attr_2.copy()] * len_2 if isinstance(attr_2, list) else attr_2.copy())
attribute_constant = False
return attribute, attribute_constant
@staticmethod
def _combine_site_props(attr_1, attr_2, len_1, len_2):
"""
Helper function to combine site properties of 2 trajectories
"""
if attr_1 is None and attr_2 is None:
return None
if attr_1 is None or attr_2 is None:
new_site_properties = []
if attr_1 is None:
new_site_properties.extend([None for i in range(len_1)])
elif len(attr_1) == 1:
new_site_properties.extend([attr_1[0] for i in range(len_1)])
elif len(attr_1) > 1:
new_site_properties.extend(attr_1)
if attr_2 is None:
new_site_properties.extend([None for i in range(len_2)])
elif len(attr_2) == 1:
new_site_properties.extend([attr_2[0] for i in range(len_2)])
elif len(attr_2) > 1:
new_site_properties.extend(attr_2)
return new_site_properties
if len(attr_1) == 1 and len(attr_2) == 1:
# If both properties lists are do not change within their respective trajectory
if attr_1 == attr_2:
# If both site_properties are the same, only store one
return attr_1
new_site_properties = [attr_1[0] for i in range(len_1)]
new_site_properties.extend([attr_2[0] for i in range(len_2)])
return new_site_properties
if len(attr_1) > 1 and len(attr_2) > 1:
# Both properties have site properties that change within the trajectory, concat both together
return [*attr_1, *attr_2]
new_site_properties = []
if attr_1 is None:
new_site_properties.extend([None for i in range(len_1)])
elif len(attr_1) == 1:
new_site_properties.extend([attr_1[0] for i in range(len_1)])
elif len(attr_1) > 1:
new_site_properties.extend(attr_1)
if attr_2 is None:
new_site_properties.extend([None for i in range(len_2)])
elif len(attr_2) == 1:
new_site_properties.extend([attr_2[0] for i in range(len_2)])
elif len(attr_2) > 1:
new_site_properties.extend(attr_2)
return new_site_properties
@staticmethod
def _combine_frame_props(attr_1, attr_2, len_1, len_2):
"""
Helper function to combine frame properties such as energy or pressure
"""
if attr_1 is None and attr_2 is None:
return None
# Find all common keys
all_keys = set(attr_1.keys()).union(set(attr_2.keys()))
# Initialize dict with the common keys
new_frame_props = dict(zip(all_keys, [[] for i in all_keys]))
for key in all_keys:
if key in attr_1.keys():
new_frame_props[key].extend(attr_1[key])
else:
# If key doesn't exist in the first trajectory, append None for each index
new_frame_props[key].extend([None for i in range(len_1)])
if key in attr_2.keys():
new_frame_props[key].extend(attr_2[key])
else:
# If key doesn't exist in the second trajectory, append None for each index
new_frame_props[key].extend([None for i in range(len_2)])
return new_frame_props
def write_Xdatcar(self, filename="XDATCAR", system=None, significant_figures=6):
"""
Writes Xdatcar to a file. The supported kwargs are the same as those for
the Xdatcar_from_structs.get_string method and are passed through directly.
Args:
filename (str): name of file (It's prudent to end the filename with 'XDATCAR',
as most visualization and analysis software require this for autodetection)
system (str): Description of system
significant_figures (int): Significant figures in the output file
"""
# Ensure trajectory is in position form
self.to_positions()
if system is None:
system = f'{self[0].composition.reduced_formula}'
lines = []
format_str = "{{:.{0}f}}".format(significant_figures)
syms = [site.specie.symbol for site in self[0]]
site_symbols = [a[0] for a in itertools.groupby(syms)]
syms = [site.specie.symbol for site in self[0]]
natoms = [len(tuple(a[1])) for a in itertools.groupby(syms)]
for si, frac_coords in enumerate(self.frac_coords):
# Only print out the info block if
if self.constant_lattice and si == 0:
lines.extend([system, "1.0"])
if self.constant_lattice:
_lattice = self.lattice
else:
_lattice = self.lattice[si]
for latt_vec in _lattice:
lines.append(f'{" ".join([str(el) for el in latt_vec])}')
lines.append(" ".join(site_symbols))
lines.append(" ".join([str(x) for x in natoms]))
lines.append(f"Direct configuration= {str(si + 1)}")
for (frac_coord, specie) in zip(frac_coords, self.species):
coords = frac_coord
line = f'{" ".join([format_str.format(c) for c in coords])} {specie}'
lines.append(line)
xdatcar_string = "\n".join(lines) + "\n"
with zopen(filename, "wt") as f:
f.write(xdatcar_string)