diff --git a/trjconv.py b/trjconv.py index 8ff1d8c..90a874e 100644 --- a/trjconv.py +++ b/trjconv.py @@ -9,6 +9,9 @@ num_particles = 0 step_size = 1 inp = "trajectory.xyz" out = "trjconv.xyz" +skip = -1 +move = 0 +max_z = 0 for i, val in enumerate(sys.argv): if val == "-dt": step_size = int(sys.argv[i+1]) @@ -16,12 +19,21 @@ for i, val in enumerate(sys.argv): inp = sys.argv[i+1] if val == "-o": out = sys.argv[i+1] + if val == "-s": + skip = int(sys.argv[i+1]) + if val == "-m": + move = float(sys.argv[i+1]) + max_z = float(sys.argv[i+2]) print "Extracting every %s step from %s" % (step_size, sys.argv[1]) step_counter = 0 in_step = False outfile = open(out, "w") + +if skip != -1: + print "Skipping %s frames" % skip + with open(inp, "r") as infile: for line in infile: @@ -34,8 +46,17 @@ with open(inp, "r") as infile: if not line.startswith("atom"): step_counter += 1 - if step_counter == 0 or (step_counter % step_size) == 0: - outfile.write(line) + if step_counter > skip and step_counter == 0 or (step_counter % step_size) == 0: + if move > 0 and line.startswith("atom"): + split = line.split() + split[3] = float(split[3]) + move + if(split[3] > max_z): + split[3] -= max_z + if(split[3] < 0): + split[3] += max_z + outfile.write("%s %s %s %s\n" % (split[0],split[1],split[2],split[3])) + else: + outfile.write(line) if step_counter % 100 == 0: print(step_counter)