1import math
2import sys
3
4# arguments read from command line
5# name of input file
6FILENAME_ = sys.argv[1]
7# number of BIAS
8NBIAS_ = int(sys.argv[2])
9# temperature
10KBT_ = float(sys.argv[3])
11# default parameters for WHAM
12# number of WHAM iterations
13NWHAM_ = 1000
14# convergence thresold
15THRES_ = 1.0e-10
16
17
18def get_wham_weights(nbias, nframes, bias_ts, nwham=NWHAM_, thres=THRES_):
19    # find minimum bias
20    min_bias = min(bias_ts)
21    # initialize weights
22    w = []
23    for i in range(0, nframes): w.append(1.0)
24    # offset and exponential of the bias
25    expv = []
26    for i in range(0, len(bias_ts)): expv.append(math.exp((-bias_ts[i]+min_bias)/KBT_))
27    # initialize Z
28    Z = []
29    for j in range(0, nbias): Z.append(1.0)
30    # WHAM iterations
31    for iii in range(0, nwham):
32        # store Z
33        Z_old = Z[:]
34        # recompute weights
35        norm = 0.0
36        for i in range(0, len(w)):
37            ew = 0.0
38            for j in range(0, len(Z)): ew += expv[i*len(Z)+j] / Z[j]
39            w[i] = 1.0 / ew
40            norm += w[i]
41        # normalize weights
42        for i in range(0, len(w)): w[i] /= norm
43        # recompute Z
44        for j in range(0, len(Z)): Z[j] = 0.0
45        for i in range(0, len(w)):
46            for j in range(0, len(Z)): Z[j] += w[i]*expv[i*len(Z)+j]
47        # normalize Z
48        norm = sum(Z)
49        for j in range(0, len(Z)): Z[j] /= norm
50        # compute change in Z
51        eps = 0.0
52        for j in range(0, len(Z)):
53            d = math.log(Z[j]/Z_old[j])
54            eps += d*d
55        # check convergence
56        if(eps<thres): break
57    # return weights
58    return w
59
60# read FILENAME_
61bias_ts=[]
62for lines in open(FILENAME_, "r").readlines():
63    riga=lines.strip().split()
64    # skip comment lines
65    if(riga[0]=="#!"): continue
66    # read bias values
67    # umbrella-sampling typical format
68    if(len(riga) == NBIAS_):
69       i0 = 0
70       i1 = NBIAS_
71    # bias exchange typical format
72    elif(len(riga) == 2*NBIAS_+1):
73       i0 = NBIAS_+1
74       i1 = 2*NBIAS_+1
75    # unknown format
76    else:
77       print( FILENAME_,"format is unknown!")
78       exit()
79    for i in range(i0, i1):
80        bias_ts.append(float(riga[i]))
81
82# number of frames
83nframes = int(len(bias_ts) / NBIAS_)
84
85# printout
86print( "Number of frames::", nframes )
87print( "Number of bias entries::", len(bias_ts))
88
89# get wham weights
90ws = get_wham_weights(NBIAS_, nframes, bias_ts)
91
92# printout WHAM weights to file
93print( "Weights have been written to weights.dat")
94log = open("weights.dat", "w")
95for i in range(0, len(ws)): log.write("%16.12e\n" % ws[i])
96log.close()
97