regress ¶

def fisher_weight(n: int, alpha: float) -> float:
    """
    Compute a reliability weight using Fisher's z-transformation.

    The weight is inversely related to the confidence interval width in z-space:
    smaller confidence intervals indicate more reliable correlations.

    Args:
        n: Sample size (must be > 3 for a meaningful result).
        alpha: Significance level for the two-tailed confidence interval.

    Returns:
        A weight in [0.0, 1.0], where larger values indicate higher reliability.
    """
    from scipy import stats

    if n <= 3:
        return 0.0

    se = 1.0 / sqrt(n - 3)
    zcrit = stats.norm.ppf(1 - alpha / 2)
    width = 2 * zcrit * se  # total CI width in z-space

    # Map width -> weight: smaller width = higher weight
    # 1 / (1 + width) ensures it stays in (0,1] smoothly
    return clamp(1.0 / (1.0 + float(width)), 0.0, 1.0)

def get_weight(radius: int, alpha: float, method: str) -> float:
    """
    Select a reliability weight computation method.

    Args:
        radius: Neighborhood radius used for regression. Determines sample size.
        alpha: Significance/confidence level for the statistical method.
        method: Weighting method, either "threshold" (t-based) or "fisher".

    Returns:
        A weight in [0.0, 1.0].

    Raises:
        CustomValueError: If `method` is not recognized.
    """
    n = (2 * radius + 1) ** 2  # window size

    if method == "threshold":
        return t_based_weight(n, alpha)

    if method == "fisher":
        return fisher_weight(n, alpha)

    raise CustomValueError("Unknown weight_method", get_weight, method)

def reconstruct(
    clip: vs.VideoNode,
    r: RegressClips,
    radius: int = 1,
    mean: MeanMode | MeanFunction = lambda c, radius: box_blur(c, radius, mode=ConvMode.HV),
    eps: float = 1e-7,
    func: FuncExcept | None = None,
) -> vs.VideoNode:
    """
    Reconstruct a predicted clip from regression results.

    Uses the regression slope, intercept (if available), and correlation to reconstruct
    the dependent variable from the input clip.

    Args:
        clip: Input clip (independent variable).
        r: Regression results (slope, intercept, correlation).
        radius: Neighborhood radius for smoothing regression coefficients. Defaults to 1.
        mean: Function or mode used to compute local means. Defaults to box blur.
        eps: Small constant to avoid division by zero. Defaults to 1e-7.
        func: Optional function reference for error handling context.

    Returns:
        A new clip representing the reconstructed dependent variable.
    """
    func = func or reconstruct

    UnsupportedSampleTypeError.check(
        [clip, *filter(None, r)], vs.FLOAT, func, "All input clips must have a FLOAT sample type."
    )

    if isinstance(mean, MeanMode):
        mean = mean.single

    corr_sum = mean(r.correlation, radius)

    slope_pm = norm_expr([r.slope, r.correlation], "x y *", func=func)
    slope_pm_sum = mean(slope_pm, radius)

    if r.intercept:
        intercept_pm = norm_expr([r.intercept, r.correlation], "x y *", func=func)
        intercept_pm_sum = mean(intercept_pm, radius)
    else:
        intercept_pm_sum = clip.std.BlankClip(keep=True, color=[0] * clip.format.num_planes)

    return norm_expr([clip, slope_pm_sum, intercept_pm_sum, corr_sum], f"x y * z + a {eps} + /", func=func)

def regression(
    x: vs.VideoNode,
    *ys: vs.VideoNode,
    radius: int = 1,
    mean: MeanMode | MeanFunction = lambda c, radius: box_blur(c, radius, mode=ConvMode.HV),
    alpha: float = 0.20,
    weight_method: Literal["threshold", "fisher"] = "threshold",
    intercept: float = False,
    eps: float = 1e-7,
    func: FuncExcept | None = None,
) -> list[RegressClips]:
    """
    Perform per-pixel linear regression between a reference clip `x` and one or more target clips `ys`,
    using local neighborhoods for mean/variance estimates.

    More information:
        - [Simple_linear_regression](https://en.wikipedia.org/wiki/Simple_linear_regression)

    Each regression fits the model:

        y ≈ x * slope + intercept

    on a per-pixel basis, with slope, intercept (optional),
    and correlation computed over neighborhoods of size `(2 * radius + 1)^2`.

    Args:
        x: Reference clip (independent variable).
        *ys: One or more clips to regress on `x` (dependent variables).
        radius: Neighborhood radius for mean/variance estimates. Must be > 0.
        mean: Function or mode used to compute local means. Defaults to box blur.
        alpha: Significance/confidence level for reliability weighting. Defaults to 0.20.
        weight_method: Method for computing reliability weight, "threshold" or "fisher".
        intercept: If True, compute regression intercepts. If False, intercepts are None.
        eps: Small constant to avoid division by zero. Defaults to 1e-7.
        func: An optional function to use for error handling.

    Raises:
        CustomValueError: If `radius` or `alpha` are out of bounds, or if an unknown `weight_method` is provided.

    Returns:
        A list of `RegressClips`, one for each input clip in `ys`, containing slope, intercept (or None),
        and correlation.
    """
    func = func or regression

    UnsupportedSampleTypeError.check([x, *ys], vs.FLOAT, func, "All input clips must have a FLOAT sample type.")

    if not radius > 0:
        raise CustomValueError('"radius" must be greater than zero.', func, radius)

    if not 0 < alpha <= 1.0:
        raise CustomValueError('"alpha" must be greater than zero and lower than 1.0 (inclusive).', func, alpha)

    if isinstance(mean, MeanMode):
        mean = mean.single

    mean_x, *mean_ys = (mean(c, radius) for c in (x, *ys))
    mean_xx, *mean_yys = (mean(norm_expr(c, "src0 dup *", func=func), radius) for c in (x, *ys))
    mean_xys = [mean(norm_expr([x, y], "src0 src1 *", func=func), radius) for y in ys]

    var_x, *var_ys = (
        norm_expr([mean_zz, mean_z], "src0 src1 dup * - 0 max", func=func)
        for mean_zz, mean_z in zip([mean_xx, *mean_yys], [mean_x, *mean_ys])
    )

    cov_xys = [norm_expr([xy, mean_x, y], "src0 src1 src2 * -", func=func) for xy, y in zip(mean_xys, mean_ys)]

    slopes = [norm_expr([cov, var_x], f"src0 src1 {eps} + /", func=func) for cov in cov_xys]

    if intercept:
        intercepts = [
            norm_expr([y, slope, mean_x], f"src0 src1 src2 * - {float(intercept)} /", func=func)
            for y, slope in zip(mean_ys, slopes)
        ]
    else:
        intercepts = [None] * len(slopes)

    w = get_weight(radius, alpha, weight_method)

    corrs = [
        norm_expr(
            [cov_xy, var_x, var_y],
            [f"src0 dup * src1 src2 * {eps} + / sqrt", f"{1 - w} - {w} / 0 max" if w != 1.0 else ""],
            func=func,
        ).std.SetFrameProp("RegressWeight", floatval=w)
        for cov_xy, var_y in zip(cov_xys, var_ys)
    ]

    return [RegressClips(s, i, c) for s, i, c in zip(slopes, intercepts, corrs)]

def t_based_weight(n: int, alpha: float) -> float:
    """
    Compute a reliability weight based on the t-distribution critical correlation.

    The weight decreases with larger critical correlation thresholds (harder to reach significance)
    and increases with sample size.

    Args:
        n: Sample size (must be > 3 for a meaningful result).
        alpha: Significance level for the two-tailed test.

    Returns:
        A weight in [0.0, 1.0], where larger values indicate higher reliability.
    """
    from scipy import stats

    if n <= 3:
        return 0.0

    df = n - 2
    tcrit = stats.t.ppf(1 - alpha / 2, df)
    rcrit = tcrit / sqrt(tcrit**2 + df)

    # Smaller critical r (easier to detect significance) = higher reliability.
    # Map rcrit in [0,1] to weight in [0,1]: weight = 1 - rcrit
    return clamp(1.0 - float(rcrit), 0.0, 1.0)