Solve an equation given a QR factored matrix and a diagonal.

Parameters:
r     - **input-output** n-by-n array. The full lower triangle contains
        the full lower triangle of R. On output, the strict upper
        triangle contains the transpose of the strict lower triangle of
        S.
pmut  - n-vector describing the permutation matrix P.
ddiag - n-vector containing the diagonal of the matrix D in the base
        problem (see below).
bqt   - n-vector containing the first n elements of B Q^T.
sdiag - output n-vector. It is filled with the diagonal of S. Should
        be preallocated by the caller -- can result in somewhat greater
        efficiency if the vector is reused from one call to the next.

Returns:
x     - n-vector solving the equation.

Compute the n-vector x such that

A^T x = B, D x = 0

where A is an n-by-m matrix, B is an m-vector, and D is an n-by-n
diagonal matrix. We are given information about pivoted QR
factorization of A with permutation, such that

A P = R Q

where P is a permutation matrix, Q has orthogonal rows, and R is lower
triangular with nonincreasing diagonal elements. Q is m-by-m, R is
n-by-m, and P is n-by-n. If x = P z, then we need to solve

R z = B Q^T,
P^T D P z = 0 (why the P^T? and do these need to be updated for the transposition?)

If the system is rank-deficient, these equations are solved as well as
possible in a least-squares sense. For the purposes of the LM
algorithm we also compute the lower triangular n-by-n matrix S such
that

P^T (A^T A + D D) P = S^T S. (transpose?)

def _qrd_solve(r, pmut, ddiag, bqt, sdiag):
    """Solve an equation given a QR factored matrix and a diagonal.

Parameters:
r     - **input-output** n-by-n array. The full lower triangle contains
        the full lower triangle of R. On output, the strict upper
        triangle contains the transpose of the strict lower triangle of
        S.
pmut  - n-vector describing the permutation matrix P.
ddiag - n-vector containing the diagonal of the matrix D in the base
        problem (see below).
bqt   - n-vector containing the first n elements of B Q^T.
sdiag - output n-vector. It is filled with the diagonal of S. Should
        be preallocated by the caller -- can result in somewhat greater
        efficiency if the vector is reused from one call to the next.

Returns:
x     - n-vector solving the equation.

Compute the n-vector x such that

A^T x = B, D x = 0

where A is an n-by-m matrix, B is an m-vector, and D is an n-by-n
diagonal matrix. We are given information about pivoted QR
factorization of A with permutation, such that

A P = R Q

where P is a permutation matrix, Q has orthogonal rows, and R is lower
triangular with nonincreasing diagonal elements. Q is m-by-m, R is
n-by-m, and P is n-by-n. If x = P z, then we need to solve

R z = B Q^T,
P^T D P z = 0 (why the P^T? and do these need to be updated for the transposition?)

If the system is rank-deficient, these equations are solved as well as
possible in a least-squares sense. For the purposes of the LM
algorithm we also compute the lower triangular n-by-n matrix S such
that

P^T (A^T A + D D) P = S^T S. (transpose?)
"""

    n, m = r.shape

    # "Copy r and bqt to preserve input and initialize s.  In
    # particular, save the diagonal elements of r in x."  Recall that
    # on input only the full lower triangle of R is meaningful, so we
    # can mirror that into the upper triangle without issues.

    for i in range(n):
        r[i,i:] = r[i:,i]

    x = r.diagonal().copy()
    zwork = bqt.copy()

    # "Eliminate the diagonal matrix d using a Givens rotation."

    for i in range(n):
        # "Prepare the row of D to be eliminated, locating the
        # diagonal element using P from the QR factorization."

        li = pmut[i]
        if ddiag[li] == 0:
            sdiag[i] = r[i,i]
            r[i,i] = x[i]
            continue

        sdiag[i:] = 0
        sdiag[i] = ddiag[li]

        # "The transformations to eliminate the row of d modify only a
        # single element of (q transpose)*b beyond the first n, which
        # is initially zero."

        bqtpi = 0.

        for j in range(i, n):
            # "Determine a Givens rotation which eliminates the
            # appropriate element in the current row of D."

            if sdiag[j] == 0:
                continue

            if abs(r[j,j]) < abs(sdiag[j]):
                cot = r[j,j] / sdiag[j]
                sin = 0.5 / np.sqrt(0.25 + 0.25 * cot**2)
                cos = sin * cot
            else:
                tan = sdiag[j] / r[j,j]
                cos = 0.5 / np.sqrt(0.25 + 0.25 * tan**2)
                sin = cos * tan

            # "Compute the modified diagonal element of r and the
            # modified element of ((q transpose)*b,0)."
            r[j,j] = cos * r[j,j] + sin * sdiag[j]
            temp = cos * zwork[j] + sin * bqtpi
            bqtpi = -sin * zwork[j] + cos * bqtpi
            zwork[j] = temp

            # "Accumulate the transformation in the row of s."
            if j + 1 < n:
                temp = cos * r[j,j+1:] + sin * sdiag[j+1:]
                sdiag[j+1:] = -sin * r[j,j+1:] + cos * sdiag[j+1:]
                r[j,j+1:] = temp

        # Save the diagonal of S and restore the diagonal of R
        # from its saved location in x.
        sdiag[i] = r[i,i]
        r[i,i] = x[i]

    # "Solve the triangular system for z.  If the system is singular
    # then obtain a least squares solution."

    nsing = n

    for i in range(n):
        if sdiag[i] == 0.:
            nsing = i
            zwork[i:] = 0
            break

    if nsing > 0:
        zwork[nsing-1] /= sdiag[nsing-1] # Degenerate case
        # "Reverse loop"
        for i in range(nsing - 2, -1, -1):
            s = np.dot(zwork[i+1:nsing], r[i,i+1:nsing])
            zwork[i] = (zwork[i] - s) / sdiag[i]

    # "Permute the components of z back to components of x."
    x[pmut] = zwork
    return x