def main(): args = _args().parse_args() data = h5py.File(args.trj, "r") idx, box, trj, masses = files_io.prepare_h5md( data, args.group, args.begin, args.end, no_image=True, step=args.every_frame ) half_box = 0.5 * box trj = fix_pbc(trj, box, half_box, args.N, args.molecules) if args.out_ee: print ("Calculating end-end distance...") ee = [] for frame in trj: ee.extend(compute_ee(args.molecules, args.N, frame)) out_ee = args.out_ee np.savetxt(out_ee, ee) print ("Saved to {}".format(out_ee)) if args.out_ee_points: print ("Save end-to-end points (raw data)...") ee_points = [] for frame in trj: for ch in xrange(args.molecules): b1, b2 = frame[ch * args.N], frame[ch * args.N + args.N - 1] ee_points.append([b1, b2]) cPickle.dump(ee_points, open(args.out_ee_points, "wb")) if args.out_rg: print ("Calculating radius of gyration") rgs = [] tot_masses = np.sum(masses[: args.N]) for frame in trj: rg = calculate_rg(args.molecules, args.N, masses, box, half_box, tot_masses, frame) rgs.extend(rg) out_rg = args.out_rg np.savetxt(out_rg, rgs) print ("Saved Rg^2 to {}".format(out_rg)) if args.out_int: print ("Calculate internal distances.") int_distances = [] fidx = 0 for frame in trj: sys.stdout.write("{}\r".format(fidx)) sys.stdout.flush() int_distances.append(bonds.calculate_msd_internal_distance(frame, args.molecules, args.N, box, half_box)) fidx += 1 out_int = args.out_int np.savetxt(out_int, np.average(np.array(int_distances), axis=0)) print ("Saved internal distance to {}".format(out_int))
def _h5_processing(args, q_vector, start_stop_list, pos_cons): nt = MPI.COMM_WORLD.size at_group = args.group if nt > 1: pid = MPI.COMM_WORLD.rank if pid >= len(start_stop_list): return [] start_stop = start_stop_list[pid] h5file = h5py.File(args.trj, 'r', driver='mpio', comm=MPI.COMM_WORLD) else: h5file = h5py.File(args.trj, 'r') start_stop = start_stop_list[0] _, box, trj, _ = files_io.prepare_h5md(h5file, at_group, start_stop[0], start_stop[1]) half_box = 0.5*box vectors = [] process_tuples = [ [vectors, q_vector, bond_libs.calculate_bond_vec], ] frame_idx = start_stop[0] for frame in trj: print('Frame {}'.format(frame_idx)) for output, atom_ids, functor in process_tuples: if atom_ids is None or atom_ids.size == 0: continue atoms = frame[atom_ids] atom_tuples = [] for d_tuple in atoms: atom_tuples.append([np.array(x) for x in d_tuple]) output.append(functor(np.array(atom_tuples), box, half_box)) frame_idx += 1 return vectors