forked from lijiext/lammps
379 lines
14 KiB
Python
379 lines
14 KiB
Python
"""Classes to deal with output of simulation data.
|
|
|
|
Copyright (C) 2013, Joshua More and Michele Ceriotti
|
|
|
|
This program is free software: you can redistribute it and/or modify
|
|
it under the terms of the GNU General Public License as published by
|
|
the Free Software Foundation, either version 3 of the License, or
|
|
(at your option) any later version.
|
|
|
|
This program is distributed in the hope that it will be useful,
|
|
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
GNU General Public License for more details.
|
|
|
|
You should have received a copy of the GNU General Public License
|
|
along with this program. If not, see <http.//www.gnu.org/licenses/>.
|
|
|
|
|
|
Holds classes to deal with the output of different properties, trajectories
|
|
and the restart files.
|
|
|
|
Classes:
|
|
PropertyOutput: Deals with outputting properties.
|
|
TrajectoryOutput: Deals with outputting trajectories.
|
|
CheckpointOutput: Deals with outputting restart files.
|
|
"""
|
|
|
|
import os
|
|
import numpy as np
|
|
import ipi.inputs.simulation
|
|
from ipi.utils.depend import *
|
|
from ipi.utils.io.io_xml import *
|
|
from ipi.engine.properties import getkey
|
|
|
|
__all__ = [ 'PropertyOutput', 'TrajectoryOutput', 'CheckpointOutput' ]
|
|
|
|
class PropertyOutput(dobject):
|
|
"""Class dealing with outputting a set of properties to file.
|
|
|
|
Does not do any calculation, just manages opening a file, getting data
|
|
from a Properties object and outputting with the desired stride.
|
|
|
|
Attributes:
|
|
filename: The name of the file to output to.
|
|
outlist: A list of the properties to be output.
|
|
stride: The number of steps that should be taken between outputting the
|
|
data to file.
|
|
flush: How often we should flush to disk.
|
|
nout: Number of steps since data was last flushed.
|
|
out: The output stream on which to output the properties.
|
|
simul: The simulation object to get the data to be output from.
|
|
"""
|
|
|
|
|
|
def __init__(self, filename="out", stride=1, flush=1, outlist=None):
|
|
"""Initializes a property output stream opening the corresponding
|
|
file name.
|
|
|
|
Also writes out headers.
|
|
|
|
Args:
|
|
filename: A string giving the name of the file to be output to.
|
|
stride: An integer giving how many steps should be taken between
|
|
outputting the data to file.
|
|
flush: Number of writes to file between flushing data.
|
|
outlist: A list of all the properties that should be output.
|
|
"""
|
|
|
|
if outlist is None:
|
|
outlist = np.zeros(0,np.dtype('|S1024'))
|
|
self.filename = filename
|
|
self.outlist = np.asarray(outlist,np.dtype('|S1024'))
|
|
self.stride = stride
|
|
self.flush = flush
|
|
self.nout = 0
|
|
self.out = None
|
|
|
|
def bind(self, simul):
|
|
"""Binds output proxy to simulation object.
|
|
|
|
Args:
|
|
simul: A simulation object to be bound.
|
|
"""
|
|
|
|
self.simul = simul
|
|
|
|
# Checks as soon as possible if some asked-for properties are
|
|
# missing or mispelled
|
|
for what in self.outlist:
|
|
key = getkey(what)
|
|
if not key in self.simul.properties.property_dict.keys():
|
|
print "Computable properties list: ", self.simul.properties.property_dict.keys()
|
|
raise KeyError(key + " is not a recognized property")
|
|
|
|
self.open_stream()
|
|
|
|
def open_stream(self):
|
|
"""Opens the output stream."""
|
|
|
|
try:
|
|
self.out = open(self.filename, "a")
|
|
except:
|
|
raise ValueError("Could not open file " + self.filename + " for output")
|
|
|
|
# print nice header if information is available on the properties
|
|
if (self.simul.step == 0) :
|
|
icol = 1
|
|
for what in self.outlist:
|
|
ohead = "# "
|
|
key = getkey(what)
|
|
prop = self.simul.properties.property_dict[key]
|
|
|
|
if "size" in prop and prop["size"] > 1:
|
|
ohead += "cols. %3d-%-3d" % ( icol, icol+prop["size"] - 1 )
|
|
icol += prop["size"]
|
|
else:
|
|
ohead += "column %3d " % ( icol )
|
|
icol += 1
|
|
ohead += " --> %s " % (what)
|
|
if "help" in prop:
|
|
ohead += ": " + prop["help"]
|
|
self.out.write(ohead + "\n")
|
|
|
|
def close_stream():
|
|
"""Closes the output stream."""
|
|
|
|
self.out.close()
|
|
|
|
def write(self):
|
|
"""Outputs the required properties of the system.
|
|
|
|
Note that properties are outputted using the same format as for the
|
|
output to the xml checkpoint files, as specified in io_xml.
|
|
|
|
Raises:
|
|
KeyError: Raised if one of the properties specified in the output list
|
|
are not contained in the property_dict member of properties.
|
|
"""
|
|
|
|
if not (self.simul.step + 1) % self.stride == 0:
|
|
return
|
|
self.out.write(" ")
|
|
for what in self.outlist:
|
|
try:
|
|
quantity = self.simul.properties[what]
|
|
except KeyError:
|
|
raise KeyError(what + " is not a recognized property")
|
|
if not hasattr(quantity,"__len__") :
|
|
self.out.write(write_type(float, quantity) + " ")
|
|
else:
|
|
for el in quantity:
|
|
self.out.write(write_type(float, el) + " ")
|
|
|
|
self.out.write("\n")
|
|
|
|
self.nout += 1
|
|
if self.flush > 0 and self.nout >= self.flush :
|
|
self.out.flush()
|
|
os.fsync(self.out) # we REALLY want to print out! pretty please OS let us do it.
|
|
self.nout = 0
|
|
|
|
|
|
class TrajectoryOutput(dobject):
|
|
"""Class dealing with outputting atom-based properties as a
|
|
trajectory file.
|
|
|
|
Does not do any calculation, just manages opening a file, getting data
|
|
from a Trajectories object and outputting with the desired stride.
|
|
|
|
Attributes:
|
|
filename: The (base) name of the file to output to.
|
|
format: The format of the trajectory file to be created.
|
|
what: The trajectory that needs to be output.
|
|
stride: The number of steps that should be taken between outputting the
|
|
data to file.
|
|
out: The output stream on which to output the trajectories.
|
|
flush: How often we should flush to disk.
|
|
nout: Number of steps since data was last flushed.
|
|
ibead: Index of the replica to print the trajectory of.
|
|
cell_units: The units that the cell parameters are given in.
|
|
simul: The simulation object to get the data to be output from.
|
|
"""
|
|
|
|
def __init__(self, filename="out", stride=1, flush=1, what="", format="xyz", cell_units="atomic_unit", ibead=-1):
|
|
""" Initializes a property output stream opening the corresponding
|
|
file name.
|
|
|
|
Also writes out headers.
|
|
|
|
Args:
|
|
filename: A string giving the name of the file to be output to.
|
|
stride: An integer giving how many steps should be taken between
|
|
outputting the data to file.
|
|
flush: How often we should flush to disk
|
|
what: A string specifying what trajectory should be output.
|
|
format: A string specifying the type of trajectory file to be created.
|
|
cell_units: A string specifying the units that the cell parameters are
|
|
given in.
|
|
ibead: If positive, prints out only the selected bead. If negative, prints out one file per bead.
|
|
"""
|
|
|
|
self.filename = filename
|
|
self.what = what
|
|
self.stride = stride
|
|
self.flush = flush
|
|
self.ibead = ibead
|
|
self.format = format
|
|
self.cell_units = cell_units
|
|
self.out = None
|
|
self.nout = 0
|
|
|
|
def bind(self, simul):
|
|
"""Binds output proxy to simulation object.
|
|
|
|
Args:
|
|
simul: A simulation object to be bound.
|
|
"""
|
|
|
|
self.simul = simul
|
|
|
|
# Checks as soon as possible if some asked-for trajs are missing or mispelled
|
|
key = getkey(self.what)
|
|
if not key in self.simul.trajs.traj_dict.keys():
|
|
print "Computable trajectories list: ", self.simul.trajs.traj_dict.keys()
|
|
raise KeyError(key + " is not a recognized output trajectory")
|
|
|
|
self.open_stream()
|
|
|
|
def open_stream(self):
|
|
"""Opens the output stream(s)."""
|
|
|
|
if getkey(self.what) in [ "positions", "velocities", "forces", "extras" ]:
|
|
# must write out trajectories for each bead, so must create b streams
|
|
self.out = []
|
|
for b in range(self.simul.beads.nbeads):
|
|
# zero-padded bead number
|
|
padb = ( ("%0" + str(int(1 + np.floor(np.log(self.simul.beads.nbeads)/np.log(10)))) + "d") % (b) )
|
|
try:
|
|
if (self.ibead < 0 or self.ibead == b):
|
|
if getkey(self.what) == "extras":
|
|
self.out.append(open(self.filename + "_" + padb, "a"))
|
|
else:
|
|
self.out.append(open(self.filename + "_" + padb + "." + self.format, "a"))
|
|
else:
|
|
self.out.append(None) # creates null outputs if a
|
|
# single bead output is chosen
|
|
except:
|
|
raise ValueError("Could not open file " + self.filename + "_" + padb + "." + self.format + " for output")
|
|
else:
|
|
try:
|
|
self.out = ( open(self.filename + "." + self.format, "a") )
|
|
except:
|
|
raise ValueError("Could not open file " + self.filename + "." + self.format + " for output")
|
|
|
|
def close_stream():
|
|
"""Closes the output stream."""
|
|
|
|
if hasattr(self.out, "__getitem__"):
|
|
for o in self.out:
|
|
o.close()
|
|
else:
|
|
self.out.close()
|
|
|
|
def write(self):
|
|
"""Writes out the required trajectories."""
|
|
|
|
if not (self.simul.step + 1) % self.stride == 0:
|
|
return
|
|
|
|
doflush = False
|
|
self.nout += 1
|
|
if self.flush > 0 and self.nout >= self.flush :
|
|
doflush = True
|
|
self.nout = 0
|
|
|
|
# quick-and-dirty way to check if a trajectory is "global" or per-bead
|
|
# Checks to see if there is a list of files or just a single file.
|
|
if hasattr(self.out, "__getitem__"):
|
|
if self.ibead < 0:
|
|
for b in range(len(self.out)):
|
|
self.simul.trajs.print_traj(self.what, self.out[b], b, format=self.format, cell_units=self.cell_units, flush=doflush)
|
|
elif self.ibead < len(self.out):
|
|
self.simul.trajs.print_traj(self.what, self.out[self.ibead], self.ibead, format=self.format, cell_units=self.cell_units, flush=doflush)
|
|
else:
|
|
raise ValueError("Selected bead index " + str(self.ibead) + " does not exist for trajectory " + self.what)
|
|
else:
|
|
self.simul.trajs.print_traj(self.what, self.out, b=0, format=self.format, cell_units=self.cell_units, flush=doflush)
|
|
|
|
|
|
class CheckpointOutput(dobject):
|
|
"""Class dealing with outputting checkpoints.
|
|
|
|
Saves the complete status of the simulation at regular intervals.
|
|
|
|
Attributes:
|
|
filename: The (base) name of the file to output to.
|
|
step: the number of times a checkpoint has been written out.
|
|
stride: The number of steps that should be taken between outputting the
|
|
data to file.
|
|
overwrite: If True, the checkpoint file is overwritten at each output.
|
|
If False, will output to 'filename_step'. Note that no check is done
|
|
on whether 'filename_step' exists already.
|
|
simul: The simulation object to get the data to be output from.
|
|
status: An input simulation object used to write out the checkpoint file.
|
|
"""
|
|
|
|
|
|
def __init__(self, filename="restart", stride=1000, overwrite=True, step=0):
|
|
"""Initializes a checkpoint output proxy.
|
|
|
|
Args:
|
|
filename: A string giving the name of the file to be output to.
|
|
stride: An integer giving how many steps should be taken between
|
|
outputting the data to file.
|
|
overwrite: If True, the checkpoint file is overwritten at each output.
|
|
If False, will output to 'filename_step'. Note that no check is done
|
|
on whether 'filename_step' exists already.
|
|
step: The number of checkpoint files that have been created so far.
|
|
"""
|
|
|
|
self.filename = filename
|
|
self.step = step
|
|
self.stride = stride
|
|
self.overwrite = overwrite
|
|
|
|
def bind(self, simul):
|
|
"""Binds output proxy to simulation object.
|
|
|
|
Args:
|
|
simul: A simulation object to be bound.
|
|
"""
|
|
|
|
self.simul = simul
|
|
self.status = ipi.inputs.simulation.InputSimulation()
|
|
self.status.store(simul)
|
|
|
|
def store(self):
|
|
"""Stores the current simulation status.
|
|
|
|
Used so that, if halfway through a step a kill signal is received,
|
|
we can output a checkpoint file corresponding to the beginning of the
|
|
current step, which is the last time that both the velocities and
|
|
positions would have been consistent.
|
|
"""
|
|
|
|
self.status.store(self.simul)
|
|
|
|
def write(self, store=True):
|
|
"""Writes out the required trajectories.
|
|
|
|
Used for both the checkpoint files and the soft-exit restart file.
|
|
We have slightly different behaviour for these two different types of
|
|
checkpoint file, as the soft-exit files have their store() function
|
|
called automatically, and we do not want this to be updated as the
|
|
status of the simulation after a soft-exit call is unlikely to be in
|
|
a consistent state. On the other hand, the standard checkpoint files
|
|
are not automatically updated in this way, and we must manually store the
|
|
current state of the system before writing them.
|
|
|
|
Args:
|
|
store: A boolean saying whether the state of the system should be
|
|
stored before writing the checkpoint file.
|
|
"""
|
|
|
|
if not (self.simul.step + 1) % self.stride == 0:
|
|
return
|
|
|
|
if self.overwrite:
|
|
filename = self.filename
|
|
else:
|
|
filename = self.filename + "_" + str(self.step)
|
|
|
|
if store:
|
|
self.step += 1 # advances the step counter before saving, so next time the correct index will be loaded.
|
|
self.store()
|
|
check_file = open(filename, "w")
|
|
check_file.write(self.status.write(name="simulation"))
|
|
check_file.close()
|