pytme 0.2.9.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl

tme/filters/whitening.py

@@ -1,8 +1,9 @@
-""" Implements class BandPassFilter to create Fourier filter representations.
+"""
+Implements class BandPassFilter to create Fourier filter representations.
 
-    Copyright (c) 2024 European Molecular Biology Laboratory
+Copyright (c) 2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from typing import Tuple, Dict
@@ -14,7 +15,7 @@ from scipy.ndimage import map_coordinates
 from ..types import BackendArray
 from ..backends import backend as be
 from .compose import ComposableFilter
-from ._utils import fftfreqn, compute_fourier_shape
+from ._utils import fftfreqn, compute_fourier_shape, shift_fourier
 
 __all__ = ["LinearWhiteningFilter"]
 
@@ -23,11 +24,6 @@ class LinearWhiteningFilter(ComposableFilter):
     """
     Compute Fourier power spectrums and perform whitening.
 
-    Parameters
-    ----------
-    **kwargs : Dict, optional
-        Additional keyword arguments.
-
     References
     ----------
     .. [1] de Teresa-Trueba, I.; Goetz, S. K.; Mattausch, A.; Stojanovska, F.; Zimmerli, C. E.;
@@ -38,7 +34,7 @@ class LinearWhiteningFilter(ComposableFilter):
         13375 (2023)
     """
 
-    def __init__(self, **kwargs):
+    def __init__(self, *args, **kwargs):
         pass
 
     @staticmethod
@@ -103,13 +99,6 @@ class LinearWhiteningFilter(ComposableFilter):
         shape_is_real_fourier: bool = True,
         order: int = 1,
     ) -> BackendArray:
-        """
-        References
-        ----------
-        .. [1]  M. L. Chaillet, G. van der Schot, I. Gubins, S. Roet,
-            R. C. Veltkamp, and F. Förster, Int. J. Mol. Sci. 24,
-            13375 (2023)
-        """
         grid = fftfreqn(
             shape=shape,
             sampling_rate=0.5,
@@ -123,11 +112,13 @@ class LinearWhiteningFilter(ComposableFilter):
 
     def __call__(
         self,
+        shape: Tuple[int],
         data: BackendArray = None,
         data_rfft: BackendArray = None,
         n_bins: int = None,
         batch_dimension: int = None,
         order: int = 1,
+        return_real_fourier: bool = True,
         **kwargs: Dict,
     ) -> Dict:
         """
@@ -135,6 +126,8 @@ class LinearWhiteningFilter(ComposableFilter):
 
         Parameters
         ----------
+        shape : tuple of ints
+            Shape of the returned whitening filter.
         data : BackendArray, optional
             The input data, defaults to None.
         data_rfft : BackendArray, optional
@@ -143,49 +136,51 @@ class LinearWhiteningFilter(ComposableFilter):
             The number of bins for computing the spectrum, defaults to None.
         batch_dimension : int, optional
             Batch dimension to average over.
-        order : int, optional
-            Interpolation order to use.
+        return_real_fourier : tuple of int
+            Return a shape compliant with rfft, i.e., omit the negative frequencies
+            terms resulting in a return shape (*shape[:-1], shape[-1]//2+1)
         **kwargs : Dict
             Additional keyword arguments.
 
         Returns
         -------
-        Dict
-           Filter data and associated parameters.
+        dict
+            data: BackendArray
+                The filter mask.
+            shape: tuple of ints
+                The requested filter shape
+            return_real_fourier: bool
+                Whether data is compliant with rfftn.
+            is_multiplicative_filter: bool
+                Whether the filter is multiplicative in Fourier space.
         """
         if data_rfft is None:
-            data_rfft = np.fft.rfftn(be.to_numpy_array(data))
+            data_rfft = be.rfftn(data)
 
         data_rfft = be.to_numpy_array(data_rfft)
-
         bins, radial_averages = self._compute_spectrum(
             data_rfft, n_bins, batch_dimension
         )
+        shape = tuple(int(x) for i, x in enumerate(shape) if i != batch_dimension)
 
-        if order is None:
-            cutoff = bins < radial_averages.size
-            filter_mask = np.zeros(bins.shape, radial_averages.dtype)
-            filter_mask[cutoff] = radial_averages[bins[cutoff]]
-        else:
-            shape = bins.shape
-            if kwargs.get("shape", False):
-                shape = compute_fourier_shape(
-                    shape=kwargs.get("shape"),
-                    shape_is_real_fourier=kwargs.get("shape_is_real_fourier", False),
-                )
-
-            filter_mask = self._interpolate_spectrum(
-                spectrum=radial_averages,
+        shape_filter = shape
+        if return_real_fourier:
+            shape_filter = compute_fourier_shape(
                 shape=shape,
-                shape_is_real_fourier=True,
+                shape_is_real_fourier=False,
             )
 
-        filter_mask = np.fft.ifftshift(
-            filter_mask,
-            axes=tuple(i for i in range(data_rfft.ndim - 1) if i != batch_dimension),
+        ret = self._interpolate_spectrum(
+            spectrum=radial_averages,
+            shape=shape_filter,
+            shape_is_real_fourier=return_real_fourier,
         )
 
+        ret = shift_fourier(data=ret, shape_is_real_fourier=return_real_fourier)
+
         return {
-            "data": be.to_backend_array(filter_mask),
+            "data": be.to_backend_array(ret),
+            "shape": shape,
+            "return_real_fourier": return_real_fourier,
             "is_multiplicative_filter": True,
         }

tme/matching_data.py

@@ -1,8 +1,9 @@
-""" Class representation of template matching data.
+"""
+Class representation of template matching data.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import warnings
@@ -174,7 +175,7 @@ class MatchingData:
         target_pad: NDArray = None,
         template_pad: NDArray = None,
         invert_target: bool = False,
-    ) -> "MatchingData":
+    ) -> Tuple["MatchingData", Tuple]:
         """
         Subset class instance based on slices.
 
@@ -193,6 +194,8 @@ class MatchingData:
         -------
         :py:class:`MatchingData`
             Newly allocated subset of class instance.
+        Tuple
+            Translation offset to merge analyzers.
 
         Examples
         --------
@@ -250,8 +253,9 @@ class MatchingData:
         target_offset[mask] = [x.start for x in target_slice]
         mask = np.subtract(1, self._target_batch).astype(bool)
         template_offset = np.zeros(len(self._output_template_shape), dtype=int)
-        template_offset[mask] = [x.start for x in template_slice]
-        ret._translation_offset = tuple(x for x in target_offset)
+        template_offset[mask] = [x.start for x, b in zip(template_slice, mask) if b]
+
+        translation_offset = tuple(x for x in target_offset)
 
         ret.target_filter = self.target_filter
         ret.template_filter = self.template_filter
@@ -261,7 +265,7 @@ class MatchingData:
             template_dim=getattr(self, "_template_dim", None),
         )
 
-        return ret
+        return ret, translation_offset
 
     def to_backend(self):
         """
@@ -322,11 +326,6 @@ class MatchingData:
 
         target_ndim -= len(target_dims)
         template_ndim -= len(template_dims)
-
-        if target_ndim != template_ndim:
-            raise ValueError(
-                f"Dimension mismatch: Target ({target_ndim}) Template ({template_ndim})."
-            )
         self._set_matching_dimension(
             target_dims=target_dims, template_dims=template_dims
         )
@@ -491,29 +490,26 @@ class MatchingData:
     def _fourier_padding(
         target_shape: Tuple[int],
         template_shape: Tuple[int],
-        pad_fourier: bool,
+        pad_target: bool = False,
         batch_mask: Tuple[int] = None,
     ) -> Tuple[Tuple, Tuple, Tuple, Tuple]:
-        fourier_pad = template_shape
-        fourier_shift = np.zeros_like(template_shape)
-
         if batch_mask is None:
             batch_mask = np.zeros_like(template_shape)
         batch_mask = np.asarray(batch_mask)
 
-        if not pad_fourier:
-            fourier_pad = np.ones(len(fourier_pad), dtype=int)
+        fourier_pad = np.ones(len(template_shape), dtype=int)
         fourier_pad = np.multiply(fourier_pad, 1 - batch_mask)
         fourier_pad = np.add(fourier_pad, batch_mask)
 
+        # Avoid padding batch dimensions
         pad_shape = np.maximum(target_shape, template_shape)
+        pad_shape = np.maximum(target_shape, np.multiply(1 - batch_mask, pad_shape))
         ret = be.compute_convolution_shapes(pad_shape, fourier_pad)
         conv_shape, fast_shape, fast_ft_shape = ret
 
         template_mod = np.mod(template_shape, 2)
-        if not pad_fourier:
-            fourier_shift = 1 - np.divide(template_shape, 2).astype(int)
-            fourier_shift = np.subtract(fourier_shift, template_mod)
+        fourier_shift = 1 - np.divide(template_shape, 2).astype(int)
+        fourier_shift = np.subtract(fourier_shift, template_mod)
 
         shape_diff = np.multiply(
             np.subtract(target_shape, template_shape), 1 - batch_mask
@@ -522,34 +518,27 @@ class MatchingData:
         if np.sum(shape_mask):
             shape_shift = np.divide(shape_diff, 2)
             offset = np.mod(shape_diff, 2)
-            if pad_fourier:
-                offset = -np.subtract(
-                    offset,
-                    np.logical_and(np.mod(target_shape, 2) == 0, template_mod == 1),
-                )
-            else:
-                warnings.warn(
-                    "Template is larger than target and padding is turned off. Consider "
-                    "swapping them or activate padding. Correcting the shift for now."
-                )
+            warnings.warn(
+                "Template is larger than target and padding is turned off. Consider "
+                "swapping them or activate padding. Correcting the shift for now."
+            )
             shape_shift = np.multiply(np.add(shape_shift, offset), shape_mask)
             fourier_shift = np.subtract(fourier_shift, shape_shift).astype(int)
 
-        fourier_shift = tuple(np.multiply(fourier_shift, 1 - batch_mask).astype(int))
+        if pad_target:
+            fourier_shift = np.subtract(fourier_shift, np.subtract(1, template_mod))
 
+        fourier_shift = tuple(np.multiply(fourier_shift, 1 - batch_mask).astype(int))
         return tuple(conv_shape), tuple(fast_shape), tuple(fast_ft_shape), fourier_shift
 
-    def fourier_padding(
-        self, pad_fourier: bool = False
-    ) -> Tuple[Tuple, Tuple, Tuple, Tuple]:
+    def fourier_padding(self, pad_target: bool = False) -> Tuple:
         """
         Computes efficient shape four Fourier transforms and potential associated shifts.
 
         Parameters
         ----------
-        pad_fourier : bool, optional
-            If true, returns the shape of the full-convolution defined as sum of target
-            shape and template shape minus one, False by default.
+        pad_target : bool, optional
+            Whether the target has been padded to the full convolution shape.
 
         Returns
         -------
@@ -564,7 +553,7 @@ class MatchingData:
             target_shape=be.to_numpy_array(self._output_target_shape),
             template_shape=be.to_numpy_array(self._output_template_shape),
             batch_mask=be.to_numpy_array(self._batch_mask),
-            pad_fourier=pad_fourier,
+            pad_target=pad_target,
         )
 
     def computation_schedule(

tme/matching_exhaustive.py

@@ -1,8 +1,9 @@
-""" Implements cross-correlation based template matching using different metrics.
+"""
+Implements cross-correlation based template matching using different metrics.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import sys
@@ -19,6 +20,7 @@ from .filters import Compose
 from .backends import backend as be
 from .matching_utils import split_shape
 from .types import CallbackClass, MatchingData
+from .analyzer.proxy import SharedAnalyzerProxy
 from .matching_scores import MATCHING_EXHAUSTIVE_REGISTER
 from .memory import MatchingMemoryUsage, MATCHING_MEMORY_REGISTRY
 
@@ -147,7 +149,7 @@ def scan(
     n_jobs: int = 4,
     callback_class: CallbackClass = None,
     callback_class_args: Dict = {},
-    pad_fourier: bool = True,
+    pad_target: bool = True,
     pad_template_filter: bool = True,
     interpolation_order: int = 3,
     jobs_per_callback_class: int = 8,
@@ -174,8 +176,8 @@ def scan(
         Analyzer class pointer to operate on computed scores.
     callback_class_args : dict, optional
         Arguments passed to the callback_class. Default is an empty dictionary.
-    pad_fourier: bool, optional
-        Whether to pad target and template to the full convolution shape.
+    pad_target: bool, optional
+        Whether to pad target to the full convolution shape.
     pad_template_filter: bool, optional
         Whether to pad potential template filters to the full convolution shape.
     interpolation_order : int, optional
@@ -208,17 +210,17 @@ def scan(
     >>> )
 
     """
-    matching_data = matching_data.subset_by_slice(
+    matching_data, translation_offset = matching_data.subset_by_slice(
         target_slice=target_slice,
         template_slice=template_slice,
-        target_pad=matching_data.target_padding(pad_target=pad_fourier),
+        target_pad=matching_data.target_padding(pad_target=pad_target),
     )
 
     matching_data.to_backend()
     template_shape = matching_data._batch_shape(
         matching_data.template.shape, matching_data._template_batch
     )
-    conv, fwd, inv, shift = matching_data.fourier_padding(pad_fourier=False)
+    conv, fwd, inv, shift = matching_data.fourier_padding(pad_target=pad_target)
 
     template_filter = _setup_template_filter_apply_target_filter(
         matching_data=matching_data,
@@ -229,19 +231,20 @@ def scan(
 
     default_callback_args = {
         "shape": fwd,
-        "offset": matching_data._translation_offset,
+        "offset": translation_offset,
         "fourier_shift": shift,
         "fast_shape": fwd,
         "targetshape": matching_data._output_shape,
         "templateshape": template_shape,
         "convolution_shape": conv,
         "thread_safe": n_jobs > 1,
-        "convolution_mode": "valid" if pad_fourier else "same",
+        "convolution_mode": "valid" if pad_target else "same",
         "shm_handler": shm_handler,
         "only_unique_rotations": True,
         "aggregate_axis": matching_data._batch_axis(matching_data._batch_mask),
         "n_rotations": matching_data.rotations.shape[0],
     }
+    callback_class_args["inversion_mapping"] = n_jobs == 1
     default_callback_args.update(callback_class_args)
 
     setup = matching_setup(
@@ -257,13 +260,17 @@ def scan(
     matching_data._free_data()
     be.free_cache()
 
-    # Some analyzers cannot be shared across processes
-    if not getattr(callback_class, "is_shareable", False):
-        jobs_per_callback_class = 1
-
     n_callback_classes = max(n_jobs // jobs_per_callback_class, 1)
     callback_classes = [
-        callback_class(**default_callback_args) if callback_class else None
+        (
+            SharedAnalyzerProxy(
+                callback_class,
+                default_callback_args,
+                shm_handler=shm_handler if n_jobs > 1 else None,
+            )
+            if callback_class
+            else None
+        )
         for _ in range(n_callback_classes)
     ]
     ret = Parallel(n_jobs=n_jobs)(
@@ -276,14 +283,9 @@ def scan(
         )
         for index, rotation in enumerate(matching_data._split_rotations_on_jobs(n_jobs))
     )
-
-    # TODO: Make sure peak callers are thread safe to begin with
-    if not getattr(callback_class, "is_shareable", False):
-        callback_classes = ret
-
     callbacks = [
-        tuple(callback._postprocess(**default_callback_args))
-        for callback in callback_classes
+        callback.result(**default_callback_args)
+        for callback in ret[:n_callback_classes]
         if callback
     ]
     be.free_cache()
@@ -423,14 +425,17 @@ def scan_subsets(
     splits = tuple(product(target_splits, template_splits))
 
     outer_jobs, inner_jobs = job_schedule
-    if hasattr(be, "scan"):
+    if be._backend_name == "jax":
+        func = be.scan
+
         corr_scoring = MATCHING_EXHAUSTIVE_REGISTER.get("CORR", (None, None))[1]
-        results = be.scan(
+        results = func(
             matching_data=matching_data,
             splits=splits,
             n_jobs=outer_jobs,
             rotate_mask=matching_score != corr_scoring,
             callback_class=callback_class,
+            callback_class_args=callback_class_args,
         )
     else:
         results = Parallel(n_jobs=outer_jobs, verbose=verbose)(
@@ -445,7 +450,7 @@ def scan_subsets(
                     callback_class=callback_class,
                     callback_class_args=callback_class_args,
                     interpolation_order=interpolation_order,
-                    pad_fourier=pad_target_edges,
+                    pad_target=pad_target_edges,
                     gpu_index=index % outer_jobs,
                     pad_template_filter=pad_template_filter,
                     target_slice=target_split,
@@ -454,7 +459,6 @@ def scan_subsets(
                 for index, (target_split, template_split) in enumerate(splits)
             ]
         )
-
     matching_data._free_data()
     if callback_class is not None:
         return callback_class.merge(results, **callback_class_args)

tme/matching_optimization.py

@@ -1,8 +1,9 @@
-""" Implements methods for non-exhaustive template matching.
+"""
+Implements methods for non-exhaustive template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import warnings
@@ -1308,4 +1309,4 @@ def optimize_match(
         result.x = np.zeros_like(result.x)
     translation, rotation = result.x[:ndim], result.x[ndim:]
     rotation_matrix = euler_to_rotationmatrix(rotation)
-    return translation, rotation_matrix, result.fun
+    return translation, rotation_matrix, float(result.fun)

tme/matching_scores.py

@@ -1,8 +1,9 @@
-""" Implements a range of cross-correlation coefficients.
+"""
+Implements a range of cross-correlation coefficients.
 
-    Copyright (c) 2023-2024 European Molecular Biology Laboratory
+Copyright (c) 2023-2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import warnings
@@ -592,20 +593,27 @@ def corr_scoring(
         **_fftargs,
     )
 
+    center = be.divide(be.to_backend_array(template.shape) - 1, 2)
     unpadded_slice = tuple(slice(0, stop) for stop in template.shape)
+
+    template_rot = be.zeros(template.shape, be._float_dtype)
     for index in range(rotations.shape[0]):
+        # d+1, d+1 rigid transform matrix from d,d rotation matrix
         rotation = rotations[index]
-        arr = be.fill(arr, 0)
-        arr, _ = be.rigid_transform(
+        matrix = be._rigid_transform_matrix(rotation_matrix=rotation, center=center)
+        template_rot, _ = be.rigid_transform(
             arr=template,
-            rotation_matrix=rotation,
-            out=arr,
-            use_geometric_center=True,
+            rotation_matrix=matrix,
+            out=template_rot,
             order=interpolation_order,
-            cache=False,
+            cache=True,
         )
-        arr = template_filter_func(arr, ft_temp, template_filter)
-        norm_template(arr[unpadded_slice], template_mask, mask_sum)
+
+        template_rot = template_filter_func(template_rot, ft_temp, template_filter)
+        norm_template(template_rot, template_mask, mask_sum)
+
+        arr = be.fill(arr, 0)
+        arr[unpadded_slice] = template_rot
 
         ft_temp = rfftn(arr, ft_temp)
         ft_temp = be.multiply(ft_target, ft_temp, out=ft_temp)
@@ -728,7 +736,7 @@ def flc_scoring(
             out_mask=temp,
             use_geometric_center=True,
             order=interpolation_order,
-            cache=False,
+            cache=True,
         )
 
         n_obs = be.sum(temp)
@@ -874,7 +882,7 @@ def mcc_scoring(
             out_mask=temp,
             use_geometric_center=True,
             order=interpolation_order,
-            cache=False,
+            cache=True,
         )
 
         template_filter_func(template_rot, temp_ft, template_filter)
@@ -1034,7 +1042,8 @@ def flc_scoring2(
             out_mask=tmp_sqz,
             use_geometric_center=True,
             order=interpolation_order,
-            cache=False,
+            cache=True,
+            batched=True,
         )
         n_obs = be.sum(tmp_sqz, axis=data_axes, keepdims=True)
         arr_norm = template_filter_func(arr_sqz, ft_temp, template_filter)
@@ -1154,7 +1163,8 @@ def corr_scoring2(
             out=arr_sqz,
             use_geometric_center=True,
             order=interpolation_order,
-            cache=False,
+            cache=True,
+            batched=True,
         )
         arr_norm = template_filter_func(arr_sqz, ft_sqz, template_filter)
         norm_template(arr_norm[unpadded_slice], template_mask, mask_sum, axis=data_axes)

tme/matching_utils.py

@@ -1,8 +1,9 @@
-""" Utility functions for template matching.
+"""
+Utility functions for template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import os
@@ -519,7 +520,7 @@ def apply_convolution_mode(
     elif convolution_mode == "same":
         return func(arr, s1)
     elif convolution_mode == "valid":
-        valid_shape = [s1[i] - s2[i] + s2[i] % 2 for i in range(arr.ndim)]
+        valid_shape = [s1[i] - s2[i] + 1 for i in range(arr.ndim)]
         return func(arr, valid_shape)
 
 
@@ -1056,7 +1057,7 @@ def tube_mask(
 
 def scramble_phases(
     arr: NDArray,
-    noise_proportion: float = 0.5,
+    noise_proportion: float = 1.0,
     seed: int = 42,
     normalize_power: bool = False,
 ) -> NDArray:
@@ -1068,7 +1069,7 @@ def scramble_phases(
     arr : NDArray
         Input data.
     noise_proportion : float, optional
-        Proportion of scrambled phases, 0.5 by default.
+        Proportion of scrambled phases, 1.0 by default.
     seed : int, optional
         The seed for the random phase scrambling, 42 by default.
     normalize_power : bool, optional
@@ -1079,15 +1080,22 @@ def scramble_phases(
     NDArray
         Phase scrambled version of ``arr``.
     """
+    from tme.filters._utils import fftfreqn
+
     np.random.seed(seed)
     noise_proportion = max(min(noise_proportion, 1), 0)
 
     arr_fft = np.fft.fftn(arr)
     amp, ph = np.abs(arr_fft), np.angle(arr_fft)
 
-    ph_noise = np.random.permutation(ph)
-    ph_new = ph * (1 - noise_proportion) + ph_noise * noise_proportion
-    ret = np.real(np.fft.ifftn(amp * np.exp(1j * ph_new)))
+    # Scrambling up to nyquist gives more uniform noise distribution
+    mask = np.fft.ifftshift(
+        fftfreqn(arr_fft.shape, sampling_rate=1, compute_euclidean_norm=True) <= 0.5
+    )
+
+    ph_noise = np.random.permutation(ph[mask])
+    ph[mask] = ph[mask] * (1 - noise_proportion) + ph_noise * noise_proportion
+    ret = np.real(np.fft.ifftn(amp * np.exp(1j * ph)))
 
     if normalize_power:
         np.divide(ret - ret.min(), ret.max() - ret.min(), out=ret)
@@ -1150,3 +1158,40 @@ def compute_extraction_box(
     keep = be.multiply(keep, clamp_change == 0)
 
     return obs_beg_clamp, obs_end_clamp, cand_beg, cand_end, keep
+
+
+class TqdmParallel(Parallel):
+    """
+    A minimal Parallel implementation using tqdm for progress reporting.
+
+    Parameters:
+    -----------
+    tqdm_args : dict, optional
+        Dictionary of arguments passed to tqdm.tqdm
+    *args, **kwargs:
+        Arguments to pass to joblib.Parallel
+    """
+
+    def __init__(self, tqdm_args: Dict = {}, *args, **kwargs):
+        from tqdm import tqdm
+
+        super().__init__(*args, **kwargs)
+        self.pbar = tqdm(**tqdm_args)
+
+    def __call__(self, iterable, *args, **kwargs):
+        self.n_tasks = len(iterable) if hasattr(iterable, "__len__") else None
+        return super().__call__(iterable, *args, **kwargs)
+
+    def print_progress(self):
+        if self.n_tasks is None:
+            return super().print_progress()
+
+        if self.n_tasks != self.pbar.total:
+            self.pbar.total = self.n_tasks
+            self.pbar.refresh()
+
+        self.pbar.n = self.n_completed_tasks
+        self.pbar.refresh()
+
+        if self.n_completed_tasks >= self.n_tasks:
+            self.pbar.close()