Feat/wraptools

tools/py/contract-trajectory.py

@@ -55,11 +55,11 @@ def contract_trajectory(fns_in, fn_out_template, n_new, cell_units_in, cell_unit
 
     # Open input trajectory iterators.
     trjs_in = [iter_file_name_raw(fn) for fn in fns_in]
-    mode = os.path.splitext(fn)[1]
+    mode = os.path.splitext(fn)[-1]
 
     # Open output files.
     fs_out = [open_backup(fn, "w") for fn in fns_out]
-    mode_out = os.path.splitext(fn_out_template)[1]
+    mode_out = os.path.splitext(fn_out_template)[-1]
 
     # prepare ring polymer rescaler
     rescale = nm_rescale(n, n_new)

tools/py/mux-positions.py

@@ -4,21 +4,25 @@
 
 import os
 import sys
+import numpy as np
 import argparse
 from ipi.utils.messages import verbosity
+from ipi.utils.io.io_units import auto_units, process_units
+import copy
 
-from ipi.utils.io import open_backup, iter_file_name, print_file
+from ipi.utils.io import open_backup, iter_file_name_raw, print_file
 
 
 description = """
 Read positions of individual beads from a set of trajectory files and
-multiplex them into a single output trajectory. Trajectory file formats are
-inferred from file extensions, the number of beads is given by the number of
-input files.
+a) multiplex them into a single output trajectory
+b) wrap/unwrap them.
+Trajectory file formats are inferred from file extensions, the number
+of beads is given by the number of input files.
 """
 
 
-def main(fns_in, fn_out, begin, end, stride):
+def main(fns_in, fn_out, begin, end, stride, wrap, unwrap):
 
     verbosity.level = "low"
     print('Multiplexing {:d} beads into one trajectory.'.format(len(fns_in)))
@@ -33,15 +37,20 @@ def main(fns_in, fn_out, begin, end, stride):
     print()
 
     # Open input trajectory iterators.
-    trjs_in = [iter_file_name(fn) for fn in fns_in]
+    trjs_in = [iter_file_name_raw(fn) for fn in fns_in]
+    mode = os.path.splitext(fns_in[0])[-1]
 
     # Open output file.
     f_out = open_backup(fn_out, 'w')
-    mode_out = os.path.splitext(fn_out)[1]
+    mode_out = os.path.splitext(fn_out)[-1]
 
     # Loop over all frames.
     i_frame = 0
     i_frame_saved = 0
+
+    # There can be multiple trajectories, so we store a frame_last for each trajectory
+    frame_last = [None] * len(fns_in)
+
     while True:
 
         # Check the endpoint index, exit if we're done.
@@ -53,11 +62,24 @@ def main(fns_in, fn_out, begin, end, stride):
 
         try:
             # Get the frames from all trajectories...
-            for trj in trjs_in:
+            for idx,trj in enumerate(trjs_in):
                 frame = trj.next()
+
+                # gets units from first frame
+                dimension, units, cell_units = auto_units(comment=frame["comment"])
+                frame = process_units(dimension=dimension, units=units, cell_units=cell_units, mode=mode, **frame)
+
+                if wrap:
+                    frame = wrap_positions(frame)
+                if unwrap:
+                    frame = unwrap_positions(frame,frame_last[idx])
+
+                frame_last[idx] = frame.copy()
+
+
                 # ... and possibly save them in the output trajectory.
                 if do_output:
-                    print_file(mode_out, frame['atoms'], frame['cell'], f_out)
+                    print_file(mode_out, frame['atoms'], frame['cell'], f_out, dimension=dimension, units=units, cell_units=cell_units)
             if do_output:
                 i_frame_saved += 1
         except StopIteration:
@@ -78,6 +100,36 @@ def main(fns_in, fn_out, begin, end, stride):
     print('Saved {:d} frames.'.format(i_frame_saved))
 
 
+def wrap_positions(frame):
+    pos = frame['atoms'].q.copy()
+    pos.shape = (len(pos)/3,3)
+    cell = frame['cell'].h
+    cell_inv = np.linalg.inv(cell)
+    s = np.dot(cell_inv,pos.T)
+    s -= np.round(s)
+    frame['atoms'].q = np.dot(cell,s).T.flatten()
+    return frame
+
+def unwrap_positions(frame,framelast):
+    if framelast is None:
+        return frame
+
+    poslast_uwr       = framelast['atoms'].q.copy()
+    poslast_uwr.shape = (len(poslast_uwr)/3,3)
+
+    pos               = frame['atoms'].q.copy()
+    pos.shape         = (len(pos)/3,3)
+
+    d = pos - poslast_uwr
+
+    cell            = frame['cell'].h
+    cell_inv        = np.linalg.inv(cell)
+    s = np.dot(cell_inv,d.T)
+    s -= np.round(s)
+    d = np.dot(cell,s).T
+    frame['atoms'].q = (poslast_uwr+d).flatten()
+    return frame
+
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser(description=description)
@@ -92,8 +144,12 @@ def main(fns_in, fn_out, begin, end, stride):
                         help='Step to end.')
     parser.add_argument('--stride', type=int, default=1,
                         help='Stride in steps.')
+    parser.add_argument('--wrap',  action='store_true', default=False,
+                        help='Wrap atomic positions.')
+    parser.add_argument('--unwrap',  action='store_true', default=False,
+                        help='Unwrap atomic positions.')
 
     args = parser.parse_args()
 
     # Process everything.
-    main(args.filenames, args.filename_out, args.begin, args.end, args.stride)
+    main(args.filenames, args.filename_out, args.begin, args.end, args.stride, args.wrap, args.unwrap)