Source code for pymatgen.core.trajectory

# 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)