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