Swapped diag(sqrt) for sqrtm(diag) for speed