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()