Major update for deinterleave_IDWS 96/4796/2
authorJérôme Hénin <heninj@gmail.com>
Fri, 9 Nov 2018 00:04:07 +0000 (01:04 +0100)
committerJim Phillips <jim@ks.uiuc.edu>
Fri, 9 Nov 2018 14:49:04 +0000 (08:49 -0600)
finally good behavior with ParseFEP and the default option
(subsample)

Change-Id: Id1534bb802db1653178eb61f935df80bf54cef06

lib/alch/deinterleave_idws.py

index 82109e5..9e3a8d5 100755 (executable)
@@ -9,11 +9,6 @@
 # - <base>_bwd_r.fepout (backward data, with order of windows reversed to simulate a
 #   reverse sequential transformation (0 -> 1 -> 0)
 
-# Notes on further analysis:
-# Missing samples of deltaE must be interpolated to recover the same number of samples in each window
-# This is necessary for ParseFEP up to version 2.1
-# Hopefully later versions will fix this
-
 # Note:
 # The first window samples forward only, the last window backward only
 #
@@ -33,119 +28,131 @@ try:
 except ImportError:
     interp_avail = False
 
-def interpolated_window(steps_window, deltaE_window, lambda1, lambda2, lines=[]):
-    equil = True
-    acc_exp = 0
-    acc_N = 0
-    buf = ''
-    # Create interpolation function for deltaE data
-    func_dE = InterpolatedUnivariateSpline(steps_window, deltaE_window)
-    for t in range(len(steps)):
-        s = steps[t]
-        if equil and s >= equilSteps:
-            buf += ('#{} STEPS OF EQUILIBRATION AT LAMBDA {} COMPLETED\n#STARTING COLLECTION OF ENSEMBLE AVERAGE\n').format(equilSteps, lambda1)
-            equil = False
-            acc_exp = 0
-            acc_N = 0
-        dE = float(func_dE(s))
-        acc_exp += math.exp(-beta * dE)
-        if acc_exp == 0: # catch floating-point underflow due to very high dE (clash)
-            acc_exp = 1e-16
-        acc_N += 1
-        buf += 'FepEnergy: {:6d}  {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f}\n'.format(s, 0, 0, 0, 0, dE, 0., T[t], -RT*math.log(acc_exp / acc_N))
-    return buf, -RT*math.log(acc_exp / acc_N)
-
-# Note: this doesn't subsample anymore, as I've given up the subsampling trick, which is frustratingly complicated to get right
-# and wouldn't be necessary if ParseFEP was slightly less ****
-# which hopefully it will be soon
-def subsampled_window(steps_window, deltaE_window, lambda1, lambda2, lines):
-    equil = True
-    acc_exp = 0
-    acc_N = 0
-    buf = ''
-    first = steps_window[0]
-
-    stride = 1
-    # Disabling the trick blow, as this does not reliably produce even-sampled windows
-    # and then parseFEP has a fit
-    # find first gap in data
-    #for i in range(len(steps_window)-1):
-    #    diff = steps_window[i+1] - steps_window[i]
-    #    if diff != alchOutFreq:
-    #        if diff % alchOutFreq != 0:
-    #            sys.exit('Error: gap between data points is ' + str(diff) + ' instead of a multiple of ' + str(alchOutFreq))
-    #        else:
-    #            stride = diff // alchOutFreq - 1
-    #        break
-
-    for t in range(0, len(steps_window), stride):
-        s = steps_window[t]
-        if equil and s >= equilSteps:
-            buf += ('#{} STEPS OF EQUILIBRATION AT LAMBDA {} COMPLETED\n#STARTING COLLECTION OF ENSEMBLE AVERAGE\n').format(equilSteps, lambda1)
-            equil = False
-            acc_exp = 0
-            acc_N = 0
-        dE = deltaE_window[t]
-        acc_exp += math.exp(-beta * dE)
-        if acc_exp == 0: # catch floating-point underflow due to very high dE (clash)
-            acc_exp = 1e-16
-        acc_N += 1
-        buf += 'FepEnergy:' + lines[t]
-        #buf += 'FepEnergy: {:6d}  {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f}\n'.format(s, 0, 0, 0, 0, dE, 0., T[t], -RT*math.log(acc_exp / acc_N))
-    return buf, -RT*math.log(acc_exp / acc_N)
-
-
-def process_window():
-    # Create interpolated function from data and write new data for the whole window
-    if not last_window and len(steps_f) > 0:
-        print("Forward data: " + str(lambda1) + " to " + str(lambda2) + " (" + str(len(steps_f)) + " points)")
-        outf.write("#NEW FEP WINDOW: LAMBDA SET TO {} LAMBDA2 {}\n".format(lambda1, lambda2))
-        buf, dG = resample(steps_f, deltaE_f, lambda1, lambda2, lines_f)
-        outf.write(buf)
-        if line.startswith('#Free'):
-            outf.write(line)
-        else:
-            outf.write('#Free energy change for lambda window [ {} {} ] is {} ; net change until now is 0.\n'.format(lambda1, lambda2, dG))
-
-    if len(steps_b) > 0: # only if we do have backward data in this window
-        # create new window in the backward output
-        # for every window but the first one, which is forward-only
-        print("Backward data: " + str(lambda1) + " to " + str(lambdaIDWS) + " (" + str(len(steps_b)) + " points)")
-        outb_buffer.append("#NEW FEP WINDOW: LAMBDA SET TO {} LAMBDA2 {}\n".format(lambda1, lambdaIDWS))
-        buf, dG = resample(steps_b, deltaE_b, lambda1, lambdaIDWS, lines_b)
-        outb_buffer[-1] += buf
-        outb_buffer[-1] += '#Free energy change for lambda window [ {} {} ] is {} ; net change until now is 0.\n'.format(lambda1, lambdaIDWS, dG)
-
 
 parser = argparse.ArgumentParser(description='Process fepout data from Interleaved Double-Wide Sampling')
 parser.add_argument('filenames', nargs='+', help='fepout file name(s) with lambda from 0 to 1')
 parser.add_argument('-T', '--temperature', type=float, default=300., help='temperature for FEP estimate, in K')
 parser.add_argument('-r', '--reverse', type=bool, default=True, help='reverse order of backward-sampling windows')
-parser.add_argument('-i', '--interpolate', default=True, action='store_true', help='interpolate data rather than subsampling it')
+parser.add_argument('-i', '--interpolate', action='store_true', help='interpolate data rather than subsampling it')
 args = parser.parse_args()
 
 # Store the output in the same directory as the first file
 basename = os.path.splitext(args.filenames[0])[0]
-outf = open(basename + '_fwd.fepout', 'w')
 
 interpolating = False
 if args.interpolate:
     if interp_avail:
-        resample = interpolated_window
         interpolating = True
     else:
         print('Interpolation requires the scipy.interpolate package. Defaulting to subsampling.')
-        resample = subsampled_window
-else:
-    resample = subsampled_window
 
-# Buffer to hold the contents of the backwards fepout file, by window. Will be written in reverse.
-outb_buffer = []
+
+
+class Window:
+    def __init__(self, l1, l2):
+        self.deltaE = []
+        self.steps = []
+        self.T = []
+        self.equilSteps = -1
+        self.lambda1 = l1
+        self.lambda2 = l2
+        self.alchOutFreq = -1
+
+    def interpolate(self):
+        equil = True
+        acc_exp = 0
+        acc_N = 0
+        buf = ''
+        # Create interpolation function for deltaE data
+        func_dE = InterpolatedUnivariateSpline(self.steps, self.deltaE)
+        for t in range(len(global_steps)):
+            s = global_steps[t]
+            # # Skip equilibration steps altogether
+            # if s <= self.equilSteps:
+            #     continue
+            if equil and s > self.equilSteps:
+                buf += ('#{} STEPS OF EQUILIBRATION AT LAMBDA {} COMPLETED\n#STARTING COLLECTION OF ENSEMBLE AVERAGE\n').format(self.equilSteps, self.lambda1)
+                equil = False
+                acc_exp = 0
+                acc_N = 0
+            dE = float(func_dE(s))
+            acc_exp += math.exp(-beta * dE)
+            if acc_exp == 0: # catch floating-point underflow due to very high dE (clash)
+                acc_exp = 1e-16
+            acc_N += 1
+            buf += 'FepEnergy: {:6d}  {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f}\n'.format(s, 0, 0, 0, 0, dE, 0., args.temperature, -RT*math.log(acc_exp / acc_N))
+        return buf, -RT*math.log(acc_exp / acc_N)
+
+    def subsample(self):
+        equil = True
+        acc_exp = 0
+        acc_N = 0
+        buf = ''
+
+        stride = 1
+        # # find first gap in data
+        # for i in range(len(self.steps)-1):
+        #     diff = self.steps[i+1] - self.steps[i]
+        #     if diff != self.alchOutFreq:
+        #         if diff % self.alchOutFreq != 0:
+        #             sys.exit('Error: gap between data points is ' + str(diff) + ' instead of a multiple of ' + str(alchOutFreq))
+        #         else:
+        #             stride = (diff - 1) // self.alchOutFreq
+        #         break
+
+        # if stride == 0:
+        #     stride = 1
+
+        for t in range(0, len(self.steps), stride):
+            s = self.steps[t]
+            # # Skip equilibration steps altogether
+            # if s <= self.equilSteps:
+            #     continue
+            if equil and s > self.equilSteps:
+                buf += ('#{} STEPS OF EQUILIBRATION AT LAMBDA {} COMPLETED\n#STARTING COLLECTION OF ENSEMBLE AVERAGE\n').format(self.equilSteps, self.lambda1)
+                equil = False
+                acc_exp = 0
+                acc_N = 0
+            dE = self.deltaE[t]
+            acc_exp += math.exp(-beta * dE)
+            if acc_exp == 0: # catch floating-point underflow due to very high dE (clash)
+                acc_exp = 1e-16
+            acc_N += 1
+            buf += 'FepEnergy: {:6d}  {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f} {:14.4f}\n'.format(t, 0, 0, 0, 0, dE, 0., self.T[t], -RT*math.log(acc_exp / acc_N))
+        return buf, -RT*math.log(acc_exp / acc_N)
+
+
+    if interpolating:
+        resample = interpolate
+    else:
+        resample = subsample
+
+
+    def process(self, outfile, is_backward, dG_in):
+        if len(self.steps) == 0:
+            return
+        if is_backward:
+            print("Backward data: " + str(self.lambda1) + " to " + str(self.lambda2) + " (" + str(len(self.steps)) + " points)")
+        else:
+            print("Forward data: " + str(self.lambda1) + " to " + str(self.lambda2) + " (" + str(len(self.steps)) + " points)")
+
+        outfile.write("#NEW FEP WINDOW: LAMBDA SET TO {} LAMBDA2 {}\n".format(self.lambda1, self.lambda2))
+        buf, dG = self.resample()
+        dG_out = dG_in + dG
+        outfile.write(buf)
+        outfile.write('#Free energy change for lambda window [ {} {} ] is {} ; net change until now is {}\n'.format(self.lambda1, self.lambda2, dG, dG_out))
+        return dG_out
+
+
+# Lists of forward and backward windows data
+windows_f = []
+windows_b = []
+IDWS_on = False
+# Steps for all data, for interpolation
+global_steps = []
 
 first_window = True
 last_window = False
-
 RT = 1.98720650096e-3 * args.temperature #in kcal/mol
 beta = 1. / RT
 
@@ -153,74 +160,12 @@ beta = 1. / RT
 step_counter = -1
 alchOutFreq = -1
 equilSteps = 0
-
-deltaE_f = []
-steps_f = []
-lines_f = []
-deltaE_b = []
-steps_b = []
-lines_b = []
-steps = []
-T = []
-
-equil_re = re.compile('#(\d*) STEPS OF EQUILIBRATION')
+equil_re = re.compile(r'#(\d*) STEPS OF EQUILIBRATION')
 
 for line in fileinput.input(args.filenames):
     fields = re.split(' +', line.strip())
-    if line.startswith('#NEW FEP WINDOW:'):
-        # Process previous window, if any
-        process_window()
-        # data will contain all dE data for the window
-        # make two copies, forward and backward
-        deltaE_f.clear()
-        steps_f.clear()
-        lines_f.clear()
-        deltaE_b.clear()
-        steps_b.clear()
-        lines_b.clear()
-        # Common data: step number, temperature
-        steps.clear()
-        T.clear()
-        step_counter = -1
-        first_step = -1
-        alchOutFreq = -1
-        equilSteps = 0
-
-        # Extract lambda values
-        lambda1 = float(fields[6])
-        lambda2 = float(fields[8])
-        if len(fields) == 11:
-            lambdaIDWS = float(fields[10])
-            IDWS_on = True
-        else:
-            lambdaIDWS = -1.
-            IDWS_on = False
-
-        # Check that we start from an end-point
-        if first_window:
-            if lambda1 == 1:
-                start_from_one = True
-            elif lambda1 == 0:
-                start_from_one = False
-            else:
-                sys.exit('Lambda should start at zero or one, found ' + str(lambda1))
-
-        # Skip the last window for the "forward" output, as it is backward from 1
-        if ((not start_from_one and lambda1 == 1) or (start_from_one and lambda1 == 0)):
-            last_window = True
-
-        if (last_window):
-            # special case, last window is backward
-            lambdaIDWS = lambda2
-            lambda2 = -1
-
-        # Done processing the header of the first window
-        first_window = False
-
-    match = equil_re.match(line)
-    if match != None:
-        equilSteps = int(match.group(1))
 
+    # Most lines in the file are energy values, so start with them
     # Collect all timestep numbers (both fwd and bwd) for interpolation
     if line.startswith('FepE'):
         l = len(fields)
@@ -236,41 +181,88 @@ for line in fileinput.input(args.filenames):
             first_step = int(fields[1])
             step_counter = first_step
         elif step_counter == first_step:
-            alchOutFreq = int(fields[1]) - first_step
-            step_counter += alchOutFreq
+            wf.alchOutFreq = int(fields[1]) - first_step
+            if IDWS_on:
+                wb.alchOutFreq = wf.alchOutFreq
+            step_counter += wf.alchOutFreq
 
-        steps.append(step_counter)
-        T.append(float(fields[i_temp]))
+        if first_window:
+            global_steps.append(step_counter)
 
-        if line.startswith('FepEnergy:') and not last_window:
-            steps_f.append(step_counter)
-            deltaE_f.append(float(fields[6]))
-            if not interpolating:
-                lines_f.append(line[10:])
+        # Discard every other sample if not doing IDWS or interpolating
+        if line.startswith('FepEnergy:') and (interpolating or IDWS_on or (step_counter//alchOutFreq)%2 == 0):
+            wf.steps.append(step_counter)
+            wf.deltaE.append(float(fields[6]))
+            wf.T.append(float(fields[8]))
 
-        if line.startswith('FepE_back:') or (last_window and line[:10] == 'FepEnergy:'):
-            steps_b.append(step_counter)
-            deltaE_b.append(float(fields[6]))
-            if not interpolating:
-                lines_b.append(line[10:])
+        if line.startswith('FepE_back:'):
+            wb.steps.append(step_counter)
+            wb.deltaE.append(float(fields[6]))
+            wb.T.append(float(fields[8]))
+
+        if wf.alchOutFreq > 0:
+            step_counter += wf.alchOutFreq
+
+    elif line.startswith('#NEW FEP WINDOW:'):
+
+        if step_counter >= 0:
+            first_window = False
+        step_counter = -1
 
-        if alchOutFreq > 0:
-            step_counter += alchOutFreq
+        # Extract lambda values
+        cur_lambda1 = float(fields[6])
+        cur_lambda2 = float(fields[8])
 
-# Process last window
-process_window()
+        # create new "forward" window
+        wf = Window(cur_lambda1, cur_lambda2)
 
+        if len(fields) == 11:
+            cur_lambdaIDWS = float(fields[10])
+            IDWS_on = True
+            wb = Window(cur_lambda1, cur_lambdaIDWS)
+        else:
+            cur_lambdaIDWS = -1.
+            IDWS_on = False
 
+        # Check that we start from an end-point
+        if first_window:
+            if cur_lambda1 == 1:
+                start_from_one = True
+            elif cur_lambda1 == 0:
+                start_from_one = False
+            else:
+                sys.exit('Lambda should start at zero or one, found ' + str(cur_lambda1))
+
+        # Skip the last window for the "forward" output, as it is backward from 1
+        if ((not start_from_one and cur_lambda1 == 1) or (start_from_one and cur_lambda1 == 0)):
+            last_window = True
+            # append the "forward" data to the list of backward windows
+            windows_b.append(wf)
+        else:
+            windows_f.append(wf)
+
+        if IDWS_on:
+            windows_b.append(wb)
+
+    match = equil_re.match(line)
+    if match != None:
+        equilSteps = int(match.group(1))
+        wf.equilSteps = equilSteps
+        if IDWS_on:
+            wb.equilSteps = equilSteps
+
+
+outf = open(basename + '_fwd.fepout', 'w')
+dG_acc = 0.
+for w in windows_f:
+    dG_acc = w.process(outf, False, dG_acc)
 outf.close()
 
-outb = open(basename + '_bwd.fepout', 'w')
 if args.reverse:
-    # Write backward windows in reverse order
-    for string in reversed(outb_buffer):
-        outb.write(string)
-else:
-    # Write output in same order as run
-    for string in outb_buffer:
-        outb.write(string)
-outb.close()
+    windows_b.reverse()
 
+outb = open(basename + '_bwd.fepout', 'w')
+dG_acc = 0.
+for w in windows_b:
+    dG_acc = w.process(outb, True, dG_acc)
+outb.close()