def run_udf(self,
dataset: DataSet,
udf: UDF,
roi: np.ndarray = None,
corrections: CorrectionSet = None,
progress: bool = False,
backends=None) -> Dict[str, BufferWrapper]:
"""
Run :code:`udf` on :code:`dataset`, restricted to the region of interest :code:`roi`.
.. versionchanged:: 0.5.0
Added the :code:`progress` parameter
.. versionchanged:: 0.6.0
Added the :code:`corrections` and :code:`backends` parameter
Parameters
----------
dataset
The dataset to work on
udf
UDF instance you want to run
roi : numpy.ndarray
Region of interest as bool mask over the navigation axes of the dataset
progress : bool
Show progress bar
corrections
Corrections to apply while running the UDF. If none are given,
the corrections that are part of the :code:`DataSet` are used,
if there are any.
backends : None or iterable containing 'numpy', 'cupy' and/or 'cuda'
Restrict the back-end to a subset of the capabilities of the UDF.
This can be useful for testing hybrid UDFs.
Returns
-------
dict
Return value of the UDF containing the result buffers of
type :class:`libertem.common.buffers.BufferWrapper`. Note that a
:class:`~libertem.common.buffers.BufferWrapper` can be used like
a :class:`numpy.ndarray` in many cases because it implements
:meth:`__array__`. You can access the underlying numpy array using the
:attr:`~libertem.common.buffers.BufferWrapper.data` property.
"""
if corrections is None:
corrections = dataset.get_correction_data()
results = UDFRunner([udf]).run_for_dataset(
dataset=dataset,
executor=self.executor,
roi=roi,
progress=progress,
corrections=corrections,
backends=backends,
)
return results[0]
def _make_udf_tasks(self, dataset: DataSet, roi):
for idx, partition in enumerate(dataset.get_partitions()):
if roi is not None:
roi_for_part = self._roi_for_partition(roi, partition)
if np.count_nonzero(roi_for_part) == 0:
# roi is empty for this partition, ignore
continue
udf = self._udf.copy_for_partition(partition, roi)
yield UDFTask(partition=partition, idx=idx, udf=udf, roi=roi)
def _make_udf_tasks(self, dataset: DataSet, roi, corrections, backends):
for idx, partition in enumerate(dataset.get_partitions()):
if roi is not None:
roi_for_part = self._roi_for_partition(roi, partition)
if np.count_nonzero(roi_for_part) == 0:
# roi is empty for this partition, ignore
continue
udfs = [
udf.copy_for_partition(partition, roi)
for udf in self._udfs
]
yield UDFTask(
partition=partition, idx=idx, udfs=udfs, roi=roi, corrections=corrections,
backends=backends,
)
def _run_sync(
self,
dataset: DataSet,
udf: Union[UDF, Iterable[UDF]],
roi: Optional[np.ndarray],
corrections: Optional[CorrectionSet],
progress: bool,
backends,
plots,
iterate: bool,
):
"""
Run the given UDF(s), either returning the final result (when
:code:`iterate=False` is given), or a generator that yields partial results.
"""
enable_plotting = bool(plots)
udf_is_list = isinstance(udf, Iterable)
if not isinstance(udf, Iterable):
udfs = [udf]
else:
udfs = list(udf)
if enable_plotting:
plots = self._prepare_plots(udfs, dataset, roi, plots)
if corrections is None:
corrections = dataset.get_correction_data()
if (roi is not None) and (roi.dtype is not np.dtype(bool)):
warnings.warn(
f"ROI dtype is {roi.dtype}, expected bool. Attempting cast to bool."
)
roi = roi.astype(bool)
def _run_sync_wrap() -> Generator[UDFResults, None, None]:
runner_cls = self.executor.get_udf_runner()
result_iter = runner_cls(udfs).run_for_dataset_sync(
dataset=dataset,
executor=self.executor,
roi=roi,
progress=progress,
corrections=corrections,
backends=backends,
iterate=(iterate or enable_plotting))
for udf_results in result_iter:
yield udf_results
if enable_plotting:
self._update_plots(plots,
udfs,
udf_results.buffers,
udf_results.damage.data,
force=False)
if enable_plotting:
self._update_plots(plots,
udfs,
udf_results.buffers,
udf_results.damage.data,
force=True)
if iterate:
return _run_sync_wrap()
else:
udf_results = run_gen_get_last(_run_sync_wrap())
if udf_is_list:
return udf_results.buffers
else:
return udf_results.buffers[0]
def _run_sync(
self,
dataset: DataSet,
udf: UDF,
roi: np.ndarray = None,
corrections: CorrectionSet = None,
progress: bool = False,
backends=None,
plots=None,
iterate=False,
):
"""
Run the given UDF(s), either returning the final result (when
:code:`iterate=False` is given), or a generator that yields partial results.
"""
enable_plotting = bool(plots)
udf_is_list = isinstance(udf, (tuple, list))
if not udf_is_list:
udfs = [udf]
else:
udfs = list(udf)
if enable_plotting:
plots = self._prepare_plots(udfs, dataset, roi, plots)
if corrections is None:
corrections = dataset.get_correction_data()
if (roi is not None) and (roi.dtype is not np.dtype(bool)):
warnings.warn(
f"ROI dtype is {roi.dtype}, expected bool. Attempting cast to bool."
)
roi = roi.astype(bool)
def _run_sync_wrap():
result_iter = UDFRunner(udfs).run_for_dataset_sync(
dataset=dataset,
executor=self.executor,
roi=roi,
progress=progress,
corrections=corrections,
backends=backends,
)
for udf_results in result_iter:
yield udf_results
if enable_plotting:
self._update_plots(plots,
udfs,
udf_results.buffers,
udf_results.damage,
force=False)
if enable_plotting:
self._update_plots(plots,
udfs,
udf_results.buffers,
udf_results.damage,
force=True)
if iterate:
return _run_sync_wrap()
else:
for udf_results in _run_sync_wrap():
pass
if udf_is_list:
return udf_results.buffers
else:
return udf_results.buffers[0]