Core

def run_treat2p(
    suite2p_results_path: Union[str, Path],
    *,
    plane: int = 0,
    chan: Literal[1, 2] = 1,
    rootnames: dict = {},
    extension="npy",
    backend="loky",
    n_jobs=6,
    run_name="",
    do_figures=True,
    save_results=True,
    **kwargs: Any,
) -> tuple[dict, dict, dict]:
    """Run Treat2p processing on Suite2p outputs, optionally in parallel.

    Loads Suite2p output arrays (e.g., fluorescence and neuropil traces), applies a
    pipeline of corrections/modifications (e.g., neuropil factor estimation and/or
    slow trend correction), optionally generates figures, and optionally saves the
    results back into the Suite2p folder.

    Args:
        suite2p_results_path (Union[str, pathlib.Path]): Path to the Suite2p results
            directory (containing plane folders and/or output ``.npy`` files).
        plane (int, optional): Plane index to process. Defaults to 0.
        chan (Literal[1, 2], optional): Channel number (1-based, as in Suite2p).
            Use 2 for channel #2. Defaults to 1.
        rootnames (dict, optional): Optional mapping to override default Suite2p
            file basenames used when loading arrays.
            Example: ``{"Fneu": "custom_Fneu_name"}``. Defaults to ``{}``.
        extension (str, optional): Input/output file extension (e.g., ``"npy"``).
            Defaults to ``"npy"``.
        backend (str, optional): Joblib parallel backend to use. Defaults to
            ``"loky"``.
        n_jobs (int, optional): Number of parallel workers. Use 1 to disable
            parallel processing. Defaults to 6.
        run_name (str, optional): Optional label for this run (used for naming
            outputs/figures depending on implementation). Defaults to ``""``.
        do_figures (bool, optional): Whether to generate diagnostic figures.
            Defaults to True.
        save_results (bool, optional): Whether to write modified arrays/results
            back to disk (Suite2p folder). Defaults to True.
        **kwargs: Additional keyword arguments forwarded to pipeline modules.
            Common keys include:

            Neuropil factor / ARDSIP options:
                neuropil_correction_factor (Optional[float]): If provided, uses
                    this neuropil correction factor and skips ARDSIP estimation.
                    Must be in [0, 1]. Typical values are 0.5 or 0.7. Defaults to
                    None.
                do_slow_trend_correction (bool): Whether to perform slow-trend
                    neuropil correction. Defaults to True.
                max_iterations (int): Maximum ARDSIP iterations. Defaults to 10.
                long_window_length (int): Long sliding window length. Defaults to 30.
                short_window_length (int): Short sliding window length. Defaults to 3.
                baseline_percentile (float): Baseline fluorescence percentile.
                    Defaults to 10.
                convergence_criteria (float): Convergence criterion. Defaults to 0.025.
                autoregressed_activity_threshold (float): Activity threshold used
                    by autoregression. Defaults to 0.5.
                initial_neuropil_factor (float): Initial neuropil factor.
                    Defaults to 0.5.
                downsampling (float): Downsampling factor for baseline computation.
                    1 means no downsampling; 4 means use 1 sample out of 4.
                    Larger values are faster but less faithful. Defaults to 4.
                zscoring_mid_percentile (float): Upper percentile used for z-scoring.
                    Defaults to 16.
                zscoring_low_percentile (float): Lower percentile used for z-scoring.
                    Defaults to 2.3.

            Slow-trend estimation options:
                window (int): Window (seconds) for percentile calculation.
                    Defaults to 120.
                slow_trend_percentile (int): Percentile used for slow-trend
                    calculation. Defaults to 10.
                fps (int): Acquisition sampling rate (Hz). Defaults to 30.
                step (int): Step (samples) between percentile computations.
                    Defaults to 10.
                filter_savgol (bool): Whether to apply a Savitzky-Golay filter to
                    the slow trend. Defaults to True.
                savgol_window_length (int): Savitzky-Golay window length.
                    Defaults to 50.
                savgol_polyorder (int): Savitzky-Golay polynomial order.
                    Defaults to 2.

    Returns:
        tuple[dict, dict, dict]: Three dictionaries containing outputs produced by
        the pipeline (exact contents depend on enabled modules and implementation).
    """

    plane_root = Path(suite2p_results_path) / f"plane{plane}"

    def make_file_path(_type, channel_dependant=True) -> Path:
        if _type in rootnames.keys():
            rootname = rootnames[_type]
        else:
            rootname = _type
        chan_supp = f"_chan{chan}" if chan > 1 and channel_dependant else ""
        file_path = plane_root / f"{rootname}{chan_supp}.{extension}"
        return file_path

    def make_figures_root() -> Union[Path, None]:
        if not do_figures:
            return None
        chan_supp = f"_chan{chan}" if chan > 1 else ""
        figures_root = plane_root / f"treat2p{chan_supp}"
        figures_root.mkdir(parents=True, exist_ok=True)
        return figures_root

    local_log = getLogger("treat2p")

    F_file = make_file_path("F")
    Fneu_file = make_file_path("Fneu")

    stat_file = make_file_path("stat", channel_dependant=False)
    ops_file = make_file_path("ops", channel_dependant=False)

    Fs = np.load(F_file, allow_pickle=True)
    Fneus = np.load(Fneu_file, allow_pickle=True)

    original_stats = np.load(stat_file, allow_pickle=True)
    original_ops = np.load(ops_file, allow_pickle=True).item()

    local_log.info(f"Found F, Fneu, stats and ops files at the specified root : {plane_root}")

    local_log.info(f"Running the signal treatment pipeline over {Fs.shape[0]} rois on {n_jobs} processors. Hold tight.")
    with parallel_backend(backend, n_jobs=n_jobs), tqdm_joblib(
        tqdm(desc="Treating ROIS signal", total=Fs.shape[0], file=stdout)
    ):
        results = Parallel()(
            delayed(pipeline)(F, Fneu, roi_nb, figures_root=make_figures_root(), **kwargs)
            for roi_nb, (F, Fneu) in enumerate(zip(Fs, Fneus))
        )

    local_log.info("Processing finished. About to save the output")

    # save outputs to suite2p results folder
    outputs, corr_stats, ops_list = (
        [item[0] for item in results],
        [item[1] for item in results],
        [item[2] for item in results],
    )

    # package and eventually save F treated files
    reformated_outputs = {}
    for name in outputs[0].keys():
        file_name = make_file_path(f"{name}{run_name}")
        reformated_outputs[name] = np.array([output[name] for output in outputs], dtype=np.float32)
        if save_results:
            np.save(file_name, reformated_outputs[name], allow_pickle=True)
            # save Fcorr, Fnorm etc...

    # package and eventually save stat file
    final_stats = original_stats.copy()
    for roi_nb in range(final_stats.shape[0]):
        final_stats[roi_nb].update(corr_stats[roi_nb])
    if save_results:
        final_stat_file = make_file_path(f"stat{run_name}", channel_dependant=False)
        np.save(final_stat_file, final_stats, allow_pickle=True)
        # save new stats with treat2p field

    # package and eventually save ops file
    final_ops = original_ops.copy()
    final_ops.update({"treat2p": ops_list[0]})
    if save_results:
        final_ops_file = make_file_path(f"ops{run_name}", channel_dependant=False)
        np.save(final_ops_file, np.array(final_ops), allow_pickle=True)
        # save new ops with treat2p field

    # returns them for good measure
    return reformated_outputs, final_stats, final_ops

def pipeline(
    F: np.ndarray,
    Fneu: np.ndarray,
    roi_nb: int,
    *,
    do_slow_trend_correction: bool = True,
    do_delta_f: bool = True,
    do_normalization: bool = True,
    normalization_percentile=None,
    neuropil_correction_factor=None,
    F0_index=5,
    F0_span=30,
    # if not None, set it's value to the neuropil correction and skip neuropil_factor_ARDSIP estimation.
    # Must be between 0 and 1.
    # A reasonable "good" value is usually 0.7 or .5
    figures_root: Optional[Path] = None,
    **kwargs: Any,
) -> tuple[dict, dict, dict]:
    """Run the Treat2p correction pipeline for a single ROI.

    Applies (optionally) slow-trend correction, neuropil correction, ΔF/F
    computation, and normalization to the ROI fluorescence trace using Suite2p
    outputs.

    The keyword arguments in ``kwargs`` are forwarded to the underlying helper
    functions (e.g. ARDSIP neuropil factor estimation and slow-trend
    estimation).

    Args:
        F (np.ndarray): Fluorescence trace of the ROI/cell (shape: ``(T,)``).
        Fneu (np.ndarray): Neuropil fluorescence trace associated with the ROI
            (shape: ``(T,)``).
        roi_nb (int): ROI index/identifier (used for bookkeeping/figure naming).

        do_slow_trend_correction (bool): Whether to perform slow-trend correction
            (typically on the neuropil and/or corrected signal). Defaults to True.
        do_delta_f (bool): Whether to compute ΔF/F. Defaults to True.
        do_normalization (bool): Whether to normalize the output trace. Defaults to True.

        normalization_percentile (Optional[float]): Percentile parameter used by
            the normalization step (interpretation depends on the normalization
            implementation). If None, uses the default behavior of the
            normalization function. Defaults to None.

        neuropil_correction_factor (Optional[float]): If provided, use this value
            as the neuropil correction factor and skip ARDSIP estimation. Must
            be between 0 and 1 (typical values: 0.5–0.7). Defaults to None.

        F0_index (int): Index (or method selector, depending on implementation)
            used for baseline (F0) computation. Defaults to 5.
        F0_span (int): Span/window length used for baseline (F0) computation.
            Defaults to 30.

        figures_root (Optional[Path]): If provided, save diagnostic figures under
            this directory. If None, figures are not saved (or are handled by
            defaults elsewhere). Defaults to None.

        **kwargs: Additional parameters forwarded to underlying steps.

            ARDSIP neuropil factor estimation (e.g. ``neuropil_factor_ARDSIP``):
                max_iterations (int): Maximum number of iterations. Default: 10.
                long_window_length (int): Long sliding window length. Default: 30.
                short_window_length (int): Short sliding window length. Default: 3.
                baseline_percentile (float): Baseline fluorescence percentile.
                    Default: 10.
                convergence_criteria (float): Convergence criterion. Default: 0.025.
                autoregressed_activity_threshold (float): Activity threshold.
                    Default: 0.5.
                initial_neuropil_factor (float): Initial neuropil factor.
                    Default: 0.5.
                downsampling (float): Downsampling factor for baseline computation.
                    1 means no downsampling; 4 means use 1 sample out of 4.
                    Default: 4.
                zscoring_mid_percentile (float): Upper percentile for z-scoring.
                    Default: 16.
                zscoring_low_percentile (float): Lower percentile for z-scoring.
                    Default: 2.3.

            Slow-trend estimation (e.g. ``estimate_slow_trend``):
                window (int): Time window (seconds) used for percentile calculation.
                    Default: 120.
                slow_trend_percentile (int): Percentile used for slow-trend
                    calculation. Default: 10.
                fps (int): Sampling rate (Hz). Default: 30.
                step (int): Step (samples) between percentile computations.
                    Default: 10.
                filter_savgol (bool): Whether to apply Savitzky-Golay filtering.
                    Default: True.
                savgol_window_length (int): Savitzky-Golay window length.
                    Default: 50.
                savgol_polyorder (int): Savitzky-Golay polynomial order.
                    Default: 2.

    Returns:
        tuple[dict, dict, dict]: Three dictionaries containing intermediate and/or
        final outputs produced by the pipeline (contents depend on enabled steps).
    """

    new_stats = {}
    outputs = {}

    try:
        from .plots import RoiFigure

        if figures_root is not None:
            fig = RoiFigure.from_roi_number(roi_nb)
        else:
            fig = None
    except ImportError:
        fig = None

    with ParameterRegistrator("treat2p") as ops:
        if fig is not None:
            fig.raw_ax.plot(F, label="raw", lw=0.05)
            fig.raw_ax.plot(Fneu, label="raw neuropil", lw=0.05)

        Fcorrected, stats = neuropil(F, Fneu, neuropil_correction_factor=neuropil_correction_factor, fig=fig, **kwargs)
        new_stats.update(stats)

        if do_slow_trend_correction:
            Fcorrected = slow_trend(Fcorrected, fig=fig, **kwargs)

        outputs["F_corrected"] = Fcorrected

        if do_delta_f:
            outputs["F_var"] = delta_over_F(Fcorrected, F0_index=F0_index, F0_span=F0_span, fig=fig)

        if do_normalization:
            outputs["F_norm"] = center_normalize(Fcorrected, percentile=normalization_percentile, fig=fig)

    if fig is not None and figures_root is not None:
        fig.finalize()
        fig.save(figures_root)
    ops.pop("fig", None)
    return outputs, new_stats, ops