1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
|
# -*- coding: utf-8 -*-
"""Some utility functions."""
# Authors: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#
# License: BSD-3-Clause
from collections.abc import Iterable
import os
import os.path as op
import logging
import tempfile
from threading import Thread
import time
import numpy as np
from .check import _check_option
from .config import get_config
from ._logging import logger
class ProgressBar(object):
"""Generate a command-line progressbar.
Parameters
----------
iterable : iterable | int | None
The iterable to use. Can also be an int for backward compatibility
(acts like ``max_value``).
initial_value : int
Initial value of process, useful when resuming process from a specific
value, defaults to 0.
mesg : str
Message to include at end of progress bar.
max_total_width : int | str
Maximum total message width. Can use "auto" (default) to try to set
a sane value based on the current terminal width.
max_value : int | None
The max value. If None, the length of ``iterable`` will be used.
**kwargs : dict
Additional keyword arguments for tqdm.
"""
def __init__(self, iterable=None, initial_value=0, mesg=None,
max_total_width='auto', max_value=None,
**kwargs): # noqa: D102
# The following mimics this, but with configurable module to use
# from ..externals.tqdm import auto
import tqdm
which_tqdm = get_config('MNE_TQDM', 'tqdm.auto')
_check_option('MNE_TQDM', which_tqdm[:5], ('tqdm', 'tqdm.', 'off'),
extra='beginning')
logger.debug(f'Using ProgressBar with {which_tqdm}')
if which_tqdm not in ('tqdm', 'off'):
try:
__import__(which_tqdm)
except Exception as exc:
raise ValueError(
f'Unknown tqdm backend {repr(which_tqdm)}, got: {exc}'
) from None
tqdm = getattr(tqdm, which_tqdm.split('.', 1)[1])
tqdm = tqdm.tqdm
defaults = dict(
leave=True, mininterval=0.016, miniters=1, smoothing=0.05,
bar_format='{percentage:3.0f}%|{bar}| {desc} : {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt:>11}{postfix}]', # noqa: E501
)
for key, val in defaults.items():
if key not in kwargs:
kwargs.update({key: val})
if isinstance(iterable, Iterable):
self.iterable = iterable
if max_value is None:
self.max_value = len(iterable)
else:
self.max_value = max_value
else: # ignore max_value then
self.max_value = int(iterable)
self.iterable = None
if max_total_width == 'auto':
max_total_width = None # tqdm's auto
with tempfile.NamedTemporaryFile('wb', prefix='tmp_mne_prog') as tf:
self._mmap_fname = tf.name
del tf # should remove the file
self._mmap = None
disable = logger.level > logging.INFO or which_tqdm == 'off'
self._tqdm = tqdm(
iterable=self.iterable, desc=mesg, total=self.max_value,
initial=initial_value, ncols=max_total_width,
disable=disable, **kwargs)
def update(self, cur_value):
"""Update progressbar with current value of process.
Parameters
----------
cur_value : number
Current value of process. Should be <= max_value (but this is not
enforced). The percent of the progressbar will be computed as
``(cur_value / max_value) * 100``.
"""
self.update_with_increment_value(cur_value - self._tqdm.n)
def update_with_increment_value(self, increment_value):
"""Update progressbar with an increment.
Parameters
----------
increment_value : int
Value of the increment of process. The percent of the progressbar
will be computed as
``(self.cur_value + increment_value / max_value) * 100``.
"""
self._tqdm.update(increment_value)
def __iter__(self):
"""Iterate to auto-increment the pbar with 1."""
for x in self._tqdm:
yield x
def subset(self, idx):
"""Make a joblib-friendly index subset updater.
Parameters
----------
idx : ndarray
List of indices for this subset.
Returns
-------
updater : instance of PBSubsetUpdater
Class with a ``.update(ii)`` method.
"""
return _PBSubsetUpdater(self, idx)
def __enter__(self): # noqa: D105
# This should only be used with pb.subset and parallelization
if op.isfile(self._mmap_fname):
os.remove(self._mmap_fname)
# prevent corner cases where self.max_value == 0
self._mmap = np.memmap(self._mmap_fname, bool, 'w+',
shape=max(self.max_value, 1))
self.update(0) # must be zero as we just created the memmap
# We need to control how the pickled bars exit: remove print statements
self._thread = _UpdateThread(self)
self._thread.start()
return self
def __exit__(self, type_, value, traceback): # noqa: D105
# Restore exit behavior for our one from the main thread
self.update(self._mmap.sum())
self._tqdm.close()
self._thread._mne_run = False
self._thread.join()
self._mmap = None
if op.isfile(self._mmap_fname):
os.remove(self._mmap_fname)
def __del__(self):
"""Ensure output completes."""
if getattr(self, '_tqdm', None) is not None:
self._tqdm.close()
class _UpdateThread(Thread):
def __init__(self, pb):
super(_UpdateThread, self).__init__(daemon=True)
self._mne_run = True
self._mne_pb = pb
def run(self):
while self._mne_run:
self._mne_pb.update(self._mne_pb._mmap.sum())
time.sleep(1. / 30.) # 30 Hz refresh is plenty
class _PBSubsetUpdater(object):
def __init__(self, pb, idx):
self.mmap = pb._mmap
self.idx = idx
def update(self, ii):
self.mmap[self.idx[ii - 1]] = True
|
