Neuropil Factor¶

@filter_save_kwargs("F_Cell", "F_Neu")
def neuropil_factor_ARDSIP(
    F_Cell: np.ndarray,
    F_Neu: np.ndarray,
    *,
    fig: Optional[RoiFigure] = None,
    max_iterations: int = 10,
    long_window_length: float = 30,  # in seconds
    short_window_length: float = 3,
    baseline_percentile: float = 10,
    convergence_criteria: float = 0.025,
    autoregressed_activity_threshold: float = 0.5,
    initial_neuropil_factor: float = 0.5,
    downsampling: int = 4,  # factor
    zscoring_mid_percentile: float = 16,
    zscoring_low_percentile: float = 2.3,
    fps: float = 30,  # in hertz
) -> dict:
    """This function performs neuropil correction using an AutoRegressed DownSampling Iterative Process (A.R.DS.I.P).

    Args:
        F_Cell (numpy.ndarray): The array of fluorescence intensity for the cell.
        F_Neu (numpy.ndarray): The array of fluorescence intensity for the neuropil.
        max_iterations (int, optional): The maximum number of iterations to perform. Default is 10.
        long_window_length (int, optional): The length of the long sliding window. Default is 30.
        short_window_length (int, optional): The length of the short sliding window. Default is 3.
        baseline_percentile (float, optional): The percentile of the baseline fluorescence. Default is 10.
        convergence_criteria (float, optional): The criteria for convergence. Default is 0.025.
        autoregressed_activity_threshold (float, optional): The threshold for determining if there is activity in the
            cell based on autoregression. Default is 0.5.
        initial_neuropil_factor (float, optional): The initial neuropil factor. Default is 0.5.
        downsampling (float, optional): How much downsampled is the original aray when computing baseline.
            1 means no downsampling, 4 means 1 smple out of 4 is used.
            Higher values speedup the computation at fidelity cost. Default is 4.
        zscoring_mid_percentile (float, optional): Upper percentile used for zscoring. Default is 16.
        zscoring_low_percentile (float, optional): Lower percentile used for zscoring. Default is 2.3.
            Both of zscoring_mid_percentile and zscoring_low_percentile refer to z_score_from_percentile arguments.
            See that function for more details.

    Returns:
        dict: A dictionary containing :
            the calculated neuropil factor (float),
            the final neuropil factor after bounding to [0.1, 0.7] (float),
            a list of the neuropil factor across all iterations (list[float])
            a boolean indicating whether the function converged (bool).
    """

    F_Cell = np.asarray(F_Cell)
    F_Neu = np.asarray(F_Neu)

    # loop changing variable :
    neuropil_factor = initial_neuropil_factor
    local_log = getLogger("neuropil_factor_ardsip")
    local_log.debug("Starting Neuropil Refined Correction")
    iteration_results = [initial_neuropil_factor]
    converged = False

    iteration = 0
    activity_periods_mask = None

    if fig is not None:
        fig.ardsip_figure.ax_raw.plot(F_Cell, label="cell", lw=0.05)
        fig.ardsip_figure.ax_raw.plot(F_Neu, label="neuropil", lw=0.05)

    for iteration in range(max_iterations):
        local_log.debug(f"Starting new iteration with neuropil_factor : {neuropil_factor:.2f} ")
        F_Corrected = F_Cell - (neuropil_factor * F_Neu)
        new_neuropil_factors = []

        if iteration == 0:  # activity_periods_mask is None at that stage, and for further iterations,
            # it has an array shape like F_Corrected
            activity_periods_mask = np.zeros_like(F_Corrected, dtype=bool)

        F_Corrected_activity_removed = F_Corrected.copy()
        F_Corrected_activity_removed[activity_periods_mask] = np.nan
        # suppress periods of ROI activity , in the first iteration include all data
        # ( activity_periods_mask is redefined later)

        low_percentiles = sliding_window_percentile(
            F_Corrected_activity_removed,
            percentile=baseline_percentile,
            window_length=long_window_length,
            window_sample_displacement=downsampling,
            window_internal_downsampling=downsampling + 1,
            samples_per_sec=fps,
        )

        # why not using standard zscoring here >?
        F_Corrected_normalized = z_score_percentile(
            F_Corrected - low_percentiles,
            mid_percentile=zscoring_mid_percentile,
            low_percentile=zscoring_low_percentile,
        )

        F_Autocorr = lagged_autorcorrelation(F_Corrected_normalized)

        bool_no_activity = (F_Autocorr < autoregressed_activity_threshold) & (
            F_Autocorr > -autoregressed_activity_threshold
        )  # periods without activity
        bool_state_nochange = ~np.abs(np.diff(bool_no_activity.astype(int), prepend=False)).astype(
            bool
        )  # periods of changes between activity / non activity
        bool_stable_no_activity = (
            bool_no_activity & bool_state_nochange
        )  # stable periods without activity nor instantaneous changes

        F_Cell_noactivity = F_Cell.copy()
        F_Cell_noactivity[~bool_stable_no_activity] = np.nan

        F_Neu_noactivity = F_Neu.copy()
        F_Neu_noactivity[~bool_stable_no_activity] = np.nan

        if np.sum(bool_stable_no_activity) < 0.1 * len(bool_stable_no_activity):
            local_log.warning(
                "Non activity periods cannot realistically consist of less than 10% of all time. Skipping."
            )
            iteration_results.append(np.nan)
            neuropil_factor = neuropil_factor + convergence_criteria
            continue

        activity_periods_mask = ~bool_stable_no_activity  # activity periods are the inverse of no activity periods
        # we assign it here, only after the continue, in case detected "rest" periods were too few.

        # loop's final estimation of neuropil_factor
        short_window_indices = rolling_indices(
            F_Cell_noactivity,
            window_displacement=short_window_length,
            window_length=short_window_length,
            window_internal_downsampling=1,
            samples_per_sec=fps,
        )[0]

        F_Cell_windows = F_Cell_noactivity[short_window_indices]
        F_Neu_windows = F_Neu_noactivity[short_window_indices]

        for cell_chunk, neu_chunk in zip(F_Cell_windows, F_Neu_windows):
            non_nan_vec = ~np.isnan(cell_chunk)
            if sum(non_nan_vec) > 5:  # at least 6 non nan data points needed for regression
                new_neuropil_factors.append(np.polyfit(neu_chunk[non_nan_vec], cell_chunk[non_nan_vec], 1)[0])
            else:
                new_neuropil_factors.append(np.nan)
        new_neuropil_factors = np.array(new_neuropil_factors)
        local_log.debug(
            "Finished calculating chunks regressions with "
            f"{len(new_neuropil_factors[~np.isnan(new_neuropil_factors)])} valid chunks"
        )
        new_neuropil_factor_mean = float(np.nanmean(new_neuropil_factors))

        divergence = abs(neuropil_factor - new_neuropil_factor_mean)
        neuropil_factor = new_neuropil_factor_mean
        iteration_results.append(neuropil_factor)
        if divergence < convergence_criteria:
            local_log.info(
                f"Found convergence in {iteration + 1} iterations. Ending with neuropil_factor value : "
                f"{neuropil_factor:.2f} (last run divergence : {divergence:.3f})."
            )
            converged = True
            break

    if not converged:
        local_log.info(
            f"Found no convergence in {iteration + 1} iterations. "
            f"Ending with neuropil_factor value : {neuropil_factor:.2f}."
        )

    if neuropil_factor < 0.1:
        constrained_neuropil_factor = 0.1
    elif neuropil_factor > 0.7:
        constrained_neuropil_factor = 0.7
    else:
        constrained_neuropil_factor = neuropil_factor

    if fig is not None:
        fig.ardsip_figure.ax_noactivity.plot(F_Cell_noactivity, label="cell", lw=0.05)
        fig.ardsip_figure.ax_noactivity.plot(F_Neu_noactivity, label="neuropil", lw=0.05)
        fig.ardsip_figure.ax_corrected.plot(F_Corrected, label="cell_corrected", lw=0.05)
        fig.ardsip_figure.ax_corrected_noactivity.plot(
            F_Corrected_activity_removed, label="cell_corr_no_activity", lw=0.05
        )
        fig.ardsip_figure.ax_corr_norm.plot(F_Corrected_normalized, label="cell_corr_normalized", lw=0.05)
        fig.ardsip_figure.ax_low_perc.plot(low_percentiles, label="cell_low_perc", lw=0.05)
        fig.ardsip_figure.ax_autocorr.plot(F_Autocorr, label="cell_lag_autocorrelation", lw=0.05)
        fig.ardsip_figure.ax_autocorr.axhline(
            autoregressed_activity_threshold, label="activity_threshold", ls="--", color="gray"
        )
        fig.ardsip_figure.ax_autocorr.axhline(-autoregressed_activity_threshold, ls="--", color="gray")
        bool_no_activity = bool_no_activity.astype(float)
        bool_state_nochange = bool_state_nochange.astype(float)
        bool_stable_no_activity = bool_stable_no_activity.astype(float)
        bool_no_activity[bool_no_activity < 1] = np.nan
        bool_state_nochange[bool_state_nochange < 1] = np.nan
        bool_stable_no_activity[bool_stable_no_activity < 1] = np.nan
        fig.ardsip_figure.ax_autocorr.plot(bool_no_activity, label="no_activity_periods")
        fig.ardsip_figure.ax_autocorr.plot(bool_state_nochange * 2, label="no_change_periods")
        fig.ardsip_figure.ax_autocorr.plot(bool_stable_no_activity * 3, label="stable_no_activity_periods")

        fig.ardsip_figure.ax_neuropil_factors.plot(new_neuropil_factors, label="neuropil_factor_over_time", lw=0.05)

        fig.ardsip_figure.ax_iteraction_neuropil_factor.plot(
            iteration_results, "o-", label="neuropil_factor_over_iterations"
        )
        fig.ardsip_figure.ax_iteraction_neuropil_factor.plot(
            len(iteration_results), constrained_neuropil_factor, "o", label="final_neu_factor"
        )

        fig.ardsip_figure.ax_corrected.text(
            0.2,
            0.5,
            f"neu_factor={neuropil_factor}",
            transform=fig.ardsip_figure.ax_corrected.transAxes,
            ha="left",
            va="center",
        )

    return {
        "neuropil_factor": neuropil_factor,
        "neuropil_factor_constrained": constrained_neuropil_factor,
        "neuropil_factor_iterations": iteration_results,
        "neuropil_factor_converged": converged,
    }