Spaces:
Running
Running
File size: 6,928 Bytes
3bb804c |
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 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 |
"""Some utility functions."""
# Authors: The MNE-Python contributors.
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.
import logging
import os
import os.path as op
import tempfile
import time
from collections.abc import Iterable
from threading import Thread
import numpy as np
from ._logging import logger
from .check import _check_option
from .config import get_config
class ProgressBar:
"""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.
which_tqdm : str | None
Which tqdm module to use. Can be "tqdm", "tqdm.notebook", or "off".
Defaults to ``None``, which uses the value of the MNE_TQDM environment
variable, or ``"tqdm.auto"`` if that is not set.
**kwargs : dict
Additional keyword arguments for tqdm.
"""
def __init__(
self,
iterable=None,
initial_value=0,
mesg=None,
max_total_width="auto",
max_value=None,
*,
which_tqdm=None,
**kwargs,
):
# The following mimics this, but with configurable module to use
# from ..externals.tqdm import auto
import tqdm
if which_tqdm is None:
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``.
"""
try:
self._tqdm.update(increment_value)
except TypeError: # can happen during GC on Windows
pass
def __iter__(self):
"""Iterate to auto-increment the pbar with 1."""
yield from self._tqdm
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):
try:
os.remove(self._mmap_fname)
# happens on Windows sometimes
except PermissionError: # pragma: no cover
pass
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().__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.0 / 30.0) # 30 Hz refresh is plenty
class _PBSubsetUpdater:
def __init__(self, pb, idx):
self.mmap = pb._mmap
self.idx = idx
def update(self, ii):
self.mmap[self.idx[ii - 1]] = True
|