rapidtide 3.0.11__py3-none-any.whl → 3.1.1__py3-none-any.whl

rapidtide/makelaggedtcs.py

@@ -17,24 +17,76 @@
 #
 #
 import gc
+import logging
+from typing import Any
 
 import numpy as np
+from numpy.typing import NDArray
 
 import rapidtide.genericmultiproc as tide_genericmultiproc
 
 
 def _procOneVoxelMakelagtc(
-    vox,
-    voxelargs,
-    **kwargs,
-):
+    vox: int,
+    voxelargs: list,
+    **kwargs: Any,
+) -> tuple[int, NDArray]:
+    """
+    Process a single voxel to compute lag timecourse using lag timecourse generator.
+
+    This function takes a voxel index and associated arguments to compute a lag
+    timecourse using the provided lag timecourse generator. The computation involves
+    evaluating the generator at timepoints shifted by the specified lag value.
+
+    Parameters
+    ----------
+    vox : int
+        Voxel index identifier used for tracking and debugging purposes.
+    voxelargs : list
+        List containing three elements:
+        1. `lagtcgenerator` - Lag timecourse generator object with `yfromx` method
+        2. `thelag` - Lag value to be subtracted from time axis
+        3. `timeaxis` - Time axis array for evaluation
+    **kwargs : Any
+        Additional keyword arguments that can override default options:
+        - `rt_floattype` : float type for computations (default: np.float64)
+        - `debug` : boolean for debug printing (default: False)
+
+    Returns
+    -------
+    tuple[int, NDArray]
+        Tuple containing:
+        - `vox` : Input voxel index
+        - `thelagtc` : Computed lag timecourse array
+
+    Notes
+    -----
+    The lag value is subtracted from the time axis as of 10/18. Alternative approaches
+    of adding the lag value resulted in poor performance.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>>
+    >>> # Example usage with mock generator
+    >>> class MockGenerator:
+    ...     def yfromx(self, x):
+    ...         return x**2
+    ...
+    >>> generator = MockGenerator()
+    >>> time_axis = np.array([0, 1, 2, 3, 4])
+    >>> voxel_args = [generator, 1.0, time_axis]
+    >>> result = _procOneVoxelMakelagtc(0, voxel_args)
+    >>> print(result)
+    (0, array([0., 1., 4., 9., 16.]))
+    """
     # unpack arguments
     options = {
-        "rt_floatset": np.float64,
+        "rt_floattype": np.float64,
         "debug": False,
     }
     options.update(kwargs)
-    rt_floatset = options["rt_floatset"]
+    rt_floattype = options["rt_floattype"]
     debug = options["debug"]
     (lagtcgenerator, thelag, timeaxis) = voxelargs
     if debug:
@@ -42,7 +94,7 @@ def _procOneVoxelMakelagtc(
 
     # question - should maxlag be added or subtracted?  As of 10/18, it is subtracted
     #  potential answer - tried adding, results are terrible.
-    thelagtc = rt_floatset(lagtcgenerator.yfromx(timeaxis - thelag))
+    thelagtc = (lagtcgenerator.yfromx(timeaxis - thelag)).astype(rt_floattype)
 
     return (
         vox,
@@ -50,29 +102,163 @@ def _procOneVoxelMakelagtc(
     )
 
 
-def _packvoxeldata(voxnum, voxelargs):
+def _packvoxeldata(voxnum: int, voxelargs: list) -> list:
+    """
+    Pack voxel data into a list format.
+
+    Parameters
+    ----------
+    voxnum : int
+        The index used to select an element from the second element of voxelargs.
+    voxelargs : list
+        A list containing three elements where:
+        - voxelargs[0] is the first element to be returned
+        - voxelargs[1] is a list or array from which an element is selected using voxnum
+        - voxelargs[2] is the third element to be returned
+
+    Returns
+    -------
+    list
+        A list containing three elements: [voxelargs[0], voxelargs[1][voxnum], voxelargs[2]]
+
+    Notes
+    -----
+    This function assumes that voxelargs[1] is indexable and that voxnum is a valid index
+    for accessing elements in voxelargs[1]. The function does not perform any validation
+    of the input parameters.
+
+    Examples
+    --------
+    >>> _packvoxeldata(1, [10, [20, 30, 40], 50])
+    [10, 30, 50]
+
+    >>> _packvoxeldata(0, ['a', ['b', 'c', 'd'], 'e'])
+    ['a', 'b', 'e']
+    """
     return [voxelargs[0], (voxelargs[1])[voxnum], voxelargs[2]]
 
 
-def _unpackvoxeldata(retvals, voxelproducts):
+def _unpackvoxeldata(retvals: tuple, voxelproducts: list) -> None:
+    """
+    Unpack voxel data into the specified voxel products array.
+
+    This function takes return values and assigns them to a specific location
+    in the voxel products array. It is typically used as part of a larger
+    voxel processing pipeline.
+
+    Parameters
+    ----------
+    retvals : tuple
+        A tuple containing the data to be unpacked. The first element is used
+        as an index for the voxel products array, and the second element contains
+        the actual data to be assigned.
+    voxelproducts : list
+        A list of arrays where the voxel data will be stored. The function
+        modifies the first element of this list in-place.
+
+    Returns
+    -------
+    None
+        This function modifies the voxelproducts list in-place and does not
+        return any value.
+
+    Notes
+    -----
+    The function assumes that retvals[0] is a valid index for the first element
+    of voxelproducts and that retvals[1] has compatible dimensions for assignment.
+
+    Examples
+    --------
+    >>> retvals = (5, [1, 2, 3, 4, 5])
+    >>> voxelproducts = [[0] * 10]
+    >>> _unpackvoxeldata(retvals, voxelproducts)
+    >>> print(voxelproducts[0])
+    [0, 0, 0, 0, 0, 1, 2, 3, 4, 5]
+    """
     (voxelproducts[0])[retvals[0], :] = retvals[1]
 
 
 def makelaggedtcs(
-    lagtcgenerator,
-    timeaxis,
-    lagmask,
-    lagtimes,
-    lagtc,
-    LGR=None,
-    nprocs=1,
-    alwaysmultiproc=False,
-    showprogressbar=True,
-    chunksize=1000,
-    rt_floatset=np.float64,
-    rt_floattype="float64",
-    debug=False,
-):
+    lagtcgenerator: Any,
+    timeaxis: NDArray,
+    lagmask: NDArray,
+    lagtimes: NDArray,
+    lagtc: NDArray,
+    LGR: logging.Logger | None = None,
+    nprocs: int = 1,
+    alwaysmultiproc: bool = False,
+    showprogressbar: bool = True,
+    chunksize: int = 1000,
+    rt_floattype: np.dtype = np.float64,
+    debug: bool = False,
+) -> int:
+    """
+    Generate lagged timecourses for a set of voxels using multiprocessing.
+
+    This function computes lagged timecourses for each voxel specified in the mask,
+    using the provided lag timecourse generator and time axis. It supports
+    parallel processing for performance optimization.
+
+    Parameters
+    ----------
+    lagtcgenerator : Any
+        A callable or object that generates lagged timecourses for a single voxel.
+    timeaxis : NDArray
+        1D array representing the time axis (e.g., TRs or time points).
+    lagmask : NDArray
+        3D or 4D boolean or integer array defining the voxels to process.
+        Non-zero entries indicate voxels to be processed.
+    lagtimes : NDArray
+        1D array of lag times (in seconds or time units) to be applied.
+    lagtc : NDArray
+        4D array of shape (ntimepoints, nvoxels, nlags) to store the output lagged
+        timecourses. This is updated in-place.
+    LGR : logging.Logger, optional
+        Logger instance for logging messages. If None, no logging is performed.
+    nprocs : int, optional
+        Number of processes to use for multiprocessing. Default is 1.
+    alwaysmultiproc : bool, optional
+        If True, always use multiprocessing even for single voxel processing.
+        Default is False.
+    showprogressbar : bool, optional
+        If True, display a progress bar during processing. Default is True.
+    chunksize : int, optional
+        Size of chunks to process in each step when using multiprocessing.
+        Default is 1000.
+    rt_floattype : str, optional
+        String representation of the floating-point type.
+        Default is `np.float64`.
+    debug : bool, optional
+        If True, print debug information. Default is False.
+
+    Returns
+    -------
+    int
+        Total number of voxels processed.
+
+    Notes
+    -----
+    This function uses `tide_genericmultiproc.run_multiproc` internally to
+    distribute voxel processing across multiple processes. It is designed for
+    efficient batch processing of large 4D datasets.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> timeaxis = np.arange(100)
+    >>> lagtimes = np.array([0, 1, 2])
+    >>> lagmask = np.ones((10, 10, 10), dtype=bool)
+    >>> lagtc = np.zeros((100, 1000, 3))
+    >>> result = makelaggedtcs(
+    ...     lagtcgenerator=my_generator,
+    ...     timeaxis=timeaxis,
+    ...     lagmask=lagmask,
+    ...     lagtimes=lagtimes,
+    ...     lagtc=lagtc,
+    ...     nprocs=4
+    ... )
+    >>> print(f"Processed {result} voxels")
+    """
     if debug:
         print("makelaggedtcs: Starting")
         print(f"\t{lagtc.shape=}")
@@ -103,7 +289,7 @@ def makelaggedtcs(
         alwaysmultiproc,
         showprogressbar,
         chunksize,
-        rt_floatset=rt_floatset,
+        rt_floattype=rt_floattype,
     )
     if LGR is not None:
         LGR.info(f"\nLagged timecourses created for {volumetotal} voxels")