rapidtide 2.9.6__py3-none-any.whl → 3.1.3__py3-none-any.whl

rapidtide/makelaggedtcs.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 #
-#   Copyright 2016-2024 Blaise Frederick
+#   Copyright 2016-2025 Blaise Frederick
 #
 #   Licensed under the Apache License, Version 2.0 (the "License");
 #   you may not use this file except in compliance with the License.
@@ -18,122 +18,292 @@
 #
 import gc
 import logging
+from typing import Any
 
 import numpy as np
-from tqdm import tqdm
+from numpy.typing import NDArray
 
-import rapidtide.multiproc as tide_multiproc
-
-LGR = logging.getLogger("GENERAL")
+import rapidtide.genericmultiproc as tide_genericmultiproc
 
 
 def _procOneVoxelMakelagtc(
-    vox,
-    lagtcgenerator,
-    maxlag,
-    timeaxis,
-    rt_floatset=np.float64,
-    rt_floattype="float64",
-):
+    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_floattype": np.float64,
+        "debug": False,
+    }
+    options.update(kwargs)
+    rt_floattype = options["rt_floattype"]
+    debug = options["debug"]
+    (lagtcgenerator, thelag, timeaxis) = voxelargs
+    if debug:
+        print(f"{vox=}, {thelag=}, {timeaxis=}")
+
     # 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 - maxlag))
+    thelagtc = (lagtcgenerator.yfromx(timeaxis - thelag)).astype(rt_floattype)
 
     return (
         vox,
-        thelagtc,
+        (thelagtc),
     )
 
 
+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: 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,
-    nprocs=1,
-    alwaysmultiproc=False,
-    showprogressbar=True,
-    chunksize=1000,
-    rt_floatset=np.float64,
-    rt_floattype="float64",
-):
+    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.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=}")
+        print(f"\t{lagtimes.shape=}")
+        print(f"\t{timeaxis.shape=}")
+
     inputshape = lagtc.shape
-    if nprocs > 1 or alwaysmultiproc:
-        # define the consumer function here so it inherits most of the arguments
-        def makelagtc_consumer(inQ, outQ):
-            while True:
-                try:
-                    # get a new message
-                    val = inQ.get()
-
-                    # this is the 'TERM' signal
-                    if val is None:
-                        break
-
-                    # process and send the data
-                    outQ.put(
-                        _procOneVoxelMakelagtc(
-                            val,
-                            lagtcgenerator,
-                            lagtimes[val],
-                            timeaxis,
-                            rt_floatset=rt_floatset,
-                            rt_floattype=rt_floattype,
-                        )
-                    )
-                except Exception as e:
-                    print("error!", e)
-                    break
-
-        data_out = tide_multiproc.run_multiproc(
-            makelagtc_consumer,
-            inputshape,
-            lagmask,
-            nprocs=nprocs,
-            showprogressbar=showprogressbar,
-            chunksize=chunksize,
-        )
-
-        # unpack the data
-        volumetotal = 0
-        for voxel in data_out:
-            volumetotal += 1
-            lagtc[voxel[0], :] = voxel[1]
-        del data_out
-    else:
-        volumetotal = 0
-        for vox in tqdm(
-            range(0, inputshape[0]),
-            desc="Voxel",
-            unit="voxels",
-            disable=(not showprogressbar),
-        ):
-            if lagmask[vox] > 0:
-                dothisone = True
-            else:
-                dothisone = False
-            if dothisone:
-                (
-                    dummy,
-                    lagtc[vox, :],
-                ) = _procOneVoxelMakelagtc(
-                    vox,
-                    lagtcgenerator,
-                    lagtimes[vox],
-                    timeaxis,
-                    rt_floatset=rt_floatset,
-                    rt_floattype=rt_floattype,
-                )
-                volumetotal += 1
-
-    LGR.info(f"\nLagged timecourses created for {volumetotal} voxels")
+    voxelargs = [
+        lagtcgenerator,
+        lagtimes,
+        timeaxis,
+    ]
+    voxelfunc = _procOneVoxelMakelagtc
+    packfunc = _packvoxeldata
+    unpackfunc = _unpackvoxeldata
+    voxeltargets = [lagtc]
+
+    volumetotal = tide_genericmultiproc.run_multiproc(
+        voxelfunc,
+        packfunc,
+        unpackfunc,
+        voxelargs,
+        voxeltargets,
+        inputshape,
+        lagmask,
+        LGR,
+        nprocs,
+        alwaysmultiproc,
+        showprogressbar,
+        chunksize,
+        rt_floattype=rt_floattype,
+    )
+    if LGR is not None:
+        LGR.info(f"\nLagged timecourses created for {volumetotal} voxels")
 
     # garbage collect
     uncollected = gc.collect()
     if uncollected != 0:
-        LGR.info(f"garbage collected - unable to collect {uncollected} objects")
+        if LGR is not None:
+            LGR.info(f"garbage collected - unable to collect {uncollected} objects")
     else:
-        LGR.info("garbage collected")
+        if LGR is not None:
+            LGR.info("garbage collected")
+
+    if debug:
+        print("makelaggedtcs: End\n\n")
 
     return volumetotal