From 1f02e8dda3900799ae55447b75bd42eba48ba62a Mon Sep 17 00:00:00 2001
From: Colm Talbot <colm.talbot@ligo.org>
Date: Wed, 23 Dec 2020 00:16:19 -0600
Subject: [PATCH] Resolve "Use tqdm.auto for better notebook performance"
---
bilby/core/sampler/dynesty.py | 4 ++--
bilby/core/sampler/emcee.py | 7 ++-----
bilby/core/sampler/kombine.py | 5 ++---
bilby/core/utils.py | 20 --------------------
bilby/gw/conversion.py | 3 +--
5 files changed, 7 insertions(+), 32 deletions(-)
diff --git a/bilby/core/sampler/dynesty.py b/bilby/core/sampler/dynesty.py
index 22811bccb..206ed9a42 100644
--- a/bilby/core/sampler/dynesty.py
+++ b/bilby/core/sampler/dynesty.py
@@ -6,7 +6,7 @@ import pickle
import signal
import time
-import tqdm
+from tqdm.auto import tqdm
import matplotlib.pyplot as plt
import numpy as np
from pandas import DataFrame
@@ -224,7 +224,7 @@ class Dynesty(NestedSampler):
self.kwargs['update_interval'] = int(0.6 * self.kwargs['nlive'])
if self.kwargs['print_func'] is None:
self.kwargs['print_func'] = self._print_func
- self.pbar = tqdm.tqdm(file=sys.stdout)
+ self.pbar = tqdm(file=sys.stdout)
Sampler._verify_kwargs_against_default_kwargs(self)
def _print_func(self, results, niter, ncall=None, dlogz=None, *args, **kwargs):
diff --git a/bilby/core/sampler/emcee.py b/bilby/core/sampler/emcee.py
index 17867251d..5e83178bd 100644
--- a/bilby/core/sampler/emcee.py
+++ b/bilby/core/sampler/emcee.py
@@ -1,5 +1,3 @@
-from __future__ import absolute_import, print_function
-
from collections import namedtuple
import os
import signal
@@ -12,8 +10,7 @@ from pandas import DataFrame
from distutils.version import LooseVersion
import dill as pickle
-from ..utils import (
- logger, get_progress_bar, check_directory_exists_and_if_not_mkdir)
+from ..utils import logger, check_directory_exists_and_if_not_mkdir
from .base_sampler import MCMCSampler, SamplerError
@@ -353,7 +350,7 @@ class Emcee(MCMCSampler):
self.pos0 = self.sampler.chain[:, -1, :]
def run_sampler(self):
- tqdm = get_progress_bar()
+ from tqdm.auto import tqdm
sampler_function_kwargs = self.sampler_function_kwargs
iterations = sampler_function_kwargs.pop('iterations')
iterations -= self._previous_iterations
diff --git a/bilby/core/sampler/kombine.py b/bilby/core/sampler/kombine.py
index f7c7768ec..e82129d9f 100644
--- a/bilby/core/sampler/kombine.py
+++ b/bilby/core/sampler/kombine.py
@@ -1,5 +1,4 @@
-from __future__ import absolute_import, print_function
-from ..utils import logger, get_progress_bar
+from ..utils import logger
import numpy as np
import os
from .emcee import Emcee
@@ -141,7 +140,7 @@ class Kombine(Emcee):
logger.info("Kombine auto-burnin complete. Removing {} samples from chains".format(self.nburn))
self._set_pos0_for_resume()
- tqdm = get_progress_bar()
+ from tqdm.auto import tqdm
sampler_function_kwargs = self.sampler_function_kwargs
iterations = sampler_function_kwargs.pop('iterations')
iterations -= self._previous_iterations
diff --git a/bilby/core/utils.py b/bilby/core/utils.py
index 03380552d..8fc3431ed 100644
--- a/bilby/core/utils.py
+++ b/bilby/core/utils.py
@@ -508,26 +508,6 @@ def get_version_information():
print("No version information file '.version' found")
-def get_progress_bar(module='tqdm'):
- """
- TODO: Write proper docstring
- """
- if module in ['tqdm']:
- try:
- from tqdm import tqdm
- except ImportError:
- def tqdm(x, *args, **kwargs):
- return x
- return tqdm
- elif module in ['tqdm_notebook']:
- try:
- from tqdm import tqdm_notebook as tqdm
- except ImportError:
- def tqdm(x, *args, **kwargs):
- return x
- return tqdm
-
-
def spherical_to_cartesian(radius, theta, phi):
""" Convert from spherical coordinates to cartesian.
diff --git a/bilby/gw/conversion.py b/bilby/gw/conversion.py
index ad1981434..00a6c9a4e 100644
--- a/bilby/gw/conversion.py
+++ b/bilby/gw/conversion.py
@@ -1,8 +1,7 @@
-from __future__ import division
import sys
import multiprocessing
-from tqdm import tqdm
+from tqdm.auto import tqdm
import numpy as np
from pandas import DataFrame
--
GitLab