#!/usr/bin/env python
from weakref import WeakValueDictionary
from scm.plams.core.errors import PlamsError
from scm.plams.core.settings import Settings
from scm.plams.interfaces.adfsuite.amsanalysis import AMSAnalysisJob
from scm.plams.tools.kftools import KFFile
from scm.plams.trajectories.rkffile import RKFTrajectoryFile
from scm.plams.trajectories.rkfhistoryfile import RKFHistoryFile
__all__ = ["Trajectory"]
[docs]class Trajectory:
"""
Class representing an AMS trajectory
It creates molecule objects along the trajectory in a list-like fashion.
It can also perform analysis tasks.
Basic use of the |Trajectory| class works as follows:
>>> from scm.plams import Trajectory
>>> trajec = Trajectory('ams.rkf')
>>> nsteps = len(trajec)
>>> mol100 = trajec[100]
>>> for i,mol in enumerate(trajec):
... if i == 100:
... print (mol is mol100)
...
True
As with a regular list, accessing the same element twice yields the same instance of the molecule.
A trajectory may consist of multiple RKF files, for instance after a restarted MD run.
>>> trajec = Trajectory(['md1.results/ams.rkf','md2.results/ams.rkf'])
>>> nsteps = len(trajec)
>>> print ('NSteps in each file: ',trajec.lengths)
NSteps in each file: [181, 19]
Otherwise the object behaves the same, and iteration occurs over the concatenated files.
.. note::
It is possible to change the molecule objects once they are loaded, but this will not change the underlying files.
When the altered molecule objects have been garbage collected, and the same frame is read anew, this frame
will have the original molecule attributes.
>>> import gc
>>> from scm.plams import Trajectory
>>> trajec = Trajectory('ams.rkf')
>>> mol = trajec[100]
>>> print (mol.properties.charge)
0.0
>>> mol.properties.charge = 10000.
>>> print (trajec[100].properties.charge)
1000.0
>>> # remove the reference
>>> mol = None
>>> gc.collect()
>>> print (trajec[100].properties.charge)
0.
The |Trajectory| object can also be used to perform analysis, using |AMSAnalysisJob| behind the scenes.
>>> from scm.plams import Settings, Trajectory
>>> from scm.plams import init, finish
>>> trajec = Trajectory('ams.rkf')
>>> # Define the type of analysis
>>> settings = Settings()
>>> settings.input.Task = 'RadialDistribution'
>>> settings.input.RadialDistribution.AtomsFrom.Element = ['O']
>>> settings.input.RadialDistribution.AtomsTo.Element = ['O']
>>> # Run the analysis
>>> init()
>>> plots = trajec.run_analysis(settings)
>>> finish()
>>> # Store the plot in human readable format
>>> for plot in plots :
... plot.write('%s'%(plot.name+'.txt'))
This results in a text file named RadialDistribution_1.txt, which contains data organized in columns.
"""
def __init__(self, filenames):
"""
Initiates instance of an AMS trajectory object
* ``filenames`` -- List of filepaths of RKF trajectory files or single filepath for RKF file
.. note::
The corresponding file needs to remain on disc while the object is alive
"""
if isinstance(filenames, str):
filenames = [filenames]
# Ceck is RKFHistoryFile needs to be used
classes = []
for fn in filenames:
if not fn.split(".")[-1] == "rkf":
raise PlamsError("Files need to be RKF")
kf = KFFile(fn)
if kf.reader is None:
raise PlamsError("RKF file %s not found" % (fn))
kf.reader._create_index()
if not "History" in kf.reader._sections.keys():
raise PlamsError("%s is not a trajectory file" % (fn))
if "SystemVersionHistory" in kf.reader._sections.keys():
classes.append(RKFHistoryFile)
else:
classes.append(RKFTrajectoryFile)
# Store the file objects and associated data for each file
self.files = [classes[i](fn) for i, fn in enumerate(filenames)]
self.molecules = [rkf.get_plamsmol() for rkf in self.files]
self.lengths = [len(rkf) for rkf in self.files]
# Store the items that have been read as a weak reference
self.weakrefs = WeakValueDictionary()
def __len__(self):
"""
Magic method that returns the length of the trajectory
"""
return sum(self.lengths)
def __iter__(self):
"""
Iterates over molecule objects
"""
for irkf, rkf in enumerate(self.files):
mol = self.molecules[irkf]
for istep in range(self.lengths[irkf]):
i = self._get_index(irkf, istep)
if i in self.weakrefs.keys():
yield self.weakrefs[i]
continue
crd, cell = rkf.read_frame(istep, molecule=mol)
m = mol.copy()
self.weakrefs[i] = m
yield m
def __getitem__(self, s):
"""
Returns Molecule object
* ``s`` -- Python slice object
"""
is_slice = False
if isinstance(s, int):
s = slice(s)
start, stop, step = s.indices(len(self))
indices = range(stop, stop + 1)
elif isinstance(s, slice):
is_slice = True
start, stop, step = s.indices(len(self))
indices = range(start, stop, step)
mols = []
for i in indices:
if i in self.weakrefs.keys():
mol = self.weakrefs[i]
else:
irkf, istep = self._get_filenum_and_stepnum(i)
# This is always the first molecule
mol = self.molecules[irkf].copy()
crd, cell = self.files[irkf].read_frame(istep, molecule=mol)
self.weakrefs[i] = mol
mols.append(mol)
if not is_slice:
mols = mols[0]
return mols
[docs] def run_analysis(self, settings, steprange=None):
"""
Calls the AMS analysis tool behind the scenes
* ``settings`` -- PLAMS Settings object
Example :
>>> settings = Settings()
>>> settings.input.Task = 'AutoCorrelation'
>>> settings.input.AutoCorrelation.Property = 'Velocities'
>>> settings.input.AutoCorrelation.MaxFrame = 2000
* ``steprange`` -- Start frame, end frame, and stepsize. The default (None) corresponds to all frames in the object
Returns a list of AMSAnalysisPlot objects.
Main attributes of the AMSAnalysisPlot objects are :
* ``name`` -- The name of the plot
* ``x`` -- A list of lists containing the values of the coordinate system
If the coordinate system is 1D, then it is a list containing a single list of values
x = [[x1,x2,x3,x4,...,xn]]
* ``y`` -- A list containing the function values
* ``write()`` -- A method returning a string containing all plot info
"""
# First get all the indices
s = slice(steprange)
start, stop, step = s.indices(len(self))
indices = range(start, stop, step)
if steprange is not None:
if len(steprange) == 1:
indices = range(stop, stop + 1)
stepnums = {}
for i in indices:
irkf, istep = self._get_filenum_and_stepnum(i)
if not irkf in stepnums:
stepnums[irkf] = []
stepnums[irkf].append(istep)
ranges = []
for irkf, values in stepnums.items():
ranges.append((values[0] + 1, values[-1] + 1, step))
print("ranges: ", ranges)
# Now completel the settings object
trajecsettings = []
for ikf, rkf in enumerate(self.files):
s = Settings()
s.Trajectory.KFFilename = self.files[0].file_object.path
s.Trajectory.Range = "%i %i %i" % (ranges[irkf][0], ranges[irkf][1], ranges[irkf][2])
trajecsettings.append(s)
settings.input.TrajectoryInfo = trajecsettings
# Run the analysis job and extract the plots
job = AMSAnalysisJob(settings=settings)
result = job.run()
plots = result.get_all_plots()
return plots
#################
# Private methods
#################
def _get_filenum_and_stepnum(self, i):
"""
Connects a filenumber and a stepnumber to index i
"""
irkf, istep = (-1, -1)
counter = 0
for il, length in enumerate(self.lengths):
if i < counter + length:
irkf = il
istep = i - counter
break
counter += length
if irkf < 0:
raise KeyError("Index out of range")
return irkf, istep
def _get_index(self, irkf, istep):
"""
Connects the index i to a filenumber and a stepnumber
"""
return sum([l for l in self.lengths[:irkf]]) + istep
def make_recursively_immutable(obj, verbose=False):
"""
Return a new object that is immutable
Note: Any recursive objects are NOT affected. For example, if obj is an instance of Molecule,
then this function will return an instance of _Molecule, but new_obj.atoms[0].mol is still the original Molecule.
"""
# Find all the branches of this object tree
branchlist = [branch for branch in get_branches(obj)]
if verbose:
print("Tree: ")
for branch in branchlist:
print([n[0] for n in branch])
# Find length of longest branch
maxlength = max([len(branch) for branch in branchlist])
for i in range(maxlength - 1, 0, -1):
if verbose:
print("Level: ", i)
# Find all objects at this level, and theit parent
for branch in branchlist:
if len(branch) > i:
current_obj = branch[i][1]
current_name = branch[i][0]
# Now make immutable. Take care to do this only once for an object
if current_obj.__class__.__name__[0] == "_" or isinstance(current_obj, tuple):
continue
current_obj = make_immutable(current_obj)
parent_obj = branch[i - 1][1]
if isinstance(parent_obj, list):
parent_obj[current_name] = current_obj
elif isinstance(parent_obj, dict):
parent_obj[current_name] = current_obj
else:
parent_obj.__dict__[current_name] = current_obj
# Now do the top layer
current_obj = make_immutable(obj, add_self=True)
return current_obj
def make_immutable(obj, add_self=False):
"""
Make this object immutable
Note: This is specifically written for a Molecule object
"""
def trytosetattr(self, name, value):
"""
A setattr function that does not write attributes
"""
# raise AttributeError("Cannot change object")
if hasattr(self, name):
raise AttributeError("Cannot reassign members")
self.__dict__[name] = value
def trytosetitem(self, name, value):
"""
A setitem function for a dictionary that returns an exception
"""
raise AttributeError("Cannot change object")
def mutable_copy(obj, atoms=None):
"""
Returns a copy of self, but one that is mutable
"""
if hasattr(obj, "__dict__"):
if "_orig" in obj.__dict__:
return obj._orig
else:
return None
elif "_orig" in obj:
return obj["_orig"]
else:
return None
# If the object is a built-in object, all the below does not work
if isinstance(obj, list):
return tuple(obj)
# There will be dictionaries (Settings objects). How to make those immutable?
# Lets get the name of the class
classname = obj.__class__.__name__
# Lets get all the instance variables of the input object
attribs = []
values = []
if hasattr(obj, "__dict__"):
attribs = [key for key in obj.__dict__.keys()]
values = [obj.__dict__[key] for key in attribs]
# Create a new class which is derived from the objects class
ImmutableClass = type("_" + classname, (obj.__class__,), {})
# Now create an instance of the new class with the same instance variables
im_obj = ImmutableClass.__new__(ImmutableClass)
if isinstance(obj, dict):
# This is for a dictionary or PLAMS Settings object
for key, value in obj.items():
im_obj[key] = value
else:
for key, value in zip(attribs, values):
setattr(im_obj, key, value)
# Add its original self, for copying purposes
if add_self:
im_obj._orig = obj
setattr(ImmutableClass, "copy", mutable_copy)
# Now set the method that makes overwriting attributes impossible
if isinstance(obj, dict):
setattr(ImmutableClass, "__setitem__", trytosetitem)
else:
setattr(ImmutableClass, "__setattr__", trytosetattr)
return im_obj
def get_branches(obj, subpath=None):
"""
Search recursively through instance variables of object until an immutable one is met
"""
tree = dictionary_from_object(obj)
if subpath is None:
subpath = (("top", obj),)
# Identify immutable objects and circular objects (they are leafs)
prev_objects = [node[1] for node in subpath]
im_objs = []
for n, o in tree.items():
if is_immutable(o):
im_objs.append(n)
# Avoid circular crap
elif True in [o is po for po in prev_objects]:
im_objs.append(n)
# The branch ends if the tree contains only immutable objects
if len(im_objs) == len(tree):
yield subpath
else:
for n, o in tree.items():
if n in im_objs:
continue
for path in get_branches(o, subpath + ((n, o),)):
yield path
def is_immutable(obj):
"""
Check if this object is immutable
Note: This is a first incomplete implementation
"""
typelist = [int, float, tuple, str, bool] # This does not take into account numpy integers etc
immutable = False
if obj is None:
immutable = True
for t in typelist:
if isinstance(obj, t):
immutable = True
break
return immutable
def dictionary_from_object(obj):
"""
Get the dicrtionary of object attributes
"""
if hasattr(obj, "__dict__"):
tree = obj.__dict__
elif isinstance(obj, dict):
tree = obj
elif hasattr(obj, "__iter__"):
tree = {}
for i, v in enumerate(obj):
tree[i] = v
else:
raise Exception("What is this object?", type(obj))
return tree