You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
88 lines
2.6 KiB
88 lines
2.6 KiB
import sys |
|
import argparse |
|
import matplotlib |
|
matplotlib.use('PS') |
|
#matplotlib.rc('text', usetex=True) |
|
import numpy as np |
|
import statsmodels.api as sm |
|
#matplotlib.use('Qt5agg') |
|
import matplotlib.pyplot as plt |
|
|
|
parser = argparse.ArgumentParser(description='Compute Diffusion Coeffients from MSD data (nxy format)') |
|
parser.add_argument('data', metavar='data', type=str) |
|
parser.add_argument('-n', '--noheader', action="store_true") |
|
parser.add_argument('-s', '--start', type=int, default=0, help='StartStep') |
|
parser.add_argument('-e', '--end', type=int, default=-1, help='EndStep') |
|
opt = parser.parse_args(sys.argv[1:]) |
|
|
|
|
|
data_file = opt.data |
|
|
|
x = [] |
|
ys = [] |
|
ncol = -1 |
|
lineno = 0 |
|
headers = [] |
|
with open(data_file, 'r') as f: |
|
for line in f: |
|
lineno +=1 |
|
if (line.startswith('#')): |
|
if opt.noheader: |
|
continue |
|
if (lineno == 1): |
|
headers = line.split()[1:] |
|
continue |
|
arr = line.split() |
|
x.append(float(arr[0])) |
|
ys.append(list(map(float, arr[1:]))) |
|
|
|
if (ncol == -1): |
|
ncol = len(arr) -1 |
|
else: |
|
if (ncol != (len(arr) -1)): |
|
raise ValueError('The number of columns of data file is not consistent!') |
|
|
|
|
|
print('There are {} columns of data.'.format(ncol)) |
|
print() |
|
|
|
print(headers) |
|
if (len(headers) != ncol): |
|
headers = [] |
|
for col in range(ncol): |
|
headers.append('Column {}'.format(col+1)) |
|
|
|
for col in range(ncol): |
|
print ('Column {}:'.format(col+1)) |
|
|
|
all_data_x = np.array(x, dtype=np.double).reshape(-1, 1) |
|
data_x = all_data_x[opt.start:opt.end] |
|
|
|
all_data_y = [d[col] for d in ys] |
|
data_y = np.array(all_data_y[opt.start:opt.end], dtype=np.double).reshape(-1, 1) |
|
X = sm.add_constant(data_x) |
|
ols = sm.OLS(data_y, X) |
|
ols_result = ols.fit() |
|
|
|
summ = ols_result.summary() |
|
print(summ) |
|
print('*'*80) |
|
diff_coeff = ols_result.params[1]/6.0 |
|
print(' Diffusion coefficient (1/6): {:16.10F} A^2/ps, {:17.12E} m^2/s'.format(diff_coeff, diff_coeff*1.0e-8)) |
|
diff_coeff = ols_result.params[1]/2.0 |
|
print(' Diffusion coefficient (1/2): {:16.10F} A^2/ps, {:17.12E} m^2/s'.format(diff_coeff, diff_coeff*1.0e-8)) |
|
diff_coeff = ols_result.params[1]/4.0 |
|
print(' Diffusion coefficient (1/4): {:16.10F} A^2/ps, {:17.12E} m^2/s'.format(diff_coeff, diff_coeff*1.0e-8)) |
|
print('*'*80) |
|
print() |
|
|
|
predicted_y = ols_result.predict(X) |
|
|
|
plt.plot(x, all_data_y, label=headers[col]) |
|
plt.plot(data_x, predicted_y, 'k--') |
|
|
|
plt.legend(loc=2) |
|
plt.xlabel('Time (ps)') |
|
plt.ylabel('MSD (Angstrom$^2$)') |
|
plt.savefig('msd.pdf', dpi=1200, format='pdf') |
|
# plt.show()
|
|
|