# 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"
[docs]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
[docs] def get_structure(self, i):
"""
Returns structure at specified index
Args:
i (int): Index of structure
Returns:
(Structure) pymatgen structure object
"""
return self[i]
[docs] 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
[docs] 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
[docs] 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')
[docs] 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)
[docs] @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)
[docs] @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)
[docs] 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
[docs] 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)