diff --git a/pycgtool/bondset.py b/pycgtool/bondset.py index 84494674c5c295b1b9e4f4acbd17aa998b55dd7d..b61b11d843fa9eca1911d4f714abe7d2e9815856 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 2660c1d8becb0fd7921fed01194aa1bc3e6441b7..deae69ed435a0eb96d1dff908fc8bd8b59a7e05c 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")