def main_paste(args):
"""
Correlate particles properties from trajectory files.
Example:
--------
trj.py paste.py file1.xyz:radius file2.xyz.voronoi.xyz:volume
"""
from atooms import trajectory as trj
from atooms.core.utils import tipify
f1, attr1 = args.file_inp[0].split(':')
f2, attr2 = args.file_inp[1].split(':')
if args.inp is None:
fmt1, fmt2 = None, None
else:
fmt1, fmt2 = args.inp.split(',')
t1 = trj.Trajectory(f1, fmt=fmt1)
t2 = trj.Trajectory(f2, fmt=fmt2)
# Define slice.
# We interpret --first N --last N as a request of step N
if args.last == args.first and args.last is not None:
args.last += 1
sl1 = fractional_slice(args.first, args.last, args.skip, len(t1))
sl2 = fractional_slice(args.first, args.last, args.skip, len(t2))
# Here we could you a trajectory slice t[sl] but this will load
# everything in ram (getitem doesnt provide a generator). This
# will be fixed with python 3.
ts1 = trajectory.Sliced(t1, sl1)
ts2 = trajectory.Sliced(t2, sl2)
def array_fmt(arr):
"""Remove commas and [] from numpy array repr."""
# Passing a scalar will trigger an error (gotcha: even
# when casting numpy array to list, the elements remain of
# numpy type and this function gets called! (4% slowdown)
_fmt = '%g'
try:
return ' '.join([_fmt % x for x in arr])
except:
return _fmt % arr
# except:
# return numpy.array2string(arr, precision=self.precision, separator=',')[1:-1]
import numpy
numpy.set_string_function(array_fmt, repr=False)
for step, s1, s2 in trj.utils.paste(ts1, ts2):
try:
for i in range(len(s1.particle)):
print(getattr(s1.particle[i], attr1),
getattr(s2.particle[i], attr2))
except:
print(getattr(s1, attr1), getattr(s2, attr2))
def _get_trajectories(input_files, args):
from atooms.trajectory import Sliced
from atooms.core.utils import fractional_slice
for input_file in input_files:
with Trajectory(input_file, fmt=args['fmt']) as th:
if args['center']:
th.add_callback(center)
# Caching is useful for systems with multiple species but
# it will increase the memory footprint. Use --no-cache to
# disable it
if not args['no_cache']:
th.cache = True
if args['species_layout'] is not None:
th.register_callback(change_species, args['species_layout'])
sl = fractional_slice(args['first'], args['last'], args['skip'],
len(th))
if th.block_size > 1:
sl_start = (
sl.start // th.block_size
) * th.block_size if sl.start is not None else sl.start
sl_stop = (
sl.stop // th.block_size
) * th.block_size if sl.stop is not None else sl.stop
sl = slice(sl_start, sl_stop, sl.step)
if sl != slice(None, None, 1):
ts = Sliced(th, sl)
else:
ts = th
yield ts
def scatter(args):
"""
Write frames in trajectory to individual files of the same file format
"""
from atooms import trajectory as trj
from atooms.core.utils import mkdir
for f in args.file_inp:
fmt = args.inp
t = trj.Trajectory(f, fmt=fmt)
fmt_out = t.suffix
# Define slice
# We interpret --first N --last N as a request of step N
if args.last == args.first and args.last is not None:
args.last += 1
sl = fractional_slice(args.first, args.last, args.skip, len(t))
ts = trajectory.Sliced(t, sl)
for i, system in enumerate(ts):
f_out = args.file_out.format(step=t.steps[i],
frame=i,
base=os.path.splitext(f)[0],
ext=os.path.splitext(f)[1])
mkdir(os.path.dirname(f_out))
with trj.Trajectory(f_out, fmt=fmt_out, mode='w') as th_out:
if args.fields is not None:
th_out.fields = args.fields.split(',')
else:
th_out.fields.append('radius')
th_out.fields.append('radius')
th_out.write(system, step=t.steps[i])
def main(args):
"""Convert trajectory `file_inp` to `file_out`."""
args.file_inp = args.file_inp[0]
if args.fields is not None:
args.fields = args.fields.split(',')
if args.out is not None and not args.out in trajectory.Trajectory.formats:
# available_formats()
__import__('pdb').set_trace()
raise ValueError('Unknown output format %s' % args.out)
if args.file_out == '-':
args.file_out = '/dev/stdout'
if args.verbose:
setup_logging('atooms', 20)
import atooms.core.progress
if args.file_out != '/dev/stdout':
atooms.core.progress.active = True
if args.folder:
t = trajectory.folder.Foldered(args.file_inp, cls=args.inp)
else:
t = trajectory.Trajectory(args.file_inp, fmt=args.inp)
# If no output format is provided we use the input one
if args.out is None:
out_class = t.__class__
else:
out_class = args.out
if args.precision is not None:
t.precision = args.precision
if args.flatten_steps:
t.steps = range(1, len(t) + 1)
# Reset random number generator
if args.seed:
random.seed(args.seed)
# Trick to allow some trajectory formats to set the box side.
# This way the cell is defined as we read the sample (callbacks
# will not do that).
if args.side is not None:
def fix_cell(system, side):
from atooms.system import Cell
system.cell = Cell(side)
return system
L = args.side.split(',')
if len(L) == 1:
t.add_callback(fix_cell, [L, L, L])
else:
t.add_callback(fix_cell, [float(_) for _ in L])
# Define slice.
# We interpret --first N --last N as a request of step N
if args.last == args.first and args.last is not None:
args.last += 1
sl = fractional_slice(args.first, args.last, args.skip, len(t))
# Unfold if requested
if args.unfold:
tu = trajectory.Unfolded(t) #, fix_cm=True)
else:
tu = t
# Fix CM and fold back
if args.fix_cm:
tu.add_callback(trajectory.fix_cm)
tu.add_callback(trajectory.fold)
# Here we could you a trajectory slice t[sl] but this will load
# everything in ram (getitem doesnt provide a generator). This
# will be fixed with python 3.
ts = trajectory.Sliced(tu, sl)
# Change number of particles
if args.N > 0:
def decimate_system(system, N):
from atooms.system.particle import decimate
system.particle = decimate(system.particle, N)
return system
ts.register_callback(decimate_system, args.N)
# Change density and temperature
if args.rho is not None:
ts.register_callback(trajectory.decorators.set_density, args.rho)
if args.temperature is not None:
ts.register_callback(trajectory.decorators.set_temperature,
args.temperature)
# Change species layout if requested
if args.species_layout is not None:
ts.register_callback(trajectory.decorators.change_species,
args.species_layout)
# Sort by species id
if args.species_sort:
ts.register_callback(trajectory.decorators.sort)
# We enforce regular periodicity; steps is None is trajectory is not periodic
try:
steps = check_block_size(ts.steps, ts.block_size, prune=True)
except IndexError:
print('# Warning: something wrong with periodicity check.')
print(
'# We will proceed, but you should check the converted trajectory.'
)
steps = ts.steps
#
# ---------------------
# Trajectory conversion
# ---------------------
#
include_list, exclude_list = [], []
if len(args.fields_include) > 0:
include_list = args.fields_include.split(',')
if len(args.fields_exclude) > 0:
exclude_list = args.fields_exclude.split(',')
fout = trajectory.convert(ts,
out_class,
args.file_out,
fields=args.fields,
include=include_list,
exclude=exclude_list,
steps=steps)
if args.ff:
from atooms.trajectory.hdf5 import add_interaction_hdf5
add_interaction_hdf5(fout, args.ff)
if args.verbose and args.file_out != '/dev/stdout':
print('# converted %s to %s' % (args.file_inp, fout))
t.close()