Weighted 1D iterative least squares

leastsq-1d-weighted

import numpy as np

def pearson_r(x: np.ndarray, y: np.ndarray) -> np.ndarray:
    n = len(x)
    return (n * x @ y - x.sum() * y.sum()) / np.sqrt(
        (n * (x @ x) - x.sum() ** 2) * (n * (y @ y) - y.sum() ** 2)
    )


def lstsq_1d(
    x: np.darray,
    y: np.darray,
    w_x: np.ndarray,
    w_y: np.ndarray,
    r: np.ndarray = 0,
    tolerance: float = 1e-15,
    n: int = 50,
) -> np.ndarray:
    """Weighted least squares

    :param x:
    :param y:
    :param w_x:
    :param w_y:
    :param r:
    :param tolerance:
    :param n:
    :return:
    """
    # https://birmingham-primo.hosted.exlibrisgroup.com/permalink/f/19a9mc5/TN_proquest216705896
    # DOI 10.1119/1.1632486
    assert n != 0

    # Initial least squares
    A = np.vstack((x, np.ones_like(x))).T
    (m, c), residuals, rank, s = np.linalg.lstsq(A, y)

    alpha = np.sqrt(w_x * w_y)
    i = 0
    while True:
        if i == n:
            break
        i += 1

        # Step (3)
        w = (w_x * w_y) / (w_x + m ** 2 * w_y - 2 * m * r * alpha)

        # Step (4)
        x_bar = (w @ x) / w.sum()
        y_bar = (w @ y) / w.sum()

        u = x - x_bar
        v = y - y_bar
        beta = w * (u / w_y + m * v / w_x - (m * u + v) * r / alpha)

        # Step (5)
        m_adj = ((w * beta) @ v) / ((w * beta) @ u)

        if np.abs(m - m_adj) <= tolerance:
            break

        m = m_adj

    c_adj = y_bar - m_adj * x_bar

    # Errors
    x_adj = x_bar + beta
    # y_adj = y_bar + m_adj*beta
    x_adj_bar = (w @ x_adj) / w.sum()
    u_adj = x_adj - x_adj_bar

    sigma_m = np.sqrt(1 / (w @ (u_adj ** 2)))
    sigma_c = np.sqrt(1 / w.sum() + x_bar ** 2 * sigma_m ** 2)

    return (m_adj, c_adj), (sigma_m, sigma_c)