def test_dummy_write():
# This writer does not write anything but needs
# to be able to do the dummy actions
with DummyWriter() as w:
w.write(a=1, b=2)
w.add_array('psi', (1, 2))
w.fill(np.ones((1, 2)))
w.sync()
with w.child('child') as w2:
w2.write(c=3)
assert len(w) == 0
def _open(self, filename, mode):
from ase.io.ulm import ulmopen, DummyWriter
if mode not in 'aw':
raise ValueError('mode must be "w" or "a".')
if self.master:
self.backend = ulmopen(filename, mode, tag='ASE-Trajectory')
if len(self.backend) > 0:
r = ulmopen(filename)
self.numbers = r.numbers
self.pbc = r.pbc
else:
self.backend = DummyWriter()
class TrajectoryWriter:
"""Writes Atoms objects to a .traj file."""
def __init__(self, filename, mode='w', atoms=None, properties=None,
extra=[], master=None):
"""A Trajectory writer, in write or append mode.
Parameters:
filename: str
The name of the file. Traditionally ends in .traj.
mode: str
The mode. 'r' is read mode, the file should already exist, and
no atoms argument should be specified.
'w' is write mode. The atoms argument specifies the Atoms
object to be written to the file, if not given it must instead
be given as an argument to the write() method.
'a' is append mode. It acts as write mode, except that
data is appended to a preexisting file.
atoms: Atoms object
The Atoms object to be written in write or append mode.
properties: list of str
If specified, these calculator properties are saved in the
trajectory. If not specified, all supported quantities are
saved. Possible values: energy, forces, stress, dipole,
charges, magmom and magmoms.
master: bool
Controls which process does the actual writing. The
default is that process number 0 does this. If this
argument is given, processes where it is True will write.
"""
if master is None:
master = (world.rank == 0)
self.master = master
self.atoms = atoms
self.properties = properties
self.description = {}
self.numbers = None
self.pbc = None
self.masses = None
self._open(filename, mode)
def set_description(self, description):
self.description.update(description)
def _open(self, filename, mode):
from ase.io.ulm import ulmopen, DummyWriter
if mode not in 'aw':
raise ValueError('mode must be "w" or "a".')
if self.master:
self.backend = ulmopen(filename, mode, tag='ASE-Trajectory')
if len(self.backend) > 0:
r = ulmopen(filename)
self.numbers = r.numbers
self.pbc = r.pbc
else:
self.backend = DummyWriter()
def write(self, atoms=None, **kwargs):
"""Write the atoms to the file.
If the atoms argument is not given, the atoms object specified
when creating the trajectory object is used.
Use keyword arguments to add extra properties::
writer.write(atoms, energy=117, dipole=[0, 0, 1.0])
"""
b = self.backend
if atoms is None:
atoms = self.atoms
if hasattr(atoms, 'interpolate'):
# seems to be a NEB
neb = atoms
assert not neb.parallel or world.size == 1
for image in neb.images:
self.write(image)
return
while hasattr(atoms, 'atoms_for_saving'):
# Seems to be a Filter or similar, instructing us to
# save the original atoms.
atoms = atoms.atoms_for_saving
if len(b) == 0:
b.write(version=1, ase_version=__version__)
if self.description:
b.write(description=self.description)
# Atomic numbers and periodic boundary conditions are only
# written once - in the header. Store them here so that we can
# check that they are the same for all images:
self.numbers = atoms.get_atomic_numbers()
self.pbc = atoms.get_pbc()
else:
if (atoms.pbc != self.pbc).any():
raise ValueError('Bad periodic boundary conditions!')
elif len(atoms) != len(self.numbers):
raise ValueError('Bad number of atoms!')
elif (atoms.numbers != self.numbers).any():
raise ValueError('Bad atomic numbers!')
write_atoms(b, atoms, write_header=(len(b) == 0))
calc = atoms.get_calculator()
if calc is None and len(kwargs) > 0:
calc = SinglePointCalculator(atoms)
if calc is not None:
if not hasattr(calc, 'get_property'):
calc = OldCalculatorWrapper(calc)
c = b.child('calculator')
c.write(name=calc.name)
if hasattr(calc, 'todict'):
d = calc.todict()
if d:
c.write(parameters=d)
for prop in all_properties:
if prop in kwargs:
x = kwargs[prop]
else:
if self.properties is not None:
if prop in self.properties:
x = calc.get_property(prop, atoms)
else:
x = None
else:
try:
x = calc.get_property(prop, atoms,
allow_calculation=False)
except (PropertyNotImplementedError, KeyError):
# KeyError is needed for Jacapo.
x = None
if x is not None:
if prop in ['stress', 'dipole']:
x = x.tolist()
c.write(prop, x)
info = {}
for key, value in atoms.info.items():
try:
encode(value)
except TypeError:
warnings.warn('Skipping "{0}" info.'.format(key))
else:
info[key] = value
if info:
b.write(info=info)
b.sync()
def close(self):
"""Close the trajectory file."""
self.backend.close()
def __len__(self):
return world.sum(len(self.backend))
