From dd07c2ada71de2dee19cb5c6b0f343442309ca9f Mon Sep 17 00:00:00 2001
From: James Graham <J.A.Graham@soton.ac.uk>
Date: Wed, 20 Jul 2016 17:09:50 +0100
Subject: [PATCH] Added tqdm progress bar to Boltzmann Inversion
---
pycgtool/bondset.py | 26 ++++++++++++++++++++------
pycgtool/pycgtool.py | 2 +-
2 files changed, 21 insertions(+), 7 deletions(-)
diff --git a/pycgtool/bondset.py b/pycgtool/bondset.py
index 8449467..b61b11d 100644
--- a/pycgtool/bondset.py
+++ b/pycgtool/bondset.py
@@ -8,7 +8,11 @@ import itertools
import numpy as np
import math
-import random
+
+try:
+ from tqdm import tqdm
+except ImportError:
+ pass
from .util import stat_moments, sliding, dist_with_pbc, transpose_and_sample
from .util import extend_graph_chain, cross, backup_file
@@ -351,14 +355,24 @@ class BondSet:
# TypeError is raised when residues on end of chain calc bond to next
pass
- def boltzmann_invert(self):
+ def boltzmann_invert(self, progress=False):
"""
Perform Boltzmann Inversion of all bonds to calculate equilibrium value and force constant.
+
+ :param progress: Display a progress bar using tqdm if available
"""
- for mol in self._molecules:
- for bond in self._molecules[mol]:
- bond.boltzmann_invert(temp=self._temperature,
- angle_default_fc=self._angle_default_fc)
+ bond_iter = itertools.chain(*self._molecules.values())
+ bond_iter_wrap = bond_iter
+ if progress:
+ try:
+ total = sum(map(len, self._molecules.values()))
+ bond_iter_wrap = tqdm(bond_iter, total=total)
+ except NameError:
+ pass
+
+ for bond in bond_iter_wrap:
+ bond.boltzmann_invert(temp=self._temperature,
+ angle_default_fc=self._angle_default_fc)
def dump_values(self, target_number=100000):
"""
diff --git a/pycgtool/pycgtool.py b/pycgtool/pycgtool.py
index 2660c1d..deae69e 100755
--- a/pycgtool/pycgtool.py
+++ b/pycgtool/pycgtool.py
@@ -56,7 +56,7 @@ def main(args, config):
if args.bnd:
if args.map:
logger.info("Beginning Boltzmann inversion")
- bonds.boltzmann_invert()
+ bonds.boltzmann_invert(progress=True)
if config.output_forcefield:
logger.info("Creating GROMACS forcefield directory")
ff = ForceField("ff" + config.output_name + ".ff")
--
GitLab