Repository: scipy/scipy
File: scipy/stats/_stats_py.py (Python)

Code before:

    Generate some data:

    >>> x = rng.random(10)
    >>> y = 1.6*x + rng.random(10)

    Perform the linear regression:

    >>> res = stats.linregress(x, y)

    Coefficient of determination (R-squared):

    >>> print(f"R-squared: {res.rvalue**2:.6f}")
    R-squared: 0.717533

    Plot the data along with the fitted line:

    >>> plt.plot(x, y, 'o', label='original data')
    >>> plt.plot(x, res.intercept + res.slope*x, 'r', label='fitted line')
    >>> plt.legend()
    >>> plt.show()

    Calculate 95% confidence interval on slope and intercept:

    >>> # Two-sided inverse Students t-distribution
    >>> # p - probability, df - degrees of freedom
    >>> from scipy.stats import t
    >>> tinv = lambda p, df: abs(t.ppf(p/2, df))

    >>> ts = tinv(0.05, len(x)-2)
    >>> print(f"slope (95%): {res.slope:.6f} +/- {ts*res.stderr:.6f}")
    slope (95%): 1.453392 +/- 0.743465
    >>> print(f"intercept (95%): {res.intercept:.6f}"
    ...       f" +/- {ts*res.intercept_stderr:.6f}")
    intercept (95%): 0.616950 +/- 0.544475

    """
    np = array_namespace(x, y)  # needed for vecdot pre NumPy 2.0

    TINY = 1.0e-20

    axis = -1
    n = x.shape[axis]
    xmean = np.mean(x, axis=axis, keepdims=True)
    ymean = np.mean(y, axis=axis, keepdims=True)

    # Average sums of square differences from the mean
    #   ssxm = mean( (x-mean(x))^2 )
    #   ssxym = mean( (x-mean(x)) * (y-mean(y)) )
    x_ = _demean(x, xmean, axis=axis, xp=np)
    y_ = _demean(y, ymean, axis=axis, xp=np)

Review comment (suggestion):
OK to branch on array size until we want to support`marray`. (Maybe lazy arrays will be unhappy, too? We'll see.)

Code after:
    # Average sums of square differences from the mean
    #   ssxm = mean( (x-mean(x))^2 )
    #   ssxym = mean( (x-mean(x)) * (y-mean(y)) )
    x_ = _demean(x, xmean, axis=-1, xp=xp)
    y_ = _demean(y, ymean, axis=-1, xp=xp, precision_warning=False)
    xmean = xp.squeeze(xmean, axis=-1)
    ymean = xp.squeeze(ymean, axis=-1)

    ssxm = xp.vecdot(x_, x_, axis=-1) / n
    ssym = xp.vecdot(y_, y_, axis=-1) / n
    ssxym = xp.vecdot(x_, y_, axis=-1) / n

    # R-value
    #   r = ssxym / sqrt( ssxm * ssym )
    degenerate = (ssxm == 0.0) | (ssym == 0.0)
    NaN = xp.asarray(xp.nan, dtype=ssxym.dtype)
    r = xpx.apply_where(
        ~degenerate,
        (ssxym, ssxm, ssym),
        lambda ssxym, ssxm, ssym: xp.clip(ssxym / xp.sqrt(ssxm * ssym), -1.0, 1.0),
        lambda ssxym, ssxm, ssym: xp.where(ssxym==0, NaN, 0.0)
    )

    slope = ssxym / ssxm
    intercept = ymean - slope*xmean
    if n == 2:
        # handle case when only two points are passed in
        one = xp.asarray(1.0, dtype=r.dtype)
        prob = xp.where(y[..., 0] == y[..., 1], one, 0.0)
        slope_stderr = xp.zeros_like(r)
        intercept_stderr = xp.zeros_like(r)
    else:
        df = n - 2  # Number of degrees of freedom
        # n-2 degrees of freedom because 2 has been used up
        # to estimate the mean and standard deviation
        t = r * xp.sqrt(df / ((1.0 - r + TINY)*(1.0 + r + TINY)))

        dist = _SimpleStudentT(xp.asarray(df, dtype=t.dtype))
        prob = _get_pvalue(t, dist, alternative, xp=xp)
        prob = prob[()] if prob.ndim == 0 else prob

        slope_stderr = xp.sqrt((1 - r**2) * ssym / ssxm / df)

        # Also calculate the standard error of the intercept
        # The following relationship is used:
        #   ssxm = mean( (x-mean(x))^2 )
        #        = ssx - sx*sx
        #        = mean( x^2 ) - mean(x)^2
        intercept_stderr = slope_stderr * xp.sqrt(ssxm + xmean**2)

    return LinregressResult(slope=slope[()], intercept=intercept[()], rvalue=r[()],