Merge pull request #4189 from evoyiatzis/patch-2

Update fep.py

tools/fep/fep.py

@@ -2,23 +2,40 @@
 # fep.py - calculate free energy from compute fep results
 
 import sys
+from argparse import ArgumentParser
 import math
 
-if len(sys.argv) < 2:
-    print("Free Energy Perturbation")
-    print("usage: fep.py temperature < out.fep")
-    sys.exit()
+def compute_fep():
 
-rt = 0.008314 / 4.184 * float(sys.argv[1]) # in kcal/mol
+    parser = ArgumentParser(description='A python script to calculate free energy from compute fep results')
 
-v = 1.0
-sum = 0.0
-for line in sys.stdin:
-    if line.startswith("#"):
-        continue
-    tok = line.split()
-    if len(tok) == 4:
-        v = float(tok[3])
-    sum += math.log(float(tok[2]) / v)
+    parser.add_argument("units", help="unit system can be lj, real or si")
+    parser.add_argument("Temperature", type=float, help="The temperature of the system")
+    parser.add_argument("InputFile", help="The name of the file with the data from compute fep.")
+
+    args = parser.parse_args()
+
+    r_value = {'lj': 1.0, 'real': 0.0019872036, 'si': 8.31446}
+
+    if args.units in r_value:
+        rt = r_value[args.units] * args.Temperature
+    else:
+        sys.exit("The provided units keyword is not valid")
+
+    v = 1.0
+    mysum = 0.0
+
+    with open(args.InputFile, "r") as input_file:
+        for line in input_file:
+            if line.startswith("#"):
+                continue
+            tok = line.split()
+            if len(tok) == 4:
+                v = float(tok[3])
+            mysum += math.log(float(tok[2]) / v)
+
+    print(-rt * mysum)
+
+if __name__ == "__main__":
+    compute_fep()
 
-print(-rt * sum)