Source code for scm.plams.trajectories.rkfhistoryfile

#!/usr/bin/env python

import numpy
from import PT
from ..mol.molecule import Molecule
from ..mol.atom import Atom
from ..core.settings import Settings
from ..core.errors import PlamsError
from .rkffile import RKFTrajectoryFile
from .rkffile import bohr_to_angstrom
from .rkffile import write_general_section

__all__ = ['RKFHistoryFile']

[docs]class RKFHistoryFile (RKFTrajectoryFile) : """ Class representing an RKF file containing a molecular simulation history with varying numbers of atoms An instance of this class has the following attributes: * ``file_object`` -- A PLAMS |KFFile| object, referring to the actual RKF file * ``position`` -- The frame to which the cursor is currently pointing in the RKF file * ``mode`` -- Designates whether the file is in read or write mode ('rb' or 'wb') * ``elements`` -- The elements of the atoms in the system at the current frame * ``conect`` -- The connectivity information of the current frame * ``mddata`` -- Read mode only: A dictionary containing data from the MDHistory section in the RKF file * ``read_lattice``-- Read mode only: Wether the lattice vectors will be read from the file * ``read_bonds`` -- Wether the connectivity information will be read from the file An |RKFHistoryFile| object behaves very similar to a regular file object. It has read and write methods (:meth:`read_next` and :meth:`write_next`) that read and write from/to the position of the cursor in the ``file_object`` attribute. If the file is in read mode, an additional method :meth:`read_frame` can be used that moves the cursor to any frame in the file and reads from there. The amount of information stored in memory is kept to a minimum, as only information from the latest frame is ever stored. Reading and writing to and from the files can be done as follows:: >>> from scm.plams import RKFHistoryFile >>> rkf = RKFHistoryFile('ams.rkf') >>> mol = rkf.get_plamsmol() >>> rkfout = RKFHistoryFile('new.rkf',mode='wb') >>> for i in range(rkf.get_length()) : >>> crd,cell = rkf.read_frame(i,molecule=mol) >>> rkfout.write_next(molecule=mol) >>> rkfout.close() The above script reads information from the RKF file ``ams.rkf`` into the |Molecule| object ``mol`` in a step-by-step manner.. The |Molecule| object is then passed to the :meth:`write_next` method of the new |RKFHistoryFile| object corresponding to the new rkf file ``new.rkf``. The exact same result can also be achieved by iterating over the instance as a callable >>> rkf = RKFHistoryFile('ams.rkf') >>> mol = rkf.get_plamsmol() >>> rkfout = RKFHistoryFile('new.rkf',mode='wb') >>> for crd,cell in rkf(mol) : >>> rkfout.write_next(molecule=mol) >>> rkfout.close() This procedure requires all coordinate information to be passed to and from the |Molecule| object for each frame, which can be time-consuming. Some time can be saved by bypassing the |Molecule| object:: >>> rkf = RKFHistoryFile('ams.rkf') >>> rkfout = RKFHistoryFile('new.rkf',mode='wb') >>> for crd,cell in rkf : >>> rkfout.write_next(coords=crd,cell=cell,elements=rkf.elements,conect=rkf.conect) >>> rkfout.close() The only mandatory argument to the :meth:`write_next` method is ``coords``. Further time can be saved by setting the ``read_lattice`` and ``read_bonds`` variables to False. By default the write mode will create a minimal version of the RKF file, containing only elements, coordinates, lattice, and connectivity information. This minimal file format can be read by AMSMovie. If the original RKF file contains an MDHistory section (if it resulted from a MolecularGun simulation) it is possible to store the information from that section and write it to another file. To enable this, the method :meth:`store_mddata` needs to be called after creation, and a dictionary of mddata needs to be passed to the :meth:`write_next` method. When that is done, the AMS trajectory analysis tools can be used on the file. Restarting an MD run with such a file is however currently not possible:: >>> rkf = RKFHistoryFile('ams.rkf') >>> rkf.store_mddata() >>> mol = rkf.get_plamsmol() >>> rkf_out = RKFHistoryFile('new.rkf',mode='wb') >>> rkf_out.store_mddata(rkf) >>> for i in range(rkf.get_length()) : >>> crd,cell = rkf.read_frame(i,molecule=mol) >>> rkf_out.write_next(molecule=mol,mddata=rkf.mddata) >>> rkf_out.close() """
[docs] def __init__ (self, filename, mode='rb', fileobject=None, ntap=None) : """ Initializes the RKFHistoryFile object * ``filename`` -- The path to the RKF file * ``mode`` -- The mode in which to open the RKF file ('rb' or 'wb') * ``fileobject`` -- Optionally, a file object can be passed instead (filename needs to be set to None) * ``ntap`` -- If the file is in write mode, the number of atoms can be passed here """ self.added_atoms = None self.removed_atoms = None self.chemical_systems = None RKFTrajectoryFile.__init__(self,filename,mode,fileobject,ntap) self.input_elements = self.elements[:] self.versionhistory_length = 0 self.system_version_elements = {} self.frame = -1 self.props = None # This is for writing only self.system_version_props = {} self.version_history_items = []
def set_elements (self, elements) : """ Sets the elements attribute (needed in write mode). * ``elements`` -- A list containing the element name of each atom """ if self.position > 0 : raise PLAMSError('Elements should not be changed while reading/writing is already in progress') RKFTrajectoryFile.set_elements(self, elements) self.input_elements = elements
[docs] def get_plamsmol (self) : """ Extracts a PLAMS molecule object from the RKF file """ section_dict = self.file_object.read_section('ChemicalSystem(1)') if len(section_dict) == 0 : section_dict = self.file_object.read_section('InputMolecule') if len(section_dict) == 0 : section_dict = self.file_object.read_section('Molecule') plamsmol = Molecule._mol_from_rkf_section(section_dict) return plamsmol
[docs] def _rewrite_molecule (self) : """ Overwrite the molecule section with the latest frame (called in close()) """ mol = self.get_plamsmol() crd,cell = self.read_last_frame(molecule=mol) self._write_molecule_section(crd,cell,molecule=mol)
def _read_header (self) : """ Read the start molecule data from the InputMolecule section (not the Molecule section) """ self.added_atoms = {} self.removed_atoms = {} self.chemical_systems = {0:1} secname = 'ChemicalSystem(1)' if self.file_object.reader._sections is None : self.file_object.reader._create_index() if not 'SystemVersionHistory' in self.file_object.reader._sections : secname = 'InputMolecule' RKFTrajectoryFile._read_header (self, molecule_section=secname) # Now store the added and removed atoms along the trajectory # (This might be slow?) # I could also do it on the fly, but that may be messy when we move back and forth through the file if not 'SystemVersionHistory' in self.file_object.reader._sections : return #version = 0 version = 1 self._set_system_version_elements() for i in range(self.get_length()) : new_version ='History','SystemVersion(%i)'%(i+1)) if new_version == version : continue self.added_atoms[i] = {} self.removed_atoms[i] = {} # Now look for the added and removed atoms removed_atoms = [] if 'RemovedAtoms(%i)'%(new_version) in self.file_object.reader._sections['SystemVersionHistory'] : removed_atoms ='SystemVersionHistory','RemovedAtoms(%i)'%(new_version)) if not isinstance(removed_atoms,list) : removed_atoms = [removed_atoms] added_atoms = [] if 'AddedAtoms(%i)'%(new_version) in self.file_object.reader._sections['SystemVersionHistory'] : added_atoms ='SystemVersionHistory','AddedAtoms(%i)'%(new_version)) if not isinstance(added_atoms,list) : added_atoms = [added_atoms] # Now find the corresponding elements chemSysNum ='SystemVersionHistory','SectionNum(%i)'%(new_version)) # Compare to the previous chemical system #prev_version = new_version - 1 #prevChemSysNum ='SystemVersionHistory','SectionNum(%i)'%(prev_version)) prevChemSysNum = self.chemical_systems[max(self.chemical_systems.keys())] sectionname = 'ChemicalSystem(%i)'%(prevChemSysNum) prev_elements = [PT.get_symbol(atnum) for atnum in,'AtomicNumbers')] sectionname = 'ChemicalSystem(%i)'%(chemSysNum) elements = [PT.get_symbol(atnum) for atnum in,'AtomicNumbers')] ################ # Read badly written files diff = (len(added_atoms) - len(removed_atoms)) - (len(elements)-len(prev_elements)) if diff != 0 : chemSysNum, elements = self._correct_chemical_system(elements,prev_elements,added_atoms,removed_atoms) #print ('chemSysNum: ',i,chemSysNum, len(elements),added_atoms,removed_atoms) #################### removed_elements = [prev_elements[i-1].split('.')[0] for i in removed_atoms] added_elements = [elements[i-1].split('.')[0] for i in added_atoms] # Now store the elements for iat,el in zip(removed_atoms,removed_elements) : self.removed_atoms[i][iat] = el for iat,el in zip(added_atoms,added_elements) : self.added_atoms[i][iat] = el self.chemical_systems[i] = chemSysNum version = new_version
[docs] def _correct_chemical_system (self, elements, prev_elements, added_atoms, removed_atoms) : """ Check if the referenced chemical system is correct, and if not, find one matching added/removed atoms """ #print ('Searching for chemical system elsewhere') new_elements = [el for el in prev_elements] for ra in reversed(removed_atoms) : new_elements.pop(ra-1) for aa in added_atoms : new_elements.insert(aa-1,'*') # Add a wildcard chemSysNum = self._check_for_chemical_system(new_elements) if chemSysNum is None : print ('Chemical system does not exist') chemSysNum = self._check_for_chemical_system(new_elements,compare_elements=False) if chemSysNum is None : raise Exception('Chemical system with this number of atoms does not exist') sectionname = 'ChemicalSystem(%i)'%(chemSysNum) elements = [PT.get_symbol(atnum) for atnum in,'AtomicNumbers')] ################### # Check the regions section_dict = self.file_object.read_section(sectionname) plamsmol = Molecule._mol_from_rkf_section(section_dict) print (0,plamsmol.atoms[0].properties.suffix) for aa in added_atoms : print (aa-1,plamsmol.atoms[aa-1].properties.suffix) ################### return chemSysNum, elements
# FIXME: The _write_header section writes the starting molecule to the Molecule section, # not the final molecule (like the Fortran code does) def _write_header (self, coords, cell, molecule=None) : """ Write Molecule info to file (elements, periodicity) """ # First write the general section write_general_section(self.file_object) # Then write the input molecule self._update_celldata(cell) self._write_molecule_section(coords, cell, molecule=molecule) # I think the InputMolecule is mandatory in this case self._write_molecule_section(coords, cell, section='InputMolecule', molecule=molecule) self.chemical_systems = {} if self.include_mddata : # Start setting up the MDHistory section as well self.file_object.write('MDHistory','blockSize',100) self.added_atoms = {} self.removed_atoms = {} def _read_coordinates (self, i, molecule, cell, bonds) : """ Read the coordinates at step i """ coords = self.coords.reshape(len(self.coords)*3) update_molecule = False elements = self._read_elements_for_frame(i) if elements != self.elements : update_molecule = True if isinstance(molecule,Molecule) : _, _, mol_elements, _, props = self._read_plamsmol(molecule) if mol_elements != self.elements or props != self.props : # This molecule has nothing to do with the previously read one update_molecule = True prev_frames = [iframe for iframe in self.chemical_systems.keys() if iframe<=i] ifr = 0 if len(prev_frames) > 0 : ifr = prev_frames[-1] next_frames = [iframe for iframe in self.chemical_systems.keys() if iframe>i] jfr = len(self) if len(next_frames) > 0 : jfr = next_frames[0] # If the previous frame falls before the last chemical system or after the next one, update if self.frame < ifr or self.frame >= jfr : update_molecule = True if update_molecule : self.elements = elements self.frame = i self.coords = numpy.zeros((len(elements),3)) coords = self.coords.reshape((len(elements)*3)) # Rebuild the molecule (bonds will disappear for now) if isinstance(molecule,Molecule) : self.props = props secname = 'ChemicalSystem(%i)'%(self.chemical_systems[ifr]) section_dict = self.file_object.read_section(secname) new_mol = Molecule._mol_from_rkf_section(section_dict) for at in reversed(molecule.atoms) : molecule.delete_atom(at) = Settings() for iel,el in enumerate(elements) : atom = new_mol.atoms[iel] #atom = Atom(PT.get_atomic_number(el)) molecule.add_atom(atom) # Now what if the elements found in this ChemicalSystem section # do not match the expected elements (self.elements)? new_elements = [at.symbol for at in molecule.atoms] if new_elements != self.elements : print ('Chemical section does not match',i) atnums = [PT.get_atomic_number(el) for el in self.elements] for iat,atnum in enumerate(atnums) : molecule.atoms[iat].atnum = atnum coords[:] ='History', 'Coords(%i)'%(i+1)) coords *= bohr_to_angstrom # This changes self.coords behind the scenes # Assign the data to the molecule object if isinstance(molecule,Molecule) : cell_reduced = None if cell is not None : cell_reduced = cell[:self.nvecs] self._set_plamsmol(self.coords,cell_reduced,molecule,bonds) def _read_elements_for_frame (self, frame) : """ Use the added and removed atoms to read the elements at each frame """ elements = self.input_elements[:] for i in range(frame+1) : if i in self.removed_atoms : # Insert them into the elements list, in the correct place elements = [el for iat,el in enumerate(elements) if not iat+1 in self.removed_atoms[i]] if i in self.added_atoms : for iat,el in self.added_atoms[i].items() : elements.insert(iat-1,el) return elements
[docs] def write_next (self, coords=None, molecule=None, elements=None, cell=[0.,0.,0.], conect=None, historydata=None, mddata=None) : """ Write frame to next position in trajectory file * ``coords`` -- A list or numpy array of (``ntap``,3) containing the system coordinates * ``molecule`` -- A molecule object to read the molecular data from * ``elements`` -- The element symbols of the atoms in the system * ``cell`` -- A set of lattice vectors (or cell diameters for an orthorhombic system) * ``conect`` -- A dictionary containing the connectivity info (e.g. {1:[2],2:[1]}) * ``historydata`` -- A dictionary containing additional variables to be written to the History section * ``mddata`` -- A dictionary containing the variables to be written to the MDHistory section The ``mddata`` dictionary can contain the following keys: ('TotalEnergy', 'PotentialEnergy', 'Step', 'Velocities', 'KineticEnergy', 'Charges', 'ConservedEnergy', 'Time', 'Temperature') The ``historydata`` dictionary can contain for example: ('Energy','Gradients','StressTensor') All values must be in atomic units Numpy arrays or lists of lists will be flattened before they are written to the file .. note:: Either ``coords`` and ``elements`` or ``molecule`` are mandatory arguments """ # Check for common error in the arguments if coords is not None : if isinstance(coords,Molecule) : raise PlamsError('The PLAMS molecule needs to be passed as the second argument (molecule)') props = None if isinstance(molecule,Molecule) : coords, cell, elements, conect, props = self._read_plamsmol(molecule) # Make sure that the cell consists of three vectors cell = self._convert_cell(cell) if conect is not None : if len(conect) == 0 : conect = None self.conect = conect # If this is the first step, write the header if self.position == 0 : self.elements = elements self.props = props self._write_header(coords,cell,molecule=molecule) # This is specific for the history-file object self.input_elements = elements[:] # Define some local variables counter = 1 step = self.position if mddata is not None : if 'Step' in mddata : step = mddata['Step'] # Energy should be read from mddata first, otherwise from historydata, otherwise set to zero energy = self._set_energy(mddata, historydata) if not self.include_historydata : historydata = {} historydata['Energy'] = energy # Write the history section counter = self._write_history_entry(step, coords, cell, conect, historydata) if self.include_mddata and mddata is not None : self._write_mdhistory_entry(mddata) # If a change took place, write it. # Note: The chemical system is only written if the elements change. Changes in the atom properties are not checked if elements != self.elements or props != self.props or self.position==0 : # version history has to be checked before the elements can be updated chemsysversion = self._check_for_chemical_system(elements, props=props) self._write_version_history(elements, coords, cell, props) # Now update the elements and props self.elements = elements self.props = props # Write the molecule sections if chemsysversion is None : chemsysversion = len(self.system_version_elements) self._write_molecule_section (coords, cell, section='ChemicalSystem(%i)'%(chemsysversion), molecule=molecule) self.chemical_systems[self.position] = chemsysversion counter = self._write_system_version_history_entry(counter) self.position += 1 if self.saving_freq is not None : if self.position%self.saving_freq == 0 :
def _write_version_history (self, elements, coords, cell, props) : """ Write the version history """ # Enter the correct SystemVersionHistory entry into the History section self.versionhistory_length += 1 version = self.versionhistory_length # Find the corresponding ChemicalSystem and write it chemsysversion = self._check_for_chemical_system(elements, props=props) if chemsysversion is None : #self.system_version_elements.append(elements[:]) if len(self.system_version_elements) in self.system_version_elements.keys() : raise Exception('self.system_version_elements is not consecutively numbered') self.system_version_elements[len(self.system_version_elements)] = elements[:] self.system_version_props[len(self.system_version_elements)] = props chemsysversion = len(self.system_version_elements) # Now add an entry to SystemVersionHistory added_atoms, removed_atoms = self._find_system_change(elements, props) self.file_object.write('SystemVersionHistory','nEntries',version) self.file_object.write('SystemVersionHistory','currentEntryOpen',False) if 'SectionNum' not in self.version_history_items : self.version_history_items.append('SectionNum') svh_counter = self.version_history_items.index('SectionNum') + 1 self._write_keydata_in_history('SectionNum',svh_counter,False,1,version,chemsysversion,'SystemVersionHistory') # Avoid writing empty data into this section by storing the items. if len(added_atoms) > 0 : if 'AddedAtoms' not in self.version_history_items : self.version_history_items.append('AddedAtoms') svh_counter = self.version_history_items.index('AddedAtoms') + 1 data = [iat+1 for iat in added_atoms] self._write_keydata_in_history('AddedAtoms',svh_counter,False,len(data),version,data,'SystemVersionHistory') if len(removed_atoms) > 0 : if 'RemovedAtoms' not in self.version_history_items : self.version_history_items.append('RemovedAtoms') svh_counter = self.version_history_items.index('RemovedAtoms') + 1 data = [iat+1 for iat in removed_atoms] self._write_keydata_in_history('RemovedAtoms',svh_counter,False,len(data),version,data,'SystemVersionHistory') def _write_system_version_history_entry (self, counter) : """ Write the entry for the SystemVersionHistory into the History section """ version = self.versionhistory_length self._write_keydata_in_history('SystemVersion', counter, False, 1, self.position+1, version) counter += 1 return counter def _find_system_change (self, elements, props) : """ Find out which atoms were added and/or deleted """ # First find out which elements were removed removed_atoms = [] position = 0 for i,el in enumerate(self.elements) : if position == len(elements) : removed_atoms += [j for j in range(i,len(self.elements))] break # Get the props, if relevant if self.props is None : op = None else : op = self.props[i] if props is None : p = None else : p = props[position] # Find the changes if el == elements[position] and op==p : position += 1 else : removed_atoms.append(i) # Then find out which elements were added added_atoms = [i for i in range(position,len(elements))] # Now store them (this is actually not really necessary) if len(removed_atoms) > 0 : self.removed_atoms[self.position] = {} if len(added_atoms) > 0 : self.added_atoms[self.position] = {} for iat in removed_atoms : self.removed_atoms[self.position][iat+1] = self.elements[iat] for iat in added_atoms : self.added_atoms[self.position][iat+1] = elements[iat] return added_atoms, removed_atoms def _check_for_chemical_system (self, elements, compare_elements=True, props=None) : """ Check if the new chemical system was encountered before """ iterator = self.system_version_elements.items() version = None for ind,(i,prev_elements) in enumerate(iterator) : equal = True if len(prev_elements) != len(elements) : continue if compare_elements : for el,pel in zip(elements,prev_elements) : if el != pel : if el == '*' : continue equal = False break if props is not None and equal : # This only occurs when writing, otherwise system_version_props is empty if i in self.system_version_props.keys() : if props != self.system_version_props[i] : equal = False if equal : version = i+1 return version
[docs] def _set_system_version_elements (self) : """ Store all chemical systems from the file """ self.system_version_elements = {} keys = [key for key in self.file_object.reader._sections.keys() if 'ChemicalSystem' in key] nums = [int(k.split('(')[1].split(')')[0])-1 for k in keys] for num,key in zip(nums,keys) : elements = [PT.get_symbol(atnum) for atnum in,'AtomicNumbers')] self.system_version_elements[num] = elements