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

rapidtide/simfuncfit.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.
@@ -19,8 +19,10 @@
 import bisect
 import gc
 import logging
+from typing import Any, Callable, Optional, Tuple, Union
 
 import numpy as np
+from numpy.typing import ArrayLike, NDArray
 from tqdm import tqdm
 
 import rapidtide.fit as tide_fit
@@ -30,20 +32,76 @@ LGR = logging.getLogger("GENERAL")
 
 
 def onesimfuncfit(
-    correlationfunc,
-    thefitter,
-    disablethresholds=False,
-    initiallag=None,
-    despeckle_thresh=5.0,
-    lthreshval=0.0,
-    fixdelay=False,
-    fixeddelayvalue=0.0,
-    rt_floatset=np.float64,
-    rt_floattype="float64",
-):
+    correlationfunc: ArrayLike,
+    thefitter: Any,
+    disablethresholds: bool = False,
+    initiallag: Optional[float] = None,
+    despeckle_thresh: float = 5.0,
+    lthreshval: float = 0.0,
+    fixdelay: bool = False,
+    initialdelayvalue: float = 0.0,
+    rt_floattype: np.dtype = np.float64,
+) -> Tuple[int, float, float, float, int, int, int, int]:
+    """
+    Perform a single fit on a correlation function using the provided fitter.
+
+    This function sets up the fitter with initial parameters and thresholds,
+    then performs a fit on the provided correlation function. If `fixdelay` is
+    True, the fit is skipped and a fixed delay value is used instead.
+
+    Parameters
+    ----------
+    correlationfunc : ArrayLike
+        The correlation function data to be fitted.
+    thefitter : Any
+        An object with methods `setguess`, `setrange`, `setlthresh`, and `fit`.
+    disablethresholds : bool, optional
+        If True, disables the threshold setting in the fitter. Default is False.
+    initiallag : float, optional
+        Initial guess for the lag value. If None, no initial guess is set.
+        Default is None.
+    despeckle_thresh : float, optional
+        Threshold for despeckling. Default is 5.0.
+    lthreshval : float, optional
+        Low threshold value for the fitter. Default is 0.0.
+    fixdelay : bool, optional
+        If True, uses a fixed delay value instead of performing a fit.
+        Default is False.
+    initialdelayvalue : float, optional
+        The fixed delay value to use when `fixdelay=True`. Default is 0.0.
+    rt_floattype : np.dtype, optional
+        The data type to use for floating-point values. Default is `np.float64`.
+
+    Returns
+    -------
+    tuple of (int, float, float, float, int, int, int, int)
+        A tuple containing:
+        - maxindex (int): Index of the maximum value in the correlation function.
+        - maxlag (float): The lag value at the maximum.
+        - maxval (float): The maximum value in the correlation function.
+        - maxsigma (float): The sigma (standard deviation) of the fit.
+        - maskval (int): A mask value indicating fit quality.
+        - peakstart (int): Start index of the fitted peak.
+        - peakend (int): End index of the fitted peak.
+        - failreason (int): Reason for fit failure (0 if successful).
+
+    Notes
+    -----
+    When `fixdelay=True`, the function bypasses the fitting process and returns
+    precomputed values based on `initialdelayvalue`.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> corr_func = np.random.rand(100)
+    >>> fitter = some_fitter_class()
+    >>> result = onesimfuncfit(corr_func, fitter)
+    >>> print(result)
+    (50, 0.5, 0.95, 0.02, 1, 45, 55, 0)
+    """
     if initiallag is not None:
         thefitter.setguess(True, maxguess=initiallag)
-        widthmax = despeckle_thresh
+        thefitter.setrange(-despeckle_thresh / 2.0, despeckle_thresh / 2.0)
     else:
         thefitter.setguess(False)
 
@@ -66,10 +124,10 @@ def onesimfuncfit(
     else:
         # do something different
         failreason = np.uint32(0)
-        maxlag = rt_floatset(fixeddelayvalue)
-        maxindex = np.int16(bisect.bisect_left(thefitter.corrtimeaxis, fixeddelayvalue))
-        maxval = rt_floatset(correlationfunc[maxindex])
-        maxsigma = rt_floatset(1.0)
+        maxlag = initialdelayvalue
+        maxindex = np.int16(bisect.bisect_left(thefitter.corrtimeaxis, initialdelayvalue))
+        maxval = correlationfunc[maxindex]
+        maxsigma = 1.0
         maskval = np.uint16(1)
         peakstart = maxindex
         peakend = maxindex
@@ -78,17 +136,80 @@ def onesimfuncfit(
 
 
 def _procOneVoxelFitcorr(
-    vox,
-    corr_y,
-    thefitter,
-    disablethresholds=False,
-    despeckle_thresh=5.0,
-    initiallag=None,
-    fixdelay=False,
-    fixeddelayvalue=0.0,
-    rt_floatset=np.float64,
-    rt_floattype="float64",
-):
+    vox: int,
+    corr_y: ArrayLike,
+    thefitter: Any,
+    disablethresholds: bool = False,
+    despeckle_thresh: float = 5.0,
+    initiallag: Optional[float] = None,
+    fixdelay: bool = False,
+    initialdelayvalue: float = 0.0,
+    rt_floattype: np.dtype = np.float64,
+) -> Tuple[int, int, float, float, float, NDArray, NDArray, float, int, int]:
+    """
+    Process a single voxel for correlation fitting.
+
+    This function performs correlation fitting on a single voxel's data using the provided fitter.
+    It returns fitting results including time, strength, sigma, Gaussian fit, window mask, R² value,
+    and metadata such as mask value and failure reason.
+
+    Parameters
+    ----------
+    vox : int
+        Voxel index.
+    corr_y : ArrayLike
+        Correlation data for the voxel.
+    thefitter : Any
+        Fitter object containing fitting parameters and methods.
+    disablethresholds : bool, optional
+        If True, disables thresholding during fitting. Default is False.
+    despeckle_thresh : float, optional
+        Threshold for despeckling. Default is 5.0.
+    initiallag : float, optional
+        Initial lag value for fitting. Default is None.
+    fixdelay : bool, optional
+        If True, fixes the delay during fitting. Default is False.
+    initialdelayvalue : float, optional
+        Initial delay value if `fixdelay` is True. Default is 0.0.
+    rt_floattype : np.dtype, optional
+        Type to use for real-valued floating-point arrays. Default is `np.float64`.
+
+    Returns
+    -------
+    tuple of (int, int, float, float, float, ndarray, ndarray, float, int, int)
+        A tuple containing:
+        - `vox`: Voxel index.
+        - `volumetotalinc`: 1 if fit was successful, 0 otherwise.
+        - `thetime`: Fitted time value.
+        - `thestrength`: Fitted strength value.
+        - `thesigma`: Fitted sigma value.
+        - `thegaussout`: Gaussian fit evaluated over the time axis.
+        - `thewindowout`: Binary window mask indicating the peak region.
+        - `theR2`: R-squared value of the fit.
+        - `maskval`: Mask value from the fitting process.
+        - `failreason`: Reason for failure, if any.
+
+    Notes
+    -----
+    - If `maxval > 0.3`, plotting is disabled.
+    - The function uses `onesimfuncfit` to perform the actual fitting.
+    - The `thefitter` object must have attributes like `zerooutbadfit`, `lagmod`, and `corrtimeaxis`.
+
+    Examples
+    --------
+    >>> result = _procOneVoxelFitcorr(
+    ...     vox=10,
+    ...     corr_y=corr_data,
+    ...     thefitter=fitter_obj,
+    ...     disablethresholds=False,
+    ...     despeckle_thresh=5.0,
+    ...     fixdelay=False,
+    ...     initialdelayvalue=0.0,
+    ...     rt_floattype=np.float64,
+    ... )
+    >>> print(result)
+    (10, 1, 1.23, 0.95, 0.12, array([...]), array([...]), 0.90, 1, 0)
+    """
     (
         maxindex,
         maxlag,
@@ -104,9 +225,8 @@ def _procOneVoxelFitcorr(
         disablethresholds=disablethresholds,
         despeckle_thresh=despeckle_thresh,
         fixdelay=fixdelay,
-        fixeddelayvalue=fixeddelayvalue,
+        initialdelayvalue=initialdelayvalue,
         initiallag=initiallag,
-        rt_floatset=rt_floatset,
         rt_floattype=rt_floattype,
     )
 
@@ -115,29 +235,29 @@ def _procOneVoxelFitcorr(
 
     # now tuck everything away in the appropriate output array
     volumetotalinc = 0
-    thewindowout = rt_floatset(0.0 * corr_y)
+    thewindowout = np.zeros_like(corr_y, rt_floattype)
     thewindowout[peakstart : peakend + 1] = 1.0
     if (maskval == 0) and thefitter.zerooutbadfit:
-        thetime = rt_floatset(0.0)
-        thestrength = rt_floatset(0.0)
-        thesigma = rt_floatset(0.0)
-        thegaussout = 0.0 * corr_y
-        theR2 = rt_floatset(0.0)
+        thetime = 0.0
+        thestrength = 0.0
+        thesigma = 0.0
+        thegaussout = np.zeros_like(corr_y, rt_floattype)
+        theR2 = 0.0
     else:
         volumetotalinc = 1
-        thetime = rt_floatset(np.fmod(maxlag, thefitter.lagmod))
-        thestrength = rt_floatset(maxval)
-        thesigma = rt_floatset(maxsigma)
-        thegaussout = rt_floatset(0.0 * corr_y)
-        thewindowout = rt_floatset(0.0 * corr_y)
+        thetime = np.fmod(maxlag, thefitter.lagmod)
+        thestrength = maxval
+        thesigma = maxsigma
+        thegaussout = np.zeros_like(corr_y, rt_floattype)
+        thewindowout = np.zeros_like(corr_y, rt_floattype)
         if (not fixdelay) and (maxsigma != 0.0):
-            thegaussout = rt_floatset(
-                tide_fit.gauss_eval(thefitter.corrtimeaxis, [maxval, maxlag, maxsigma])
-            )
+            thegaussout = tide_fit.gauss_eval(
+                thefitter.corrtimeaxis, [maxval, maxlag, maxsigma]
+            ).astype(rt_floattype)
         else:
-            thegaussout = rt_floatset(0.0)
-            thewindowout = rt_floatset(0.0)
-        theR2 = rt_floatset(thestrength * thestrength)
+            thegaussout = 0.0
+            thewindowout = 0.0
+        theR2 = thestrength * thestrength
 
     return (
         vox,
@@ -154,30 +274,112 @@ def _procOneVoxelFitcorr(
 
 
 def fitcorr(
-    corrtimescale,
-    thefitter,
-    corrout,
-    lagmask,
-    failimage,
-    lagtimes,
-    lagstrengths,
-    lagsigma,
-    gaussout,
-    windowout,
-    R2,
-    despeckling=False,
-    peakdict=None,
-    nprocs=1,
-    alwaysmultiproc=False,
-    fixdelay=False,
-    fixdelayvalue=0.0,
-    showprogressbar=True,
-    chunksize=1000,
-    despeckle_thresh=5.0,
-    initiallags=None,
-    rt_floatset=np.float64,
-    rt_floattype="float64",
-):
+    corrtimescale: ArrayLike,
+    thefitter: Any,
+    corrout: NDArray,
+    lagmask: NDArray,
+    failimage: NDArray,
+    lagtimes: NDArray,
+    lagstrengths: NDArray,
+    lagsigma: NDArray,
+    gaussout: NDArray,
+    windowout: NDArray,
+    R2: NDArray,
+    despeckling: bool = False,
+    peakdict: Optional[dict] = None,
+    nprocs: int = 1,
+    alwaysmultiproc: bool = False,
+    fixdelay: bool = False,
+    initialdelayvalue: Union[float, NDArray] = 0.0,
+    showprogressbar: bool = True,
+    chunksize: int = 1000,
+    despeckle_thresh: float = 5.0,
+    initiallags: Optional[NDArray] = None,
+    rt_floattype: np.dtype = np.float64,
+) -> int:
+    """
+    Fit correlation data to extract lag parameters and related statistics for each voxel.
+
+    This function performs a fitting procedure on correlation data for each voxel,
+    extracting lag times, strengths, sigma values, Gaussian fits, window functions,
+    and R² values. It supports both single-threaded and multi-threaded processing.
+
+    Parameters
+    ----------
+    corrtimescale : ArrayLike
+        Time scale of the correlation data.
+    thefitter : Any
+        Fitter object used to perform the fitting. Must have methods like `setcorrtimeaxis`.
+    corrout : NDArray
+        Correlation data for all voxels, shape (n_voxels, n_timepoints).
+    lagmask : NDArray
+        Mask indicating valid lags for each voxel, shape (n_voxels,).
+    failimage : NDArray
+        Image to store failure flags for each voxel, shape (n_voxels,).
+    lagtimes : NDArray
+        Output array for lag times, shape (n_voxels,).
+    lagstrengths : NDArray
+        Output array for lag strengths, shape (n_voxels,).
+    lagsigma : NDArray
+        Output array for lag sigma values, shape (n_voxels,).
+    gaussout : NDArray
+        Output array for Gaussian fit parameters, shape (n_voxels, n_timepoints).
+    windowout : NDArray
+        Output array for window function values, shape (n_voxels, n_timepoints).
+    R2 : NDArray
+        Output array for R² values, shape (n_voxels,).
+    despeckling : bool, optional
+        If True, performs despeckling pass, only accepting successful fits, by default False.
+    peakdict : dict, optional
+        Dictionary of peak information, by default None.
+    nprocs : int, optional
+        Number of processes to use for multiprocessing, by default 1.
+    alwaysmultiproc : bool, optional
+        If True, always use multiprocessing even for single process, by default False.
+    fixdelay : bool, optional
+        If True, fix the delay value, by default False.
+    initialdelayvalue : Union[float, NDArray], optional
+        Initial delay value(s), by default 0.0.
+    showprogressbar : bool, optional
+        If True, show progress bar, by default True.
+    chunksize : int, optional
+        Size of chunks for multiprocessing, by default 1000.
+    despeckle_thresh : float, optional
+        Threshold for despeckling, by default 5.0.
+    initiallags : NDArray, optional
+        Initial lag values for each voxel, by default None.
+    rt_floattype : np.dtype, optional
+        Floating-point type for runtime, by default np.float64.
+
+    Returns
+    -------
+    int
+        Total number of voxels successfully processed.
+
+    Notes
+    -----
+    The function modifies the input arrays (`lagtimes`, `lagstrengths`, `lagsigma`,
+    `gaussout`, `windowout`, `R2`, `lagmask`, `failimage`) in-place.
+
+    Examples
+    --------
+    >>> fitcorr(
+    ...     corrtimescale=timescale,
+    ...     thefitter=fitter,
+    ...     corrout=correlation_data,
+    ...     lagmask=lag_mask,
+    ...     failimage=fail_image,
+    ...     lagtimes=lag_times,
+    ...     lagstrengths=lag_strengths,
+    ...     lagsigma=lag_sigma,
+    ...     gaussout=gaussian_out,
+    ...     windowout=window_out,
+    ...     R2=r2_values,
+    ...     nprocs=4,
+    ...     despeckling=True,
+    ... )
+    12345
+    """
     thefitter.setcorrtimeaxis(corrtimescale)
     inputshape = np.shape(corrout)
     if initiallags is None:
@@ -208,17 +410,17 @@ def fitcorr(
                         break
 
                     # process and send the data
-                    if (lagmask is None) or (initiallags is None):
+                    if (themask is None) or (initiallags is None):
                         thislag = None
                     else:
-                        if lagmask[val] > 0:
+                        if themask[val] > 0:
                             thislag = initiallags[val]
                         else:
                             thislag = None
-                    """if initiallags is None:
-                        thislag = None
+                    if isinstance(initialdelayvalue, np.ndarray):
+                        thisinitialdelayvalue = initialdelayvalue[val]
                     else:
-                        thislag = initiallags[val]"""
+                        thisinitialdelayvalue = initialdelayvalue
                     outQ.put(
                         _procOneVoxelFitcorr(
                             val,
@@ -228,8 +430,7 @@ def fitcorr(
                             despeckle_thresh=despeckle_thresh,
                             initiallag=thislag,
                             fixdelay=fixdelay,
-                            fixeddelayvalue=fixdelayvalue,
-                            rt_floatset=rt_floatset,
+                            initialdelayvalue=thisinitialdelayvalue,
                             rt_floattype=rt_floattype,
                         )
                     )
@@ -289,28 +490,23 @@ def fitcorr(
             unit="voxels",
             disable=(not showprogressbar),
         ):
-            if themask is None:
-                dothisone = True
+            # process and send the data
+            if (themask is None) or (initiallags is None):
                 thislag = None
-            elif themask[vox] > 0:
                 dothisone = True
-                thislag = initiallags[vox]
             else:
-                dothisone = False
-                thislag = None
+                if themask[vox] > 0:
+                    thislag = initiallags[vox]
+                    dothisone = True
+                else:
+                    thislag = None
+                    dothisone = False
+            if isinstance(initialdelayvalue, np.ndarray):
+                thisinitialdelayvalue = initialdelayvalue[vox]
+            else:
+                thisinitialdelayvalue = initialdelayvalue
             if dothisone:
-                (
-                    dummy,
-                    volumetotalinc,
-                    lagtimes[vox],
-                    lagstrengths[vox],
-                    lagsigma[vox],
-                    gaussout[vox, :],
-                    windowout[vox, :],
-                    R2[vox],
-                    lagmask[vox],
-                    failreason,
-                ) = _procOneVoxelFitcorr(
+                voxel = _procOneVoxelFitcorr(
                     vox,
                     corrout[vox, :],
                     thefitter,
@@ -318,33 +514,41 @@ def fitcorr(
                     despeckle_thresh=despeckle_thresh,
                     initiallag=thislag,
                     fixdelay=fixdelay,
-                    fixeddelayvalue=fixdelayvalue,
-                    rt_floatset=rt_floatset,
+                    initialdelayvalue=thisinitialdelayvalue,
                     rt_floattype=rt_floattype,
                 )
-                volumetotal += volumetotalinc
                 if (
                     thefitter.FML_INITAMPLOW
                     | thefitter.FML_INITAMPHIGH
                     | thefitter.FML_FITAMPLOW
                     | thefitter.FML_FITAMPHIGH
-                ) & failreason:
+                ) & voxel[9]:
                     ampfails += 1
-                if (thefitter.FML_INITWIDTHLOW | thefitter.FML_FITWIDTHLOW) & failreason:
+                if (thefitter.FML_INITWIDTHLOW | thefitter.FML_FITWIDTHLOW) & voxel[9]:
                     lowwidthfails += 1
-                if (thefitter.FML_INITWIDTHHIGH | thefitter.FML_FITWIDTHHIGH) & failreason:
+                if (thefitter.FML_INITWIDTHHIGH | thefitter.FML_FITWIDTHHIGH) & voxel[9]:
                     highwidthfails += 1
-                if (thefitter.FML_INITLAGLOW | thefitter.FML_INITLAGHIGH) & failreason:
+                if (thefitter.FML_INITLAGLOW | thefitter.FML_FITLAGLOW) & voxel[9]:
                     lowlagfails += 1
-                if (thefitter.FML_INITLAGLOW | thefitter.FML_FITLAGLOW) & failreason:
-                    lowlagfails += 1
-                if (thefitter.FML_INITLAGHIGH | thefitter.FML_FITLAGHIGH) & failreason:
+                if (thefitter.FML_INITLAGHIGH | thefitter.FML_FITLAGHIGH) & voxel[9]:
                     highlagfails += 1
-                if thefitter.FML_INITFAIL & failreason:
+                if thefitter.FML_INITFAIL & voxel[9]:
                     initfails += 1
-                if thefitter.FML_FITFAIL & failreason:
+                if thefitter.FML_FITFAIL & voxel[9]:
                     fitfails += 1
 
+                # if this is a despeckle pass, only accept the new values if the fit did not fail
+                if (voxel[9] == 0) or not despeckling:
+                    volumetotal += voxel[1]
+                    lagtimes[vox] = voxel[2]
+                    lagstrengths[vox] = voxel[3]
+                    lagsigma[vox] = voxel[4]
+                    gaussout[vox, :] = voxel[5]
+                    windowout[vox, :] = voxel[6]
+                    R2[vox] = voxel[7]
+                    lagmask[vox] = voxel[8]
+                    failimage[vox] = voxel[9] & 0xFFFF
+
     LGR.info(f"\nSimilarity function fitted in {volumetotal} voxels")
     LGR.info(
         f"\tampfails: {ampfails}"