Math does not have the matrix (inverse) square root

[bgoodri]

Description

You can apply the sqrt function elementwise to a matrix, but Stan Math does not have a function to take the (symmetric) square root of a positive definite matrix (or its inverse), defined as V * diag_matrix(sqrt(v)) * V' where V is a matrix of eigenvectors and v is a vector of eigenvalues. In the case of the inverse square root of a matrix, there is a diag_matrix(inv_sqrt(v)) in the middle. Eigen already calculates these, so it is just a matter of exposing them and working through the Jacobians.

Example

data {
  int<lower = 1> K;
  matrix[K, K] X;
}
transformed data {
  matrix[K, K] sqrt_X = sqrt_spd(X);
  matrix[K, K] inv_sqrt_X = inv_sqrt_spd(X);
  real<lower = 0, upper = 1e-15> check = fabs(X - sqrt_X * sqrt_X);
}

[Note: The sqrt_spd and inv_sqrt_spd functions do not exist in Stan yet.]

Expected Output

Satisfies inequality restrictions on check.

Current Version:

v2.33

[bob-carpenter]

I'd like to boost the priority on this one and I suggest we name the function matrix_sqrt(). Sarah Heaps has an elegant unconstrained parameterization of stationary parameters for a vector autoregressive (VAR) model. It uses matrix square root. She sent along the code she was using for it in Stan, without the new tuple feature

matrix sqrtm(matrix A) {
    int m = rows(A);
    vector[m] root_root_evals = sqrt(sqrt(eigenvalues_sym(A)));
    matrix[m, m] evecs = eigenvectors_sym(A);
    matrix[m, m] eprod = diag_post_multiply(evecs, root_root_evals);
    return tcrossprod(eprod);
  }

and with:

matrix sqrtm(matrix A) {
    int m = rows(A);
    tuple(matrix[m, m], vector[m]) eigen = eigendecompose_sym(A);
    matrix[m, m] eprod = diag_post_multiply(eigen.1, sqrt(sqrt(eigen.2)));
    return tcrossprod(eprod);
  }

I'd also like to add the full constraining transform and the unconstraining transform, but that will be in a separate issue.

[bgoodri]

If we were to boost the priority of this, then https://github.com/KingJamesSong/FastDifferentiableMatSqrt is better than the closed PR I wrote a long time ago.