diff options
| -rw-r--r-- | ml_exp/__init__.py | 4 | ||||
| -rw-r--r-- | ml_exp/misc.py | 38 |
2 files changed, 20 insertions, 22 deletions
diff --git a/ml_exp/__init__.py b/ml_exp/__init__.py index d685e7ccc..c62141a58 100644 --- a/ml_exp/__init__.py +++ b/ml_exp/__init__.py @@ -24,6 +24,7 @@ from ml_exp.compound import Compound from ml_exp.representations import coulomb_matrix, lennard_jones_matrix,\ first_neighbor_matrix, adjacency_matrix, check_bond, bag_of_stuff from ml_exp.math import cholesky_solve +from ml_exp.qm7db import qm7db __all__ = ['Compound', 'coulomb_matrix', @@ -32,4 +33,5 @@ __all__ = ['Compound', 'adjacency_matrix', 'check_bond', 'bag_of_stuff', - 'cholesky_solve'] + 'cholesky_solve', + 'qm7db'] diff --git a/ml_exp/misc.py b/ml_exp/misc.py index e9142b05f..4bd68eefc 100644 --- a/ml_exp/misc.py +++ b/ml_exp/misc.py @@ -29,29 +29,25 @@ init() def printc(text, color): """ Prints texts normaly, but in color. Using colorama. - text: string with the text to print. + text: text to print. color: color to be used, same as available in colorama. """ - color_dic = {'BLACK': Fore.BLACK, - 'RED': Fore.RED, - 'GREEN': Fore.GREEN, - 'YELLOW': Fore.YELLOW, - 'BLUE': Fore.BLUE, - 'MAGENTA': Fore.MAGENTA, - 'CYAN': Fore.CYAN, - 'WHITE': Fore.WHITE, - 'RESET': Fore.RESET} - - color_dic_keys = color_dic.keys() - if color not in color_dic_keys: - print(Fore.RED - + '\'{}\' not found, using default color.'.format(color) - + Style.RESET_ALL) - actual_color = Fore.RESET - else: - actual_color = color_dic[color] - - print(actual_color + text + Style.RESET_ALL) + color = color.upper() + colors = {'BLACK': Fore.BLACK, + 'RED': Fore.RED, + 'GREEN': Fore.GREEN, + 'YELLOW': Fore.YELLOW, + 'BLUE': Fore.BLUE, + 'MAGENTA': Fore.MAGENTA, + 'CYAN': Fore.CYAN, + 'WHITE': Fore.WHITE, + 'RESET': Fore.RESET} + + av_colors = colors.keys() + if color not in av_colors: + raise KeyError(f'{color} not found.') + + print(colors[color] + text + Style.RESET_ALL) def plot_benchmarks(): |