diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 3407e2101e..401044099f 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -1,6 +1,11 @@ # Release Notes -## PyMC3 3.8 (on deck) +## PyMC3 3.9 (On deck) + +### New features +- use [fastprogress](https://github.com/fastai/fastprogress) instead of tqdm [#3693](https://github.com/pymc-devs/pymc3/pull/3693) + +## PyMC3 3.8 (November 29 2019) ### New features - Implemented robust u turn check in NUTS (similar to stan-dev/stan#2800). See PR [#3605] diff --git a/pymc3/parallel_sampling.py b/pymc3/parallel_sampling.py index 36a529cdc1..7dacf5607d 100644 --- a/pymc3/parallel_sampling.py +++ b/pymc3/parallel_sampling.py @@ -9,6 +9,7 @@ import errno import numpy as np +from fastprogress import progress_bar from . import theanof @@ -17,28 +18,31 @@ def _get_broken_pipe_exception(): import sys - if sys.platform == 'win32': - return RuntimeError("The communication pipe between the main process " - "and its spawned children is broken.\n" - "In Windows OS, this usually means that the child " - "process raised an exception while it was being " - "spawned, before it was setup to communicate to " - "the main process.\n" - "The exceptions raised by the child process while " - "spawning cannot be caught or handled from the " - "main process, and when running from an IPython or " - "jupyter notebook interactive kernel, the child's " - "exception and traceback appears to be lost.\n" - "A known way to see the child's error, and try to " - "fix or handle it, is to run the problematic code " - "as a batch script from a system's Command Prompt. " - "The child's exception will be printed to the " - "Command Promt's stderr, and it should be visible " - "above this error and traceback.\n" - "Note that if running a jupyter notebook that was " - "invoked from a Command Prompt, the child's " - "exception should have been printed to the Command " - "Prompt on which the notebook is running.") + + if sys.platform == "win32": + return RuntimeError( + "The communication pipe between the main process " + "and its spawned children is broken.\n" + "In Windows OS, this usually means that the child " + "process raised an exception while it was being " + "spawned, before it was setup to communicate to " + "the main process.\n" + "The exceptions raised by the child process while " + "spawning cannot be caught or handled from the " + "main process, and when running from an IPython or " + "jupyter notebook interactive kernel, the child's " + "exception and traceback appears to be lost.\n" + "A known way to see the child's error, and try to " + "fix or handle it, is to run the problematic code " + "as a batch script from a system's Command Prompt. " + "The child's exception will be printed to the " + "Command Promt's stderr, and it should be visible " + "above this error and traceback.\n" + "Note that if running a jupyter notebook that was " + "invoked from a Command Prompt, the child's " + "exception should have been printed to the Command " + "Prompt on which the notebook is running." + ) else: return None @@ -237,7 +241,6 @@ def __init__(self, draws, tune, step_method, chain, seed, start): tune, seed, ) - # We fork right away, so that the main process can start tqdm threads try: self._process.start() except IOError as e: @@ -346,8 +349,6 @@ def __init__( start_chain_num=0, progressbar=True, ): - if progressbar: - from tqdm import tqdm if any(len(arg) != chains for arg in [seeds, start_points]): raise ValueError("Number of seeds and start_points must be %s." % chains) @@ -369,14 +370,13 @@ def __init__( self._progress = None self._divergences = 0 + self._total_draws = 0 self._desc = "Sampling {0._chains:d} chains, {0._divergences:,d} divergences" self._chains = chains - if progressbar: - self._progress = tqdm( - total=chains * (draws + tune), - unit="draws", - desc=self._desc.format(self) - ) + self._progress = progress_bar( + range(chains * (draws + tune)), display=progressbar, auto_update=False + ) + self._progress.comment = self._desc.format(self) def _make_active(self): while self._inactive and len(self._active) < self._max_active: @@ -393,11 +393,11 @@ def __iter__(self): while self._active: draw = ProcessAdapter.recv_draw(self._active) proc, is_last, draw, tuning, stats, warns = draw - if self._progress is not None: - if not tuning and stats and stats[0].get('diverging'): - self._divergences += 1 - self._progress.set_description(self._desc.format(self)) - self._progress.update() + self._total_draws += 1 + if not tuning and stats and stats[0].get("diverging"): + self._divergences += 1 + self._progress.comment = self._desc.format(self) + self._progress.update(self._total_draws) if is_last: proc.join() @@ -423,8 +423,7 @@ def __enter__(self): def __exit__(self, *args): ProcessAdapter.terminate_all(self._samplers) - if self._progress is not None: - self._progress.close() + def _cpu_count(): """Try to guess the number of CPUs in the system. diff --git a/pymc3/sampling.py b/pymc3/sampling.py index 6c2fa4354c..92a346b9eb 100644 --- a/pymc3/sampling.py +++ b/pymc3/sampling.py @@ -40,7 +40,7 @@ from .parallel_sampling import _cpu_count from pymc3.step_methods.hmc import quadpotential import pymc3 as pm -from tqdm import tqdm +from fastprogress import progress_bar import sys @@ -568,11 +568,17 @@ def _sample_population( # create the generator that iterates all chains in parallel chains = [chain + c for c in range(chains)] sampling = _prepare_iter_population( - draws, chains, step, start, parallelize, tune=tune, model=model, random_seed=random_seed + draws, + chains, + step, + start, + parallelize, + tune=tune, + model=model, + random_seed=random_seed, ) - if progressbar: - sampling = tqdm(sampling, total=draws) + sampling = progress_bar(sampling, total=draws, display=progressbar) latest_traces = None for it, traces in enumerate(sampling): @@ -596,10 +602,10 @@ def _sample( sampling = _iter_sample(draws, step, start, trace, chain, tune, model, random_seed) _pbar_data = None - if progressbar: - _pbar_data = {"chain": chain, "divergences": 0} - _desc = "Sampling chain {chain:d}, {divergences:,d} divergences" - sampling = tqdm(sampling, total=draws, desc=_desc.format(**_pbar_data)) + _pbar_data = {"chain": chain, "divergences": 0} + _desc = "Sampling chain {chain:d}, {divergences:,d} divergences" + sampling = progress_bar(sampling, total=draws, display=progressbar) + sampling.comment = _desc.format(**_pbar_data) try: strace = None for it, (strace, diverging) in enumerate(sampling): @@ -607,12 +613,9 @@ def _sample( trace = MultiTrace([strace]) if diverging and _pbar_data is not None: _pbar_data["divergences"] += 1 - sampling.set_description(_desc.format(**_pbar_data)) + sampling.comment = _desc.format(**_pbar_data) except KeyboardInterrupt: pass - finally: - if progressbar: - sampling.close() return strace @@ -753,7 +756,7 @@ def __init__(self, steppers, parallelize): ) import multiprocessing - for c, stepper in enumerate(tqdm(steppers)): + for c, stepper in enumerate(progress_bar(steppers)): slave_end, master_end = multiprocessing.Pipe() stepper_dumps = pickle.dumps(stepper, protocol=4) process = multiprocessing.Process( @@ -1235,9 +1238,13 @@ def sample_posterior_predictive( nchain = 1 if keep_size and samples is not None: - raise IncorrectArgumentsError("Should not specify both keep_size and samples argukments") + raise IncorrectArgumentsError( + "Should not specify both keep_size and samples argukments" + ) if keep_size and size is not None: - raise IncorrectArgumentsError("Should not specify both keep_size and size argukments") + raise IncorrectArgumentsError( + "Should not specify both keep_size and size argukments" + ) if samples is None: samples = sum(len(v) for v in trace._straces.values()) @@ -1253,7 +1260,9 @@ def sample_posterior_predictive( if var_names is not None: if vars is not None: - raise IncorrectArgumentsError("Should not specify both vars and var_names arguments.") + raise IncorrectArgumentsError( + "Should not specify both vars and var_names arguments." + ) else: vars = [model[x] for x in var_names] elif vars is not None: # var_names is None, and vars is not. @@ -1266,8 +1275,7 @@ def sample_posterior_predictive( indices = np.arange(samples) - if progressbar: - indices = tqdm(indices, total=samples) + indices = progress_bar(indices, total=samples, display=progressbar) ppc_trace_t = _DefaultTrace(samples) try: @@ -1285,10 +1293,6 @@ def sample_posterior_predictive( except KeyboardInterrupt: pass - finally: - if progressbar: - indices.close() - ppc_trace = ppc_trace_t.trace_dict if keep_size: for k, ary in ppc_trace.items(): @@ -1411,8 +1415,7 @@ def sample_posterior_predictive_w( indices = np.random.randint(0, len_trace, samples) - if progressbar: - indices = tqdm(indices, total=samples) + indices = progress_bar(indices, total=samples, display=progressbar) try: ppc = defaultdict(list) @@ -1426,10 +1429,6 @@ def sample_posterior_predictive_w( except KeyboardInterrupt: pass - finally: - if progressbar: - indices.close() - return {k: np.asarray(v) for k, v in ppc.items()} diff --git a/pymc3/smc/smc.py b/pymc3/smc/smc.py index d9abbcd4f6..8a614eb895 100644 --- a/pymc3/smc/smc.py +++ b/pymc3/smc/smc.py @@ -1,7 +1,7 @@ from collections import OrderedDict import numpy as np -from tqdm import tqdm +from fastprogress import progress_bar import multiprocessing as mp import warnings from theano import function as theano_function @@ -91,7 +91,9 @@ def initialize_population(self): var_info = OrderedDict() if self.start is None: init_rnd = sample_prior_predictive( - self.draws, var_names=[v.name for v in self.model.unobserved_RVs], model=self.model + self.draws, + var_names=[v.name for v in self.model.unobserved_RVs], + model=self.model, ) else: init_rnd = self.start @@ -103,7 +105,9 @@ def initialize_population(self): for i in range(self.draws): - point = Point({v.name: init_rnd[v.name][i] for v in self.variables}, model=self.model) + point = Point( + {v.name: init_rnd[v.name][i] for v in self.variables}, model=self.model + ) population.append(self.model.dict_to_array(point)) self.posterior = np.array(floatX(population)) @@ -119,7 +123,9 @@ def setup_kernel(self): if self.kernel.lower() == "abc": warnings.warn(EXPERIMENTAL_WARNING) if len(self.model.observed_RVs) != 1: - warnings.warn("SMC-ABC only works properly with models with one observed variable") + warnings.warn( + "SMC-ABC only works properly with models with one observed variable" + ) simulator = self.model.observed_RVs[0] self.likelihood_logp = PseudoLikelihood( self.epsilon, @@ -131,7 +137,9 @@ def setup_kernel(self): self.sum_stat, ) elif self.kernel.lower() == "metropolis": - self.likelihood_logp = logp_forw([self.model.datalogpt], self.variables, shared) + self.likelihood_logp = logp_forw( + [self.model.datalogpt], self.variables, shared + ) def initialize_logp(self): """ @@ -139,7 +147,9 @@ def initialize_logp(self): """ if self.parallel and self.cores > 1: self.pool = mp.Pool(processes=self.cores) - priors = self.pool.starmap(self.prior_logp, [(sample,) for sample in self.posterior]) + priors = self.pool.starmap( + self.prior_logp, [(sample,) for sample in self.posterior] + ) likelihoods = self.pool.starmap( self.likelihood_logp, [(sample,) for sample in self.posterior] ) @@ -204,14 +214,18 @@ def update_proposal(self): cov = np.atleast_2d(cov) cov += 1e-6 * np.eye(cov.shape[0]) if np.isnan(cov).any() or np.isinf(cov).any(): - raise ValueError('Sample covariances not valid! Likely "draws" is too small!') + raise ValueError( + 'Sample covariances not valid! Likely "draws" is too small!' + ) self.proposal = MultivariateNormalProposal(cov) def tune(self): """ Tune scaling and n_steps based on the acceptance rate. """ - ave_scaling = np.exp(np.log(self.scalings.mean()) + (self.acc_per_chain.mean() - 0.234)) + ave_scaling = np.exp( + np.log(self.scalings.mean()) + (self.acc_per_chain.mean() - 0.234) + ) self.scalings = 0.5 * ( ave_scaling + np.exp(np.log(self.scalings) + (self.acc_per_chain - 0.234)) ) @@ -219,7 +233,8 @@ def tune(self): if self.tune_steps: acc_rate = max(1.0 / self.proposed, self.acc_rate) self.n_steps = min( - self.max_steps, max(2, int(np.log(1 - self.p_acc_rate) / np.log(1 - acc_rate))) + self.max_steps, + max(2, int(np.log(1 - self.p_acc_rate) / np.log(1 - acc_rate))), ) self.proposed = self.draws * self.n_steps @@ -264,7 +279,7 @@ def mutate(self): draw, *parameters ) - for draw in tqdm(range(self.draws), disable=not self.progressbar) + for draw in progress_bar(range(self.draws), display=self.progressbar) ] posterior, acc_list, priors, likelihoods = zip(*results) self.posterior = np.array(posterior) @@ -335,7 +350,9 @@ def metrop_kernel( new_tempered_logp = pl + ll * beta - q_old, accept = metrop_select(new_tempered_logp - old_tempered_logp, q_new, q_old) + q_old, accept = metrop_select( + new_tempered_logp - old_tempered_logp, q_new, q_old + ) if accept: accepted += 1 @@ -369,7 +386,9 @@ class PseudoLikelihood: Pseudo Likelihood """ - def __init__(self, epsilon, observations, function, model, var_info, distance, sum_stat): + def __init__( + self, epsilon, observations, function, model, var_info, distance, sum_stat + ): """ epsilon : float Standard deviation of the gaussian pseudo likelihood. @@ -419,7 +438,9 @@ def posterior_to_function(self, posterior): def gauss_kernel(self, value): epsilon = self.epsilon - return (-(value ** 2) / epsilon ** 2 + np.log(1 / (2 * np.pi * epsilon ** 2))) / 2.0 + return ( + -(value ** 2) / epsilon ** 2 + np.log(1 / (2 * np.pi * epsilon ** 2)) + ) / 2.0 def absolute_error(self, a, b): if self.sum_stat: diff --git a/pymc3/tuning/starting.py b/pymc3/tuning/starting.py index 4435ed6b01..96ce62470f 100644 --- a/pymc3/tuning/starting.py +++ b/pymc3/tuning/starting.py @@ -6,7 +6,7 @@ from scipy.optimize import minimize import numpy as np from numpy import isfinite, nan_to_num -from tqdm import tqdm +from fastprogress import progress_bar import pymc3 as pm from ..vartypes import discrete_types, typefilter from ..model import modelcontext, Point @@ -68,7 +68,6 @@ def find_MAP( As a result, we have greatly enhanced the initialization of NUTS and wrapped it inside pymc3.sample() and you should thus avoid this method. """ - model = modelcontext(model) if start is None: start = model.test_point @@ -126,9 +125,7 @@ def find_MAP( def grad_logp(point): return nan_to_num(-dlogp_func(point)) - opt_result = fmin( - cost_func, bij.map(start), fprime=grad_logp, *args, **kwargs - ) + opt_result = fmin(cost_func, bij.map(start), fprime=grad_logp, *args, **kwargs) else: # Check to see if minimization function uses a starting value if "x0" in getargspec(fmin).args: @@ -152,20 +149,17 @@ def grad_logp(point): cost_func, x0, method=method, jac=compute_gradient, *args, **kwargs ) mx0 = opt_result["x"] # r -> opt_result - cost_func.progress.total = cost_func.progress.n + 1 - cost_func.progress.update() except (KeyboardInterrupt, StopIteration) as e: mx0, opt_result = cost_func.previous_x, None - cost_func.progress.close() if isinstance(e, StopIteration): pm._log.info(e) finally: - cost_func.progress.close() + last_v = cost_func.n_eval + cost_func.progress.total = last_v + cost_func.progress.update(last_v) vars = get_default_varnames(model.unobserved_RVs, include_transformed) - mx = { - var.name: value for var, value in zip(vars, model.fastfn(vars)(bij.rmap(mx0))) - } + mx = {var.name: value for var, value in zip(vars, model.fastfn(vars)(bij.rmap(mx0)))} if return_raw: return mx, opt_result @@ -200,8 +194,8 @@ def __init__(self, maxeval=5000, progressbar=True, logp_func=None, dlogp_func=No self.use_gradient = True self.desc = "logp = {:,.5g}, ||grad|| = {:,.5g}" self.previous_x = None - self.progress = tqdm(total=maxeval, disable=not progressbar) - self.progress.n = 0 + self.progress = progress_bar(range(maxeval), total=maxeval, display=progressbar) + self.progress.update(0) def __call__(self, x): neg_value = np.float64(self.logp_func(pm.floatX(x))) @@ -221,11 +215,10 @@ def __call__(self, x): if self.n_eval > self.maxeval: self.update_progress_desc(neg_value, grad) - self.progress.close() raise StopIteration self.n_eval += 1 - self.progress.update(1) + self.progress.update_bar(self.n_eval) if self.use_gradient: return value, grad @@ -234,7 +227,7 @@ def __call__(self, x): def update_progress_desc(self, neg_value, grad=None): if grad is None: - self.progress.set_description(self.desc.format(neg_value)) + self.progress.comment = self.desc.format(neg_value) else: norm_grad = np.linalg.norm(grad) - self.progress.set_description(self.desc.format(neg_value, norm_grad)) + self.progress.comment = self.desc.format(neg_value, norm_grad) diff --git a/pymc3/variational/inference.py b/pymc3/variational/inference.py index 23185aec7d..0ed1d0fa45 100644 --- a/pymc3/variational/inference.py +++ b/pymc3/variational/inference.py @@ -3,12 +3,15 @@ import collections import numpy as np -import tqdm +from fastprogress import progress_bar import pymc3 as pm from pymc3.variational import test_functions from pymc3.variational.approximations import ( - MeanField, FullRank, Empirical, NormalizingFlow + MeanField, + FullRank, + Empirical, + NormalizingFlow, ) from pymc3.variational.operators import KL, KSD from . import opvi @@ -16,22 +19,22 @@ logger = logging.getLogger(__name__) __all__ = [ - 'ADVI', - 'FullRankADVI', - 'SVGD', - 'ASVGD', - 'NFVI', - 'Inference', - 'ImplicitGradient', - 'KLqp', - 'fit' + "ADVI", + "FullRankADVI", + "SVGD", + "ASVGD", + "NFVI", + "Inference", + "ImplicitGradient", + "KLqp", + "fit", ] -State = collections.namedtuple('State', 'i,step,callbacks,score') +State = collections.namedtuple("State", "i,step,callbacks,score") class Inference: - R"""**Base class for Variational Inference** + r"""**Base class for Variational Inference** Communicates Operator, Approximation and Test Function to build Objective Function @@ -57,9 +60,10 @@ def _maybe_score(self, score): if score is None: score = returns_loss elif score and not returns_loss: - warnings.warn('method `fit` got `score == True` but %s ' - 'does not return loss. Ignoring `score` argument' - % self.objective.op) + warnings.warn( + "method `fit` got `score == True` but %s " + "does not return loss. Ignoring `score` argument" % self.objective.op + ) score = False else: pass @@ -67,24 +71,20 @@ def _maybe_score(self, score): def run_profiling(self, n=1000, score=None, **kwargs): score = self._maybe_score(score) - fn_kwargs = kwargs.pop('fn_kwargs', dict()) - fn_kwargs['profile'] = True + fn_kwargs = kwargs.pop("fn_kwargs", dict()) + fn_kwargs["profile"] = True step_func = self.objective.step_function( - score=score, fn_kwargs=fn_kwargs, - **kwargs + score=score, fn_kwargs=fn_kwargs, **kwargs ) - progress = tqdm.trange(n) + progress = progress_bar(range(n)) try: for _ in progress: step_func() except KeyboardInterrupt: pass - finally: - progress.close() return step_func.profile - def fit(self, n=10000, score=None, callbacks=None, progressbar=True, - **kwargs): + def fit(self, n=10000, score=None, callbacks=None, progressbar=True, **kwargs): """Perform Operator Variational Inference Parameters @@ -129,11 +129,11 @@ def fit(self, n=10000, score=None, callbacks=None, progressbar=True, callbacks = [] score = self._maybe_score(score) step_func = self.objective.step_function(score=score, **kwargs) - with tqdm.trange(n, disable=not progressbar) as progress: - if score: - state = self._iterate_with_loss(0, n, step_func, progress, callbacks) - else: - state = self._iterate_without_loss(0, n, step_func, progress, callbacks) + progress = progress_bar(range(n), display=progressbar) + if score: + state = self._iterate_with_loss(0, n, step_func, progress, callbacks) + else: + state = self._iterate_without_loss(0, n, step_func, progress, callbacks) # hack to allow pm.fit() access to loss hist self.approx.hist = self.hist @@ -156,37 +156,37 @@ def _iterate_without_loss(self, s, _, step_func, progress, callbacks): for i in range(slclen): name_slc.append((vmap_.var, i)) index = np.where(np.isnan(current_param))[0] - errmsg = ['NaN occurred in optimization. '] - suggest_solution = 'Try tracking this parameter: ' \ - 'http://docs.pymc.io/notebooks/variational_api_quickstart.html#Tracking-parameters' + errmsg = ["NaN occurred in optimization. "] + suggest_solution = ( + "Try tracking this parameter: " + "http://docs.pymc.io/notebooks/variational_api_quickstart.html#Tracking-parameters" + ) try: for ii in index: - errmsg.append('The current approximation of RV `{}`.ravel()[{}]' - ' is NaN.'.format(*name_slc[ii])) + errmsg.append( + "The current approximation of RV `{}`.ravel()[{}]" + " is NaN.".format(*name_slc[ii]) + ) errmsg.append(suggest_solution) except IndexError: pass - raise FloatingPointError('\n'.join(errmsg)) + raise FloatingPointError("\n".join(errmsg)) for callback in callbacks: - callback(self.approx, None, i+s+1) + callback(self.approx, None, i + s + 1) except (KeyboardInterrupt, StopIteration) as e: - progress.close() if isinstance(e, StopIteration): logger.info(str(e)) - finally: - progress.close() - return State(i+s, step=step_func, - callbacks=callbacks, - score=False) + return State(i + s, step=step_func, callbacks=callbacks, score=False) def _iterate_with_loss(self, s, n, step_func, progress, callbacks): def _infmean(input_array): """Return the mean of the finite values of the array""" - input_array = input_array[np.isfinite(input_array)].astype('float64') + input_array = input_array[np.isfinite(input_array)].astype("float64") if len(input_array) == 0: return np.nan else: return np.mean(input_array) + scores = np.empty(n) scores[:] = np.nan i = 0 @@ -205,65 +205,67 @@ def _infmean(input_array): for i in range(slclen): name_slc.append((vmap_.var, i)) index = np.where(np.isnan(current_param))[0] - errmsg = ['NaN occurred in optimization. '] - suggest_solution = 'Try tracking this parameter: ' \ - 'http://docs.pymc.io/notebooks/variational_api_quickstart.html#Tracking-parameters' + errmsg = ["NaN occurred in optimization. "] + suggest_solution = ( + "Try tracking this parameter: " + "http://docs.pymc.io/notebooks/variational_api_quickstart.html#Tracking-parameters" + ) try: for ii in index: - errmsg.append('The current approximation of RV `{}`.ravel()[{}]' - ' is NaN.'.format(*name_slc[ii])) + errmsg.append( + "The current approximation of RV `{}`.ravel()[{}]" + " is NaN.".format(*name_slc[ii]) + ) errmsg.append(suggest_solution) except IndexError: pass - raise FloatingPointError('\n'.join(errmsg)) + raise FloatingPointError("\n".join(errmsg)) scores[i] = e if i % 10 == 0: - avg_loss = _infmean(scores[max(0, i - 1000):i + 1]) - progress.set_description('Average Loss = {:,.5g}'.format(avg_loss)) - avg_loss = scores[max(0, i - 1000):i + 1].mean() - progress.set_description( - 'Average Loss = {:,.5g}'.format(avg_loss)) + avg_loss = _infmean(scores[max(0, i - 1000) : i + 1]) + progress.comment = "Average Loss = {:,.5g}".format(avg_loss) + avg_loss = scores[max(0, i - 1000) : i + 1].mean() + progress.comment = "Average Loss = {:,.5g}".format(avg_loss) for callback in callbacks: - callback(self.approx, scores[:i + 1], i+s+1) + callback(self.approx, scores[: i + 1], i + s + 1) except (KeyboardInterrupt, StopIteration) as e: # pragma: no cover # do not print log on the same line - progress.close() scores = scores[:i] if isinstance(e, StopIteration): logger.info(str(e)) if n < 10: - logger.info('Interrupted at {:,d} [{:.0f}%]: Loss = {:,.5g}'.format( - i, 100 * i // n, scores[i])) + logger.info( + "Interrupted at {:,d} [{:.0f}%]: Loss = {:,.5g}".format( + i, 100 * i // n, scores[i] + ) + ) else: - avg_loss = _infmean(scores[min(0, i - 1000):i + 1]) - logger.info('Interrupted at {:,d} [{:.0f}%]: Average Loss = {:,.5g}'.format( - i, 100 * i // n, avg_loss)) + avg_loss = _infmean(scores[min(0, i - 1000) : i + 1]) + logger.info( + "Interrupted at {:,d} [{:.0f}%]: Average Loss = {:,.5g}".format( + i, 100 * i // n, avg_loss + ) + ) else: if n < 10: - logger.info( - 'Finished [100%]: Loss = {:,.5g}'.format(scores[-1])) + logger.info("Finished [100%]: Loss = {:,.5g}".format(scores[-1])) else: - avg_loss = _infmean(scores[max(0, i - 1000):i + 1]) - logger.info( - 'Finished [100%]: Average Loss = {:,.5g}'.format(avg_loss)) - finally: - progress.close() + avg_loss = _infmean(scores[max(0, i - 1000) : i + 1]) + logger.info("Finished [100%]: Average Loss = {:,.5g}".format(avg_loss)) self.hist = np.concatenate([self.hist, scores]) - return State(i+s, step=step_func, - callbacks=callbacks, - score=True) + return State(i + s, step=step_func, callbacks=callbacks, score=True) def refine(self, n, progressbar=True): """Refine the solution using the last compiled step function """ if self.state is None: - raise TypeError('Need to call `.fit` first') + raise TypeError("Need to call `.fit` first") i, step, callbacks, score = self.state - with tqdm.trange(n, disable=not progressbar) as progress: - if score: - state = self._iterate_with_loss(i, n, step, progress, callbacks) - else: - state = self._iterate_without_loss(i, n, step, progress, callbacks) + progress = progress_bar(n, display=progressbar) + if score: + state = self._iterate_with_loss(i, n, step, progress, callbacks) + else: + state = self._iterate_without_loss(i, n, step, progress, callbacks) self.state = state @@ -291,12 +293,13 @@ class KLqp(Inference): Understanding disentangling in :math:`\beta`-VAE arXiv preprint 1804.03599 """ - def __init__(self, approx, beta=1.): + + def __init__(self, approx, beta=1.0): super().__init__(KL, approx, None, beta=beta) class ADVI(KLqp): - R"""**Automatic Differentiation Variational Inference (ADVI)** + r"""**Automatic Differentiation Variational Inference (ADVI)** This class implements the meanfield ADVI, where the variational posterior distribution is assumed to be spherical Gaussian without @@ -444,7 +447,7 @@ def __init__(self, *args, **kwargs): class FullRankADVI(KLqp): - R"""**Full Rank Automatic Differentiation Variational Inference (ADVI)** + r"""**Full Rank Automatic Differentiation Variational Inference (ADVI)** Parameters ---------- @@ -489,12 +492,13 @@ class ImplicitGradient(Inference): only for large number of samples. Larger temperature is needed for small number of samples but there is no theoretical approach to choose the best one in such case. """ + def __init__(self, approx, estimator=KSD, kernel=test_functions.rbf, **kwargs): super().__init__(op=estimator, approx=approx, tf=kernel, **kwargs) class SVGD(ImplicitGradient): - R"""**Stein Variational Gradient Descent** + r"""**Stein Variational Gradient Descent** This inference is based on Kernelized Stein Discrepancy it's main idea is to move initial noisy particles so that @@ -544,18 +548,31 @@ class SVGD(ImplicitGradient): arXiv:1704.02399 """ - def __init__(self, n_particles=100, jitter=1, model=None, start=None, - random_seed=None, estimator=KSD, kernel=test_functions.rbf, **kwargs): - if kwargs.get('local_rv') is not None: - raise opvi.AEVBInferenceError('SVGD does not support local groups') + def __init__( + self, + n_particles=100, + jitter=1, + model=None, + start=None, + random_seed=None, + estimator=KSD, + kernel=test_functions.rbf, + **kwargs, + ): + if kwargs.get("local_rv") is not None: + raise opvi.AEVBInferenceError("SVGD does not support local groups") empirical = Empirical( - size=n_particles, jitter=jitter, - start=start, model=model, random_seed=random_seed) + size=n_particles, + jitter=jitter, + start=start, + model=model, + random_seed=random_seed, + ) super().__init__(approx=empirical, estimator=estimator, kernel=kernel, **kwargs) class ASVGD(ImplicitGradient): - R"""**Amortized Stein Variational Gradient Descent** + r"""**Amortized Stein Variational Gradient Descent** **not suggested to use** @@ -600,32 +617,45 @@ class ASVGD(ImplicitGradient): """ def __init__(self, approx=None, estimator=KSD, kernel=test_functions.rbf, **kwargs): - warnings.warn('You are using experimental inference Operator. ' - 'It requires careful choice of temperature, default is 1. ' - 'Default temperature works well for low dimensional problems and ' - 'for significant `n_obj_mc`. Temperature > 1 gives more exploration ' - 'power to algorithm, < 1 leads to undesirable results. Please take ' - 'it in account when looking at inference result. Posterior variance ' - 'is often **underestimated** when using temperature = 1.') + warnings.warn( + "You are using experimental inference Operator. " + "It requires careful choice of temperature, default is 1. " + "Default temperature works well for low dimensional problems and " + "for significant `n_obj_mc`. Temperature > 1 gives more exploration " + "power to algorithm, < 1 leads to undesirable results. Please take " + "it in account when looking at inference result. Posterior variance " + "is often **underestimated** when using temperature = 1." + ) if approx is None: approx = FullRank( - model=kwargs.pop('model', None), - local_rv=kwargs.pop('local_rv', None) + model=kwargs.pop("model", None), local_rv=kwargs.pop("local_rv", None) ) super().__init__(estimator=estimator, approx=approx, kernel=kernel, **kwargs) - def fit(self, n=10000, score=None, callbacks=None, progressbar=True, - obj_n_mc=500, **kwargs): + def fit( + self, + n=10000, + score=None, + callbacks=None, + progressbar=True, + obj_n_mc=500, + **kwargs, + ): return super().fit( - n=n, score=score, callbacks=callbacks, - progressbar=progressbar, obj_n_mc=obj_n_mc, **kwargs) + n=n, + score=score, + callbacks=callbacks, + progressbar=progressbar, + obj_n_mc=obj_n_mc, + **kwargs, + ) def run_profiling(self, n=1000, score=None, obj_n_mc=500, **kwargs): return super().run_profiling(n=n, score=score, obj_n_mc=obj_n_mc, **kwargs) class NFVI(KLqp): - R"""**Normalizing Flow based :class:`KLqp` inference** + r"""**Normalizing Flow based :class:`KLqp` inference** Normalizing flow is a series of invertible transformations on initial distribution. @@ -679,9 +709,17 @@ def __init__(self, *args, **kwargs): super().__init__(NormalizingFlow(*args, **kwargs)) -def fit(n=10000, local_rv=None, method='advi', model=None, - random_seed=None, start=None, inf_kwargs=None, **kwargs): - R"""Handy shortcut for using inference methods in functional way +def fit( + n=10000, + local_rv=None, + method="advi", + model=None, + random_seed=None, + start=None, + inf_kwargs=None, + **kwargs, +): + r"""Handy shortcut for using inference methods in functional way Parameters ---------- @@ -749,42 +787,32 @@ def fit(n=10000, local_rv=None, method='advi', model=None, else: inf_kwargs = inf_kwargs.copy() if local_rv is not None: - inf_kwargs['local_rv'] = local_rv + inf_kwargs["local_rv"] = local_rv if random_seed is not None: - inf_kwargs['random_seed'] = random_seed + inf_kwargs["random_seed"] = random_seed if start is not None: - inf_kwargs['start'] = start + inf_kwargs["start"] = start if model is None: model = pm.modelcontext(model) _select = dict( - advi=ADVI, - fullrank_advi=FullRankADVI, - svgd=SVGD, - asvgd=ASVGD, - nfvi=NFVI + advi=ADVI, fullrank_advi=FullRankADVI, svgd=SVGD, asvgd=ASVGD, nfvi=NFVI ) if isinstance(method, str): method = method.lower() - if method.startswith('nfvi='): + if method.startswith("nfvi="): formula = method[5:] - inference = NFVI( - formula, - **inf_kwargs - ) + inference = NFVI(formula, **inf_kwargs) elif method in _select: - inference = _select[method]( - model=model, - **inf_kwargs - ) + inference = _select[method](model=model, **inf_kwargs) else: - raise KeyError('method should be one of %s ' - 'or Inference instance' % - set(_select.keys())) + raise KeyError( + f"method should be one of {set(_select.keys())} or Inference instance" + ) elif isinstance(method, Inference): inference = method else: - raise TypeError('method should be one of %s ' - 'or Inference instance' % - set(_select.keys())) + raise TypeError( + f"method should be one of {set(_select.keys())} or Inference instance" + ) return inference.fit(n, **kwargs) diff --git a/requirements.txt b/requirements.txt index 5beda3ba46..489fb83aac 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,5 +4,5 @@ numpy>=1.13.0 scipy>=0.18.1 pandas>=0.18.0 patsy>=0.4.0 -tqdm>=4.8.4 +fastprogress>=0.1.21 h5py>=2.7.0