Computing square root of diag(u)-uu'?

Question

I need an efficient way to take square root of a matrix which is a sum of diagonal matrix and rank-1 matrix.

More specifically it's the following matrix

$$A=D-uu'=\text{diag}(u)-uu'$$

Where entries of $u$ are non-negative and $\sum_i^d u_i=1$. This is also known as the covariance of the multinomial distribution with mean $u$

My first approach was to use Eigenvalue decomposition, but it was about 50 times too slow for my application.

$$ A^{1/2} = US^{1/2}U'$$ The second approach was to to ignore $uu'$ term. This was fast enough, but incurred large error.

$$ A^{1/2} \approx D^{1/2}$$

My third approach was to do some algebraic rearrangement on the error term and come up with the following rank-1 approximation

$$A^{1/2} \approx D^{1/2}+\frac{\sqrt{1+u'D^{-1}u}-1}{u'D^{-1}u}D^{-\frac{1}{4}}uu'D^{-\frac{1}{4}}$$

This is fast enough, and had a small approximation error on a couple of random examples (notebook). However, I'm not sure if there are cases when this approximation breaks down.

At this point I feel like I may be reinventing the wheel -- are there known results for this problem? Specifically, a cheap way to compute it exactly, or approximately with a guaranteed small approximation error.

The values of $u$ are highly skewed in my case, they decay at least as fast as $O(1/k)$ and possibly exponentially.

Answer 1

Would a decomposition of the form $A = XX^T$ suffice? This would be enough, e.g., if the end goal is sampling from the Gaussian distribution with this given covariance.

If so, you can use the following formula, which is quite similar to your approximation: $$X = D^{1/2} + \frac{\sqrt{u^T D^{-1} u+1}-1}{u^T D^{-1} u} u u^T D^{-1/2}$$ This follows from whitening the diagonal term \begin{align} D + uu^T &= D^{1/2}\left(I + D^{-1/2}uu^TD^{-1/2}\right)D^{1/2}, \end{align} then finding the square root of the middle factor, which is a rank-1 update of the identity.

---

If you really need the square root, you may consider using rational approximations, $$(D + uu^T)^{1/2} \approx c_0 + c_1 (\sigma_1 I + D + uu^T)^{-1} + c_2 (\sigma_2 I + D + uu^T)^{-1} + \dots.$$ then inverting each of these terms with the Sherman-Morrison formula. The number of terms you need for a very good approximation grows as $$O(\log \kappa)$$ where $\kappa$ is the condition number of $D + uu^T$. So, instead of a diagonal plus a rank-1 matrix, this becomes a sum of a small number of diagonal plus rank one matrices.

For more details on the rational approximation, here is a wonderful paper:

Hale, Nicholas, Nicholas J. Higham, and Lloyd N. Trefethen. "Computing A^α,\log(A), and related matrix functions by contour integrals." SIAM Journal on Numerical Analysis 46.5 (2008): 2505-2523. http://eprints.maths.manchester.ac.uk/834/1/hale_higham_trefethen.pdf

See, in particular, equation (4.4) and method 3 in that paper.

---

Edit: Here I wrote some Python code that does the rational method:

# Adaptation of Method 3 from Hale, Higham, and Trefethen, Computing f(A)b by contour integrals. SIAM 2008
import numpy as np
import scipy.linalg as sla
from scipy.special import *


def hht_isqrt_weights_and_poles(min_eigenvalue_m, max_eigenvalue_M, number_of_rational_terms_N):
    # 1/sqrt(z) = w0/(z - p0) + w1/(z - p1) + ...
    m = min_eigenvalue_m
    M = max_eigenvalue_M
    N = number_of_rational_terms_N
    k2 = m/M
    Kp = ellipk(1-k2)
    t = 1j * np.arange(0.5, N) * Kp/N
    sn, cn, dn, ph = ellipj(t.imag,1-k2)
    cn = 1./cn
    dn = dn * cn
    sn = 1j * sn * cn
    w = np.sqrt(m) * sn
    dzdt = cn * dn

    poles = (w**2).real
    weights = (2 * Kp * np.sqrt(m) / (np.pi*N)) * dzdt
    rational_function = lambda zz: np.dot(1. / (zz.reshape((-1,1)) - poles), weights)
    return weights, poles, rational_function


def diagonal_smw(dd, u):
    # (diag(dd)+uu^T)^-1 = diag(dd) - vv^T
    v = (u / dd)/np.sqrt(1. + np.dot(u, u / dd))
    return v


def diagonal_plus_rank_one_inverse_sqrt(dd, u, num_terms):
    # inv(sqrtm(diag(dd) + u*u')) = diag(ss) - V*V' + small error
    lambda_min_bound = np.min(dd) # These bounds could be improved
    lambda_max_bound = np.max(dd) + np.linalg.norm(u)**2
    ww, pp, _ = hht_isqrt_weights_and_poles(lambda_min_bound, lambda_max_bound, num_terms)
    dd_shift = [dd - p for p in pp]
    vv0 = [diagonal_smw(d_shift, u) for d_shift in dd_shift]
    ss = np.sum([w*(1./d_shift) for w, d_shift in zip(ww, dd_shift)], axis=0)
    V = np.vstack([np.sqrt(w)*v0 for w, v0 in zip(ww, vv0)]).T
    return ss, V

def diagonal_woodburyish(ss, V):
    # (diag(ss) - VV^T)^-1 = diag(1./ss) + W*M*W^T
    r = V.shape[1]
    W = V / ss.reshape([-1, 1]) # inv(diag(dd))*V
    capacitance_mtx = np.eye(r)-np.dot(V.T, W)
    M = np.linalg.inv(capacitance_mtx) # inverse of very small matrix, e.g., 5x5. Could do Cholesky if desired..
    return M, W

def diagonal_plus_rank_one_sqrt(dd, u, num_terms):
    # inv(sqrtm(diag(dd) + u*u')) = diag(iss) + W*M*W^T + small error
    ss, V = diagonal_plus_rank_one_inverse_sqrt(dd, u, num_terms)
    M, W = diagonal_woodburyish(ss, V)
    iss = 1. / ss
    return iss, M, W

n = 500
kappa_ish = 1e3
dd = kappa_ish * np.random.rand(n)
u = np.random.randn(n)

A = np.diag(dd) + np.outer(u, u)
sqrtA = sla.sqrtm(A)

for num_terms in 1+np.arange(10):
    iss, M, W = diagonal_plus_rank_one_sqrt(dd, u, num_terms)
    sqrtA2 = np.diag(iss) + np.dot(W, np.dot(M, W.T))
    err = np.linalg.norm(sqrtA - sqrtA2) / np.linalg.norm(sqrtA)
    print('num_terms=', num_terms, ', err=', err)

It works quite well, here is the output error for 1 to 10 terms in the rational approximation:

num_terms= 1 , err= 0.28856754578036703
num_terms= 2 , err= 0.04084537953169016
num_terms= 3 , err= 0.005770763245002039
num_terms= 4 , err= 0.000624851647505238
num_terms= 5 , err= 7.234944285553648e-05
num_terms= 6 , err= 1.0059955401611735e-05
num_terms= 7 , err= 1.22949322111201e-06
num_terms= 8 , err= 1.3750242705729841e-07
num_terms= 9 , err= 1.7564866678109768e-08
num_terms= 10 , err= 2.289881883503377e-09