-
Notifications
You must be signed in to change notification settings - Fork 5
/
plot_forces.py
executable file
·129 lines (107 loc) · 4.76 KB
/
plot_forces.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
#!/usr/bin/env python
# This script will plot a series of force-matched potentials if they are written out as lammps directores.
# Here's an example of how to generate the directories:
# for i in range(10):
# fm.force_match()
# fm.write_lammps_directory(folder='batch_%d' % i)
#
# A plot may be generated by running plot_forces.py "baatch_" from the directory containing the batch directories.
# This code requires matplotlib.
#
#
#
#
#
import os
import numpy as np
import matplotlib.pyplot as plt
pot_types = {'bond': 'cg_force_bond.table',
'pair': 'cg_force_pair.table'}
def remaining_file_len(f):
pos = f.tell()
i = 0
for l in f:
i += 1
f.seek(pos)
return i
def count_remaining_tables(file, table_points=10000, header_length=5, footer_length=2):
count = remaining_file_len(file) / (header_length + footer_length + table_points)
assert count * (header_length + footer_length + table_points) == remaining_file_len(file)
return count
def walk_files(pot_type, dir_prefix="batch_"):
if not pot_type in pot_types:
print '{} is not a valid force type.'.format(pot_type)
i = 0
while True:
pass_dir = '{}{}'.format(dir_prefix, i)
if not os.path.exists(pass_dir):
break
yield open(os.path.join(pass_dir, pot_types[pot_type]))
i += 1
def consume_lines(f,l):
for i in range(l):
f.readline()
def gen_tables(files, index = -1, table_points=10000, header_length=5, footer_length=2, columns=3, return_extra=False):
if(type(files) == file):
files = [files]
for f in files:
f.seek(0)
if index >= count_remaining_tables(f, table_points, header_length, footer_length):
raise IOError('Not enough tables in {}'.format(index))
for tindex in range(count_remaining_tables(f, table_points, header_length, footer_length)):
if(tindex == index or index == -1):
matrix = np.empty( (table_points, columns) )
#THIS STUFF HERE IS NOT KINDA LAZY. SHOULD ACTUALLY INFER THIS STUFF
name = f.readline().split('#')[1]
f.readline()
lammps_header = ''.join([f.readline() for x in range(2)])
f.readline()
for i in range(table_points):
sline = f.readline().split()
matrix[i,] = [float(x) for x in sline[1:]]
consume_lines(f, footer_length)
if(return_extra):
yield matrix, name, lammps_header
else:
yield matrix
else:
consume_lines(f, table_points + header_length + footer_length)
def write_table(f, matrix, name, lammps_header, conversion_factors):
f.write('#{}\n'.format(name))
f.write(lammps_header)
f.write('\n')
for i,v in enumerate(matrix):
f.write('{}'.format(i))
for conversion, vi in zip(v, conversion_factors):
f.write(' {}'.format(conversion * vi))
f.write('\n')
f.write('\n\n')
def plot_tables(ms, N=15, plot_column=1, cmap='RdBu'):
color = plt.get_cmap(cmap, N)
index = 0
for m in ms:
plt.plot(m[:,0], m[:,plot_column], color=color(index))
plt.axis([0, max(m[:,0]), -20, 20])
#plt.axis([0, max(m[:,0]), -np.amax(m[:,plot_column]), np.amax(m[:,plot_column])])
index += 1
yield plt
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Plot multiple ForcePy force-mathcing lammps directories over time")
parser.add_argument('directory_prefix', help='directories containing lammps tables should be prefix0 prefix1 prefix2... etc. Note that the first directory must be prefix0')
parser.add_argument('-table_points', type=int, default=10000)
pargs = parser.parse_args()
prefix = pargs.directory_prefix
for k in pot_types.iterkeys():
ntables = max((count_remaining_tables(x, table_points=pargs.table_points) for x in walk_files(k, prefix)))
passes = sum(1 for x in walk_files(k,prefix))
print 'There are {} potentials for the {} type which were refined {} times'.format(ntables, k, passes)
for i in range(ntables):
plt.figure(figsize=(16,9))
for p in plot_tables(gen_tables(walk_files(k, prefix), i, table_points=pargs.table_points), N=passes):
pass
p.savefig('{}_potential_{}.pdf'.format(k, i), bbox_inches=0)
plt.figure(figsize=(16,9))
for p in plot_tables(gen_tables(walk_files(k, prefix), i, table_points=pargs.table_points), N=passes, plot_column=2):
pass
p.savefig('{}_force_{}.pdf'.format(k, i), bbox_inches=0)