rapidtide 3.0.10__py3-none-any.whl → 3.1__py3-none-any.whl

rapidtide/multiproc.py

@@ -20,7 +20,10 @@ import multiprocessing as mp
 import sys
 import threading as thread
 from platform import python_version, system
+from typing import Any, Callable, List, Optional, Tuple
 
+import numpy as np
+from numpy.typing import NDArray
 from tqdm import tqdm
 
 try:
@@ -29,14 +32,86 @@ except ImportError:
     import Queue as thrQueue
 
 
-def maxcpus(reservecpu=True):
+def maxcpus(reservecpu: bool = True) -> int:
+    """Return the maximum number of CPUs that can be used for parallel processing.
+
+    This function returns the total number of CPU cores available on the system,
+    with an option to reserve one CPU core for system operations.
+
+    Parameters
+    ----------
+    reservecpu : bool, default=True
+        If True, reserves one CPU core for system operations by returning
+        `cpu_count() - 1`. If False, returns the total number of CPU cores
+        available without reservation.
+
+    Returns
+    -------
+    int
+        The maximum number of CPUs available for parallel processing.
+        If `reservecpu=True`, returns `cpu_count() - 1`.
+        If `reservecpu=False`, returns `cpu_count()`.
+
+    Notes
+    -----
+    This function uses `multiprocessing.cpu_count()` to determine the number
+    of available CPU cores. The reserved CPU core helps maintain system
+    responsiveness during parallel processing tasks.
+
+    Examples
+    --------
+    >>> maxcpus()
+    7
+    >>> maxcpus(reservecpu=False)
+    8
+    """
     if reservecpu:
         return mp.cpu_count() - 1
     else:
         return mp.cpu_count()
 
 
-def _process_data(data_in, inQ, outQ, showprogressbar=True, chunksize=10000):
+def _process_data(
+    data_in: List[Any], inQ: Any, outQ: Any, showprogressbar: bool = True, chunksize: int = 10000
+) -> List[Any]:
+    """Process input data in chunks using multiprocessing queues.
+
+    This function distributes data into chunks and processes them using
+    provided input and output queues. It supports progress tracking and
+    handles both complete chunks and a final remainder chunk.
+
+    Parameters
+    ----------
+    data_in : List[Any]
+        Input data to be processed.
+    inQ : Any
+        Input queue for sending data to worker processes.
+    outQ : Any
+        Output queue for receiving processed data from worker processes.
+    showprogressbar : bool, optional
+        If True, display a progress bar during processing. Default is True.
+    chunksize : int, optional
+        Size of data chunks to process at a time. Default is 10000.
+
+    Returns
+    -------
+    List[Any]
+        List of processed data items retrieved from the output queue.
+
+    Notes
+    -----
+    This function assumes that `inQ` and `outQ` are properly configured
+    multiprocessing queues and that worker processes are running and
+    consuming from `inQ` and producing to `outQ`.
+
+    Examples
+    --------
+    >>> from multiprocessing import Queue
+    >>> data = list(range(1000))
+    >>> in_q = Queue()
+    >>> out_q = Queue()
+    >>> result = _process_data(data, in_q, out_q)
+    """
     # send pos/data to workers
     data_out = []
     totalnum = len(data_in)
@@ -84,16 +159,72 @@ def _process_data(data_in, inQ, outQ, showprogressbar=True, chunksize=10000):
 
 
 def run_multiproc(
-    consumerfunc,
-    inputshape,
-    maskarray,
-    nprocs=1,
-    verbose=True,
-    indexaxis=0,
-    procunit="voxels",
-    showprogressbar=True,
-    chunksize=1000,
-):
+    consumerfunc: Callable[[Any, Any], None],
+    inputshape: Tuple[int, ...],
+    maskarray: Optional[NDArray] = None,
+    nprocs: int = 1,
+    verbose: bool = True,
+    indexaxis: int = 0,
+    procunit: str = "voxels",
+    showprogressbar: bool = True,
+    chunksize: int = 1000,
+) -> List[Any]:
+    """
+    Execute a function in parallel across multiple processes using multiprocessing.
+
+    This function initializes a set of worker processes and distributes input data
+    across them for parallel processing. It supports optional masking of data
+    along a specified axis and provides progress reporting.
+
+    Parameters
+    ----------
+    consumerfunc : callable
+        Function to be executed in parallel. Must accept two arguments: an input queue
+        and an output queue for inter-process communication.
+    inputshape : tuple of int
+        Shape of the input data along all axes. The dimension along `indexaxis` is
+        used to determine the number of items to process.
+    maskarray : ndarray, optional
+        Boolean or binary mask array used to filter indices. Only indices where
+        `maskarray[d] > 0.5` are processed. If None, all indices are processed.
+    nprocs : int, optional
+        Number of worker processes to use. Default is 1 (single-threaded).
+    verbose : bool, optional
+        If True, print information about the number of units being processed.
+        Default is True.
+    indexaxis : int, optional
+        Axis along which to iterate for processing. Default is 0.
+    procunit : str, optional
+        Unit of processing, used for logging messages. Default is "voxels".
+    showprogressbar : bool, optional
+        If True, display a progress bar during processing. Default is True.
+    chunksize : int, optional
+        Number of items to process in each chunk. Default is 1000.
+
+    Returns
+    -------
+    list
+        List of results returned by the worker processes.
+
+    Notes
+    -----
+    - On Python 3.8+ and non-Windows systems, the function uses the 'fork' context
+      for better performance.
+    - The function will exit with an error if `maskarray` is provided but its
+      length does not match the size of the `indexaxis` dimension of `inputshape`.
+
+    Examples
+    --------
+    >>> def worker_func(inQ, outQ):
+    ...     while True:
+    ...         item = inQ.get()
+    ...         if item is None:
+    ...             break
+    ...         outQ.put(item * 2)
+    ...
+    >>> shape = (100, 100)
+    >>> result = run_multiproc(worker_func, shape, nprocs=4)
+    """
     # initialize the workers and the queues
     __spec__ = None
     n_workers = nprocs
@@ -148,16 +279,73 @@ def run_multiproc(
 
 
 def run_multithread(
-    consumerfunc,
-    inputshape,
-    maskarray,
-    verbose=True,
-    nprocs=1,
-    indexaxis=0,
-    procunit="voxels",
-    showprogressbar=True,
-    chunksize=1000,
-):
+    consumerfunc: Callable[[Any, Any], None],
+    inputshape: Tuple[int, ...],
+    maskarray: Optional[NDArray] = None,
+    verbose: bool = True,
+    nprocs: int = 1,
+    indexaxis: int = 0,
+    procunit: str = "voxels",
+    showprogressbar: bool = True,
+    chunksize: int = 1000,
+) -> List[Any]:
+    """
+    Execute a multithreaded processing task using a specified consumer function.
+
+    This function initializes a set of worker threads that process data in parallel
+    according to the provided consumer function. It supports optional masking,
+    progress tracking, and configurable chunking for efficient processing.
+
+    Parameters
+    ----------
+    consumerfunc : callable
+        A function that takes two arguments (input queue, output queue) and
+        processes data in a loop until signaled to stop.
+    inputshape : tuple of int
+        Shape of the input data along all axes. The dimension along `indexaxis`
+        determines how many items will be processed.
+    maskarray : ndarray, optional
+        Boolean or integer array used to filter which indices are processed.
+        Must match the size of the axis specified by `indexaxis`.
+    verbose : bool, optional
+        If True, print information about the number of items being processed
+        and the number of threads used. Default is True.
+    nprocs : int, optional
+        Number of worker threads to spawn. Default is 1.
+    indexaxis : int, optional
+        Axis along which the indexing is performed. Default is 0.
+    procunit : str, optional
+        Unit of processing, used in verbose output. Default is "voxels".
+    showprogressbar : bool, optional
+        If True, display a progress bar during processing. Default is True.
+    chunksize : int, optional
+        Number of items to process in each chunk. Default is 1000.
+
+    Returns
+    -------
+    list
+        A list of results returned by the consumer function for each processed item.
+
+    Notes
+    -----
+    - The function uses `threading.Queue` for inter-thread communication.
+    - If `maskarray` is provided, only indices where `maskarray[d] > 0` are processed.
+    - The `consumerfunc` is expected to read from `inQ` and write to `outQ` until
+      a `None` is received on `inQ`, signaling the end of processing.
+
+    Examples
+    --------
+    >>> def my_consumer(inQ, outQ):
+    ...     while True:
+    ...         item = inQ.get()
+    ...         if item is None:
+    ...             break
+    ...         result = item * 2
+    ...         outQ.put(result)
+    ...
+    >>> shape = (100, 50)
+    >>> result = run_multithread(my_consumer, shape, nprocs=4)
+    """
     # initialize the workers and the queues
     n_workers = nprocs
     inQ = thrQueue.Queue()

rapidtide/patchmatch.py

@@ -21,8 +21,10 @@ import math
 import os
 import sys
 import warnings
+from typing import Any
 
 import numpy as np
+from numpy.typing import NDArray
 from scipy.interpolate import griddata
 from scipy.ndimage import distance_transform_edt, gaussian_filter1d
 from skimage.filters import threshold_multiotsu
@@ -31,7 +33,9 @@ from skimage.segmentation import flood_fill
 import rapidtide.io as tide_io
 
 
-def interpolate_masked_voxels(data, mask, method="linear", extrapolate=True):
+def interpolate_masked_voxels(
+    data: NDArray, mask: NDArray, method: str = "linear", extrapolate: bool = True
+) -> NDArray:
     """
     Replaces masked voxels in a 3D numpy array with interpolated values
     from the unmasked region. Supports boundary extrapolation and multiple interpolation methods.
@@ -87,18 +91,41 @@ def interpolate_masked_voxels(data, mask, method="linear", extrapolate=True):
     return interpolated_data
 
 
-def get_bounding_box(mask, value, buffer=0):
+def get_bounding_box(mask: NDArray, value: int, buffer: int = 0) -> tuple[tuple, tuple]:
     """
     Computes the 3D bounding box that contains all the voxels in the mask with value value.
 
-    Parameters:
-        mask (np.ndarray): A 3D binary mask where non-zero values indicate the masked region.
-        value (int): The masked region value.
+    Parameters
+    ----------
+    mask : np.ndarray
+        A 3D binary mask where non-zero values indicate the masked region.
+    value : int
+        The masked region value to compute the bounding box for.
+    buffer : int, optional
+        Buffer to add around the bounding box in all directions. Default is 0.
 
-    Returns:
-        tuple: Two tuples defining the bounding box:
-               ((min_x, min_y, min_z), (max_x, max_y, max_z)),
-               where min and max are inclusive coordinates of the bounding box.
+    Returns
+    -------
+    tuple of tuple of int
+        Two tuples defining the bounding box:
+        ((min_x, min_y, min_z), (max_x, max_y, max_z)),
+        where min and max are inclusive coordinates of the bounding box.
+
+    Notes
+    -----
+    The function handles edge cases where the buffer extends beyond the mask boundaries
+    by clamping the coordinates to the valid range [0, shape[axis]-1].
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> mask = np.zeros((10, 10, 10), dtype=int)
+    >>> mask[3:7, 3:7, 3:7] = 1
+    >>> get_bounding_box(mask, 1)
+    ((3, 3, 3), (6, 6, 6))
+
+    >>> get_bounding_box(mask, 1, buffer=1)
+    ((2, 2, 2), (7, 7, 7))
     """
     if mask.ndim != 3:
         raise ValueError("Input mask must be a 3D array.")
@@ -119,21 +146,94 @@ def get_bounding_box(mask, value, buffer=0):
     return tuple(min_coords), tuple(max_coords)
 
 
-def flood3d(
-    image,
-    newvalue,
-):
+def flood3d(image: NDArray, newvalue: int) -> NDArray:
+    """
+    Apply flood fill to each slice of a 3D image.
+
+    This function performs a connected-component flood fill operation on each
+    2D slice of a 3D image, starting from the top-left corner (0, 0).
+
+    Parameters
+    ----------
+    image : NDArray
+        Input 3D image array of shape (height, width, depth)
+    newvalue : int
+        The value to fill the connected component with
+
+    Returns
+    -------
+    NDArray
+        3D image array of the same shape as input, with flood fill applied
+        to each slice
+
+    Notes
+    -----
+    - Uses 4-connectivity (rook-style connectivity) for flood fill
+    - Each slice is processed independently
+    - The fill operation starts from position (0, 0) in each slice
+    - Original image values are preserved in the output where fill did not occur
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> image = np.array([[[1, 1, 0],
+    ...                    [1, 0, 0],
+    ...                    [0, 0, 0]],
+    ...                   [[1, 1, 0],
+    ...                    [1, 0, 0],
+    ...                    [0, 0, 0]]])
+    >>> result = flood3d(image, 5)
+    >>> print(result)
+    """
     filledim = image * 0
     for slice in range(image.shape[2]):
         filledim[:, :, slice] = flood_fill(image[:, :, slice], (0, 0), newvalue, connectivity=1)
     return filledim
 
 
-def invertedflood3D(image, newvalue):
+def invertedflood3D(image: NDArray, newvalue: int) -> NDArray:
+    """
+    Apply inverted flood fill operation to a 3D image.
+
+    This function performs an inverted flood fill by adding the new value to the
+    original image and subtracting the result of a standard flood3d operation.
+
+    Parameters
+    ----------
+    image : NDArray
+        Input 3D image array to process
+    newvalue : int
+        Value to be added during the inverted flood fill operation
+
+    Returns
+    -------
+    NDArray
+        Resulting image after inverted flood fill operation
+
+    Notes
+    -----
+    The function relies on a `flood3d` function which is assumed to be defined
+    elsewhere in the codebase. The inverted flood fill is computed as:
+    result = image + newvalue - flood3d(image, newvalue)
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> image = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+    >>> result = invertedflood3D(image, 10)
+    >>> print(result)
+    """
     return image + newvalue - flood3d(image, newvalue)
 
 
-def growregion(image, location, value, separatedimage, regionsize, debug=False):
+def growregion(
+    image: NDArray,
+    location: tuple[int, int, int],
+    value: int,
+    separatedimage: NDArray,
+    regionsize: int,
+    debug: bool = False,
+) -> int:
     separatedimage[location[0], location[1], location[2]] = value
     regionsize += 1
     if debug:
@@ -157,7 +257,7 @@ def growregion(image, location, value, separatedimage, regionsize, debug=False):
     return regionsize
 
 
-def separateclusters(image, sizethresh=0, debug=False):
+def separateclusters(image: NDArray, sizethresh: int = 0, debug: bool = False) -> NDArray:
     separatedclusters = image * 0
     stop = False
     value = 1
@@ -219,25 +319,50 @@ def separateclusters(image, sizethresh=0, debug=False):
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-def clamp(low, high, value):
-    """bound an integer to a range
+def clamp(low: int, high: int, value: int) -> int:
+    """
+    Bound an integer to a range.
+
+    This function clamps a value to ensure it falls within the inclusive range [low, high].
+    If the value is less than low, it returns low. If the value is greater than high,
+    it returns high. Otherwise, it returns the value unchanged.
 
     Parameters
     ----------
     low : int
+        The lower bound of the range (inclusive).
     high : int
+        The upper bound of the range (inclusive).
     value : int
+        The value to be clamped.
 
     Returns
     -------
-    result : int
+    int
+        The clamped value within the range [low, high].
+
+    Notes
+    -----
+    The function assumes that `low <= high`. If this condition is not met,
+    the behavior is undefined and may return unexpected results.
+
+    Examples
+    --------
+    >>> clamp(0, 10, 5)
+    5
+    >>> clamp(0, 10, -1)
+    0
+    >>> clamp(0, 10, 15)
+    10
     """
     return max(low, min(high, value))
 
 
-def dehaze(fdata, level, debug=False):
-    """use Otsu to threshold https://scikit-image.org/docs/stable/auto_examples/segmentation/plot_multiotsu.html
-        n.b. threshold used to mask image: dark values are zeroed, but result is NOT binary
+def dehaze(fdata: NDArray, level: int, debug: bool = False) -> NDArray:
+    """
+    use Otsu to threshold https://scikit-image.org/docs/stable/auto_examples/segmentation/plot_multiotsu.html
+    n.b. threshold used to mask image: dark values are zeroed, but result is NOT binary
+
     Parameters
     ----------
     fdata : numpy.memmap from Niimg-like object
@@ -272,8 +397,15 @@ def dehaze(fdata, level, debug=False):
 
 
 # https://github.com/nilearn/nilearn/blob/1607b52458c28953a87bbe6f42448b7b4e30a72f/nilearn/image/image.py#L164
-def _smooth_array(arr, affine, fwhm=None, ensure_finite=True, copy=True):
-    """Smooth images by applying a Gaussian filter.
+def _smooth_array(
+    arr: NDArray,
+    affine: NDArray | None,
+    fwhm: float | NDArray | tuple | list | str | None = None,
+    ensure_finite: bool = True,
+    copy: bool = True,
+) -> NDArray:
+    """
+    Smooth images by applying a Gaussian filter.
 
     Apply a Gaussian filter along the three first dimensions of `arr`.
 
@@ -355,8 +487,10 @@ def _smooth_array(arr, affine, fwhm=None, ensure_finite=True, copy=True):
     return arr
 
 
-def binary_zero_crossing(fdata):
-    """binarize (negative voxels are zero)
+def binary_zero_crossing(fdata: NDArray) -> NDArray:
+    """
+    binarize (negative voxels are zero)
+
     Parameters
     ----------
     fdata : numpy.memmap from Niimg-like object
@@ -372,8 +506,11 @@ def binary_zero_crossing(fdata):
     return edge
 
 
-def difference_of_gaussian(fdata, affine, fwhmNarrow, ratioopt=True, debug=False):
-    """Apply Difference of Gaussian (DoG) filter.
+def difference_of_gaussian(
+    fdata: NDArray, affine: NDArray, fwhmNarrow: float, ratioopt: bool = True, debug: bool = False
+) -> NDArray:
+    """
+    Apply Difference of Gaussian (DoG) filter.
     https://en.wikipedia.org/wiki/Difference_of_Gaussians
     https://en.wikipedia.org/wiki/Marr–Hildreth_algorithm
     D. Marr and E. C. Hildreth. Theory of edge detection. Proceedings of the Royal Society, London B, 207:187-217, 1980
@@ -415,8 +552,16 @@ def difference_of_gaussian(fdata, affine, fwhmNarrow, ratioopt=True, debug=False
 # We are operating on data in memory that are closely associated with the source
 # NIFTI files, so the affine and sizes fields are easy to come by, but unlike the
 # original library, we are not working directly with NIFTI images.
-def calc_DoG(thedata, theaffine, thesizes, fwhm=3, ratioopt=True, debug=False):
-    """Find edges of a NIfTI image using the Difference of Gaussian (DoG).
+def calc_DoG(
+    thedata: NDArray,
+    theaffine: NDArray,
+    thesizes: tuple,
+    fwhm: float = 3,
+    ratioopt: bool = True,
+    debug: bool = False,
+) -> NDArray:
+    """
+    Find edges of a NIfTI image using the Difference of Gaussian (DoG).
     Parameters
     ----------
     thedata : 3D data array
@@ -440,7 +585,15 @@ def calc_DoG(thedata, theaffine, thesizes, fwhm=3, ratioopt=True, debug=False):
     return difference_of_gaussian(dehazed_data, theaffine, fwhm, ratioopt=ratioopt, debug=debug)
 
 
-def getclusters(theimage, theaffine, thesizes, fwhm=5, ratioopt=True, sizethresh=10, debug=False):
+def getclusters(
+    theimage: NDArray,
+    theaffine: NDArray,
+    thesizes: tuple,
+    fwhm: float = 5,
+    ratioopt: bool = True,
+    sizethresh: int = 10,
+    debug: bool = False,
+) -> NDArray:
     if debug:
         print("Detecting clusters..")
         print(f"\t{theimage.shape=}")
@@ -455,7 +608,54 @@ def getclusters(theimage, theaffine, thesizes, fwhm=5, ratioopt=True, sizethresh
     )
 
 
-def interppatch(img_data, separatedimage, method="linear", debug=False):
+def interppatch(
+    img_data: NDArray, separatedimage: NDArray, method: str = "linear", debug: bool = False
+) -> tuple[NDArray, NDArray]:
+    """
+    Interpolate voxel values within labeled regions of a 3D image.
+
+    This function applies interpolation to each labeled region in a separated image,
+    using the specified interpolation method. It returns both the interpolated image
+    and a copy of the original image with the same spatial extent.
+
+    Parameters
+    ----------
+    img_data : NDArray
+        A 3D array representing the input image data to be interpolated.
+    separatedimage : NDArray
+        A 3D array of integers where each unique positive integer represents a
+        distinct region. Zero values are treated as background.
+    method : str, optional
+        The interpolation method to use. Default is "linear". Other options may
+        include "nearest", "cubic", etc., depending on the implementation of
+        `interpolate_masked_voxels`.
+    debug : bool, optional
+        If True, print debug information for each region being processed.
+        Default is False.
+
+    Returns
+    -------
+    tuple[NDArray, NDArray]
+        A tuple containing:
+        - `interpolated`: The image with interpolated values in each region.
+        - `justboxes`: A copy of the original image data, with the same shape
+          as `img_data`, used for reference or visualization purposes.
+
+    Notes
+    -----
+    - Each region is processed independently using its bounding box.
+    - The function modifies `img_data` only within the bounds of each region.
+    - The `interpolate_masked_voxels` function is assumed to handle the actual
+      interpolation logic for masked voxels.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> img = np.random.rand(10, 10, 10)
+    >>> labels = np.zeros((10, 10, 10))
+    >>> labels[3:7, 3:7, 3:7] = 1
+    >>> interpolated, boxes = interppatch(img, labels, method="linear")
+    """
     interpolated = img_data + 0.0
     justboxes = img_data * 0.0
     numregions = np.max(separatedimage)
@@ -493,7 +693,8 @@ def interppatch(img_data, separatedimage, method="linear", debug=False):
 
 
 if __name__ == "__main__":
-    """Apply Gaussian smooth to image
+    """
+    Apply Gaussian smooth to image
     Parameters
     ----------
     fnm : str