def _node_io_loop(self, func, *args, **kwargs):
"""
Call the provided function over a list of nodes.
If possible, group nodes by common data files to speed
things up. This should work like __iter__, except we call
a function instead of yielding.
Parameters
----------
func : function
Function to be called on an array of nodes.
pbar : optional, string or yt.funcs.TqdmProgressBar
A progress bar to be updated with each iteration.
If a string, a progress bar will be created and the
finish function will be called. If a progress bar is
provided, the finish function will not be called.
Default: None (no progress bar).
root_nodes : optional, array of root TreeNodes
Array of nodes over which the function will be called.
If None, the list will be self.trees (i.e., all
root_nodes).
Default: None.
store : optional, string
If not None, any return value captured from the function
will be stored in an attribute with this name associated
with the TreeNode.
Default: None.
"""
pbar = kwargs.pop("pbar", None)
root_nodes = kwargs.pop("root_nodes", None)
if root_nodes is None:
root_nodes = self.trees
store = kwargs.pop("store", None)
data_files, node_list = self._node_io_loop_prepare(root_nodes)
nnodes = sum([nodes.size for nodes in node_list])
finish = True
if pbar is None:
pbar = fake_pbar("", nnodes)
elif not isinstance(pbar, TqdmProgressBar):
pbar = get_pbar(pbar, nnodes)
else:
finish = False
for data_file, nodes in zip(data_files, node_list):
self._node_io_loop_start(data_file)
for node in nodes:
rval = func(node, *args, **kwargs)
if store is not None:
setattr(node, store, rval)
pbar.update(1)
self._node_io_loop_finish(data_file)
if finish:
pbar.finish()
def f_text_block(f,
block_size=4096,
file_size=None,
sep="\n",
pbar_string=None):
"""
Read lines from a file faster than f.readlines().
"""
start = f.tell()
if file_size is None:
f.seek(0, 2)
file_size = f.tell() - start
f.seek(start)
nblocks = np.ceil(float(file_size) / block_size).astype(np.int64)
read_size = file_size + start
lbuff = ""
if pbar_string is None:
pbar = fake_pbar()
else:
pbar = get_pbar(pbar_string, file_size)
for ib in range(nblocks):
offset = f.tell()
my_block = min(block_size, read_size - offset)
if my_block <= 0: break
buff = f.read(my_block)
linl = -1
for ih in range(buff.count(sep)):
inl = buff.find(sep, linl + 1)
if inl < 0:
lbuff += buff[linl + 1:]
continue
else:
line = lbuff + buff[linl + 1:inl]
loc = offset - len(lbuff) + linl + 1
lbuff = ""
linl = inl
pbar.update(loc + len(line) - start + 1)
yield line, loc
lbuff += buff[linl + 1:]
if lbuff:
loc = f.tell() - len(lbuff)
pbar.update(loc + len(lbuff) - start + 1)
yield lbuff, loc
pbar.finish()
def _node_io_loop(self, func, *args, **kwargs):
"""
Call the provided function over a list of nodes.
If possible, group nodes by common data files to speed
things up.
Parameters
----------
func : function
Function to be called on an array of nodes.
pbar : optional, string or yt.funcs.TqdmProgressBar
A progress bar to be updated with each iteration.
If a string, a progress bar will be created and the
finish function will be called. If a progress bar is
provided, the finish function will not be called.
Default: None (no progress bar).
root_nodes : optional, array of root TreeNodes
Array of nodes over which the function will be called.
If None, the list will be self[:] (i.e., all
root_nodes).
Default: None.
Returns
-------
rvals : list
return values from calling func on each node.
These will have the same order as the original node list.
"""
self._plant_trees()
pbar = kwargs.pop("pbar", None)
root_nodes = kwargs.pop("root_nodes", None)
data_files, node_list, return_order = \
self._node_io_loop_prepare(root_nodes)
nnodes = sum([nodes.size for nodes in node_list])
finish = True
if pbar is None:
pbar = fake_pbar("", nnodes)
elif not isinstance(pbar, TqdmProgressBar):
pbar = get_pbar(pbar, nnodes)
else:
finish = False
rvals = []
c = 0
for data_file, nodes in zip(data_files, node_list):
self._node_io_loop_start(data_file)
# if we're doing all of them, just give the indices
if root_nodes is None:
my_nodes = nodes
else:
my_nodes = root_nodes[nodes]
for node in self._yield_root_nodes(my_nodes):
rval = func(node, *args, **kwargs)
rvals.append(rval)
c += 1
pbar.update(c)
self._node_io_loop_finish(data_file)
if finish:
pbar.finish()
if return_order is not None:
rvals = [rvals[i] for i in return_order]
return rvals