Square root of an arrowhead block matrix

Question

How do I find a square root of a symmetric matrix:

$W=\begin{bmatrix}D_1 & A \\ A^T & D_2\end{bmatrix}$

where $D_1$ and $D_2$ are diagonal matrices, $D_1$ very large (5000x5000), $D_2$ quite small (3x3 to 100x100), $A$ is a slim sparse matrix (5000x3 to 5000x100). $W$ thus looks like an arrowhead. I need to compute $W^{1 \over 2}$ as efficiently as possible (this computation will repeat many times).

PS: not sure if it will be useful information: the matrix is a Hessian matrix of a continuous likelihood function.

---

EDIT: I found a great paper - Eisenfeld 1976 - Block Diagonalization and Eigenvalues. The Corollary 3.1 is awesome, but it only works for upper triangular matrices. So I implemented the general iterative computation of root matrices in Theorem 4.1, but even with very beautiful matrices with all eigenvalues positive and not small and full rank, the iterations strongly diverged, leading to exponentially higher and higher differences (see the code below):

# Implementation of Eisenfeld 1976 - Block Diagonalization and Eigenvalues
# Generate matrix

<- 8
q <- 3
n <- p+q

p <- 100
q <- 10
n <- p+q

p <- 2900
q <- 100
n <- p+q

# generate matrix
#pi <- runif(p, 0, 1)
fi <- rnorm(p)
gi <- rnorm(q)
z <- 2*q # pocet occassions
xcov <- matrix(rnorm(z*q), nrow = z, ncol = q)
pi <- inv.logit(matrix(fi, nrow = p, ncol = z) + matrix(xcov %*% gi, nrow = p, ncol = z, byrow = TRUE))
W <- diag(n)
diag(W)[1:p] <- apply(pi*(1-pi), 1, sum)
#W[1:p,(p+1):n] <- xcov * matrix(pi*(1-pi), nrow = p, ncol = q)
for (j in 1:q) {
    W[1:p,p+j] <- (pi*(1-pi)) %*% xcov[,j,drop = FALSE]
}
W[(p+1):n,1:p] <- t(W[1:p,(p+1):n])
for (i in 1:q) { # nevim jak toto vektorizovat... 3D matice :-D (nebo lze vice nasobenimi?)
    for (j in 1:q) {
        W[p+i,p+j] <- sum( (pi*(1-pi)) %*% (xcov[,i]*xcov[,j]) )
    }
}


# notation by Eisenfeld:

A <- W
stopifnot(all(dim(A) == p+q))

E <- A[1:p,1:p]
F <- A[1:p,(p+1):n]
G <- A[(p+1):n,1:p]
H <- A[(p+1):n,(p+1):n]

# computation by Eisenfeld

# E is diagonal
E.inv <- E
diag(E.inv) <- 1/diag(E.inv)

# in my case also t(F) == G, and H is symmetric

# Theorem 4.1 - iterative computation of the roots R1, R2

# test the conditions
#a =  E.inv %*% t(F) det(F)*det(E.inv)

it <- 0
itmax <- 10

tol <- 1e-6
Y1diff <- Y2diff <- 1
Y1.new <- matrix(0, nrow = q, ncol = p)
Y2.new <- matrix(0, nrow = p, ncol = q)

#Y1.new <- matrix(rnorm(p*q)/1000, nrow = q, ncol = p)
#Y2.new <- matrix(rnorm(p*q)/1000, nrow = p, ncol = q)
while ((Y1diff > tol || Y2diff > tol) && it < itmax) {
    it <- it + 1
    Y1 <- Y1.new
    Y2 <- Y2.new

    # f1(Y) = -YFYE-l + HYE-1 + GE-l
    Y1.new <- -Y1 %*% F %*% Y1 %*% E.inv + H %*% Y1 %*% E.inv + G %*% E.inv

    # f2(Y) = E-1YGY + E-1YH - E-1F
    Y2.new <- E.inv %*% Y2 %*% G %*% Y2 + E.inv %*% Y2 %*% H - E.inv %*% F

    Y1diff <- max(abs(Y1.new - Y1))
    Y2diff <- max(abs(Y2.new - Y2))
    cat("Y1diff = ", Y1diff, ", Y2diff = ", Y2diff, "\n")
}
# doesn't converge even for beautiful positive definite matrix

FINAL SOLUTION (but not of the square root): In the end, I got a great advice that in my case, I don't need a square root decomposition in particular, Cholesky would do as well. And that saved me. The Cholesky decomposition of my matrix can actually be done much simpler and very fast!! I used the Block LU decomposition and this completely solved my problem :)

I don't post this as an answer, since it did not answer the original question of the square root; but it solved my problem, Cholesky is perfect fit in this context as well.

Answer 1

$W$ is $\geq 0$, that is $D_1\geq 0,(I-D_1(D_1)^+)A=0,D_2-A^T(D_1)^+A\geq 0$ where $()^+$ denotes the Moore-Penrose inverse.

EDIT 1.

In particular, $A$ is "small". Roughly speaking, $W$ has a "large" diagonal but $W$ is not necessarily diagonally dominant; on the other hand, the fact that $(D_i)$ is diagonal seems to be useless.

If $W$ is invertible, then a stable method is the Denman Beavers algorithm.

$Y_0=W,Z_0=I$ and the iteration

$Y_{k+1}=0.5(Y_k+(Z_k)^{-1}),Z_{k+1}=0.5(Z_k+(Y_k)^{-1})$.

The complexity is in $\approx 2kn^3$ where $k$ is the number of used iterations; for example, if $k=20$, then the complexity is in $\approx 40n^3$. To compare, the complexity of the standard method (by diagonalization of the symmetric matrix W) is in $≈30n^3$ and is not stable when the matrix is badly conditioned.

EDIT 2. If you have really a fast method to calculate the Cholesky decomposition $W=L^TL$, then you can do as follows (note that $L$ is a full matrix).

i) Calculate the $QR$ decomposition $L=QR$, where $R$ is $>0$ symmetric. For stability, use a Newton's algorithm with $k$ iterations; the complexity is $\approx kn^3$.

ii) Then $W=R^TQ^TQR=R^2$.

Of course, the total complexity is in $O(n^3)$.