Feature/Sylvester solver

stan/math/fwd/fun.hpp

@@ -113,6 +113,7 @@
 #include <stan/math/fwd/fun/sqrt.hpp>
 #include <stan/math/fwd/fun/square.hpp>
 #include <stan/math/fwd/fun/sum.hpp>
+#include <stan/math/fwd/fun/sylvester_solver.hpp>
 #include <stan/math/fwd/fun/tan.hpp>
 #include <stan/math/fwd/fun/tanh.hpp>
 #include <stan/math/fwd/fun/tcrossprod.hpp>

stan/math/fwd/fun/sylvester_solver.hpp

@@ -0,0 +1,108 @@
+#ifndef STAN_MATH_FWD_FUN_SYLVESTER_SOLVER_HPP
+#define STAN_MATH_FWD_FUN_SYLVESTER_SOLVER_HPP
+
+#include <stan/math/fwd/meta.hpp>
+#include <stan/math/fwd/core.hpp>
+#include <stan/math/fwd/fun/value_of.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <stan/math/prim/fun/to_ref.hpp>
+#include <stan/math/prim/fun/sylvester_solver.hpp>
+#include <complex>
+#include <type_traits>
+
+namespace stan {
+namespace math {
+
+/**
+ * Solve the Sylvester equation A X + X B = C for X, forward mode.
+ *
+ * The value X solves the equation on the values of A, B and C. Differentiating
+ * the equation gives the tangent as the solution of another Sylvester equation
+ * with the same coefficients,
+ *   A X' + X' B = C' - A' X - X B'.
+ * Both solves run on the inner scalar type, so no Schur decomposition is taken
+ * on fvar values.
+ *
+ * @tparam MatA type of A (real or complex, fvar or arithmetic scalars)
+ * @tparam MatB type of B (real or complex, fvar or arithmetic scalars)
+ * @tparam MatC type of C (real or complex, fvar or arithmetic scalars)
+ * @param[in] A n x n matrix
+ * @param[in] B m x m matrix
+ * @param[in] C n x m matrix
+ * @return complex n x m solution X with std::complex<fvar> entries
+ * @throw std::invalid_argument if A or B is not square, or if the
+ *   dimensions of C do not match
+ */
+template <typename MatA, typename MatB, typename MatC,
+          require_all_matrix_t<MatA, MatB, MatC>* = nullptr,
+          require_any_bt_fvar<MatA, MatB, MatC>* = nullptr>
+inline auto sylvester_solver(const MatA& A, const MatB& B, const MatC& C) {
+  using fvar_t
+      = return_type_t<base_type_t<MatA>, base_type_t<MatB>, base_type_t<MatC>>;
+  using ret_type
+      = Eigen::Matrix<std::complex<fvar_t>, Eigen::Dynamic, Eigen::Dynamic>;
+
+  check_square("sylvester_solver", "A", A);
+  check_square("sylvester_solver", "B", B);
+  check_size_match("sylvester_solver", "rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("sylvester_solver", "rows of B", B.rows(), "cols of C",
+                   C.cols());
+  if (unlikely(A.rows() == 0 || B.rows() == 0)) {
+    return ret_type(A.rows(), B.rows());
+  }
+
+  using complex_inner_t = std::complex<typename fvar_t::Scalar>;
+
+  // tangent of an fvar or complex<fvar> matrix, staying real or complex
+  auto tangent = [](const auto& x) {
+    return x
+        .unaryExpr([](const auto& v) {
+          if constexpr (is_complex<std::decay_t<decltype(v)>>::value) {
+            return std::complex<std::decay_t<decltype(v.real().d_)>>(
+                v.real().d_, v.imag().d_);
+          } else {
+            return v.d_;
+          }
+        })
+        .eval();
+  };
+
+  const auto& A_ref = to_ref(A);
+  const auto& B_ref = to_ref(B);
+  const auto& C_ref = to_ref(C);
+  const auto& A_val = to_ref(value_of(A_ref));
+  const auto& B_val = to_ref(value_of(B_ref));
+  using inner_mat_t
+      = Eigen::Matrix<complex_inner_t, Eigen::Dynamic, Eigen::Dynamic>;
+
+  auto X = sylvester_solver(A_val, B_val, value_of(C_ref));
+
+  inner_mat_t rhs = inner_mat_t::Zero(X.rows(), X.cols());
+  if constexpr (is_fvar<base_type_t<MatC>>::value) {
+    rhs += tangent(C_ref);
+  }
+  if constexpr (is_fvar<base_type_t<MatA>>::value) {
+    rhs -= tangent(A_ref) * X;
+  }
+  if constexpr (is_fvar<base_type_t<MatB>>::value) {
+    rhs -= X * tangent(B_ref);
+  }
+  inner_mat_t X_d = sylvester_solver(A_val, B_val, rhs);
+  const Eigen::Index X_rows = X.rows();
+  const Eigen::Index X_cols = X.cols();
+  ret_type res(X_rows, X_cols);
+  for (Eigen::Index j = 0; j < X_cols; ++j) {
+    for (Eigen::Index i = 0; i < X_rows; ++i) {
+      res(i, j)
+          = std::complex<fvar_t>(fvar_t(X(i, j).real(), X_d(i, j).real()),
+                                 fvar_t(X(i, j).imag(), X_d(i, j).imag()));
+    }
+  }
+  return res;
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/prim/fun.hpp

@@ -313,6 +313,7 @@
 #include <stan/math/prim/fun/svd.hpp>
 #include <stan/math/prim/fun/svd_U.hpp>
 #include <stan/math/prim/fun/svd_V.hpp>
+#include <stan/math/prim/fun/sylvester_solver.hpp>
 #include <stan/math/prim/fun/symmetrize_from_lower_tri.hpp>
 #include <stan/math/prim/fun/symmetrize_from_upper_tri.hpp>
 #include <stan/math/prim/fun/tail.hpp>

stan/math/prim/fun/sylvester_solver.hpp

@@ -0,0 +1,75 @@
+#ifndef STAN_MATH_PRIM_FUN_SYLVESTER_SOLVER_HPP
+#define STAN_MATH_PRIM_FUN_SYLVESTER_SOLVER_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <Eigen/Eigenvalues>
+#include <complex>
+#include <type_traits>
+
+namespace stan {
+namespace math {
+
+/**
+ * Solve the Sylvester equation A X + X B = C for X.
+ *
+ * With complex Schur decompositions A = U S U^H and B = V T V^H the
+ * equation reduces to S Y + Y T = U^H C V, solved column by column,
+ * and X = U Y V^H. A unique solution requires lambda_i(A) + mu_j(B) != 0
+ * for all eigenvalues of A and B.
+ *
+ * @tparam MatA type of A (real or complex arithmetic scalars)
+ * @tparam MatB type of B (real or complex arithmetic scalars)
+ * @tparam MatC type of C (real or complex arithmetic scalars)
+ * @param[in] A n x n matrix
+ * @param[in] B m x m matrix
+ * @param[in] C n x m matrix
+ * @return complex n x m solution X
+ * @throw std::invalid_argument if A or B is not square, or if the
+ *   dimensions of C do not match
+ */
+template <typename MatA, typename MatB, typename MatC,
+          require_all_matrix_t<MatA, MatB, MatC>* = nullptr,
+          require_all_bt_arithmetic<MatA, MatB, MatC>* = nullptr>
+inline auto sylvester_solver(MatA&& A, MatB&& B, MatC&& C) {
+  using complex_mat_t
+      = Eigen::Matrix<std::complex<double>, Eigen::Dynamic, Eigen::Dynamic>;
+
+  check_square("sylvester_solver", "A", A);
+  check_square("sylvester_solver", "B", B);
+  check_size_match("sylvester_solver", "rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("sylvester_solver", "rows of B", B.rows(), "cols of C",
+                   C.cols());
+
+  if (unlikely(A.rows() == 0 || B.rows() == 0)) {
+    return complex_mat_t(A.rows(), B.rows());
+  }
+
+  const Eigen::Index m = B.rows();
+
+  // ComplexSchur needs a plain matrix (same pattern as complex_schur_decompose)
+  Eigen::ComplexSchur<plain_type_t<MatA>> schur_A{eval(std::forward<MatA>(A))};
+  auto&& U = schur_A.matrixU();
+  auto&& S = schur_A.matrixT();
+
+  Eigen::ComplexSchur<plain_type_t<MatB>> schur_B{eval(std::forward<MatB>(B))};
+  auto&& V = schur_B.matrixU();
+  auto&& T = schur_B.matrixT();
+
+  // Y holds D = U^H C V initially, overwritten column by column
+  auto&& C_ref = to_ref(std::forward<MatC>(C));
+  auto Y = (U.adjoint() * C_ref * V).eval();
+  plain_type_t<decltype(S)> Sshift = S;
+  for (Eigen::Index i = 0; i < m; ++i) {
+    Y.col(i).noalias() -= Y.leftCols(i) * T.col(i).head(i);
+    Sshift.diagonal() = S.diagonal().array() + T(i, i);
+    Sshift.template triangularView<Eigen::Upper>().solveInPlace(Y.col(i));
+  }
+  return complex_mat_t(U * Y * V.adjoint());
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/prim/meta/base_type.hpp

@@ -3,6 +3,8 @@
 
 #include <stan/math/prim/fun/Eigen.hpp>
 #include <stan/math/prim/meta/is_eigen.hpp>
+#include <stan/math/prim/meta/is_var.hpp>
+#include <stan/math/prim/meta/require_helpers.hpp>
 #include <stan/math/prim/meta/value_type.hpp>
 #include <stan/math/prim/meta/is_vector.hpp>
 #include <type_traits>
@@ -52,5 +54,33 @@ struct base_type<T, std::enable_if_t<is_eigen<T>::value>> {
   using type = base_type_t<typename std::decay_t<T>::Scalar>;
 };
 
+/*! \ingroup require_stan_scalar_real */
+/*! \addtogroup var_types */
+/*! @{ */
+
+/*! \brief Require any of the base types satisfy @ref is_var */
+/*! Unlike @ref require_any_st_var this also matches `std::complex<var>`. */
+/*! Defined here because is_var.hpp cannot include base_type.hpp. */
+/*! @tparam Types The types with a valid overload of @ref base_type available
+ */
+template <typename... Types>
+using require_any_bt_var
+    = require_any_t<is_var<base_type_t<std::decay_t<Types>>>...>;
+/*! @} */
+
+/*! \ingroup require_stan_scalar_real */
+/*! \addtogroup arithmetic_types */
+/*! @{ */
+
+/*! \brief Require all of the base types satisfy std::is_arithmetic */
+/*! Unlike @ref require_all_st_arithmetic this also matches */
+/*! `std::complex<double>`. */
+/*! @tparam Types The types with a valid overload of @ref base_type available
+ */
+template <typename... Types>
+using require_all_bt_arithmetic
+    = require_all_t<std::is_arithmetic<base_type_t<std::decay_t<Types>>>...>;
+/*! @} */
+
 }  // namespace stan
 #endif

stan/math/prim/meta/is_fvar.hpp

@@ -1,6 +1,7 @@
 #ifndef STAN_MATH_PRIM_META_IS_FVAR_HPP
 #define STAN_MATH_PRIM_META_IS_FVAR_HPP
 
+#include <stan/math/prim/meta/base_type.hpp>
 #include <stan/math/prim/meta/require_helpers.hpp>
 #include <stan/math/prim/meta/scalar_type.hpp>
 #include <stan/math/prim/meta/value_type.hpp>
@@ -106,6 +107,14 @@ using require_any_st_fvar
 template <typename... Types>
 using require_all_not_st_fvar
     = require_all_not_t<is_fvar<scalar_type_t<std::decay_t<Types>>>...>;
+
+/*! \brief Require any of the base types satisfy @ref is_fvar */
+/*! Unlike @ref require_any_st_fvar this also matches `std::complex<fvar>` */
+/*! @tparam Types The types with a valid overload of @ref base_type available
+ */
+template <typename... Types>
+using require_any_bt_fvar
+    = require_any_t<is_fvar<base_type_t<std::decay_t<Types>>>...>;
 /*! @} */
 
 }  // namespace stan

stan/math/rev/fun.hpp

@@ -163,6 +163,7 @@
 #include <stan/math/rev/fun/stan_print.hpp>
 #include <stan/math/rev/fun/step.hpp>
 #include <stan/math/rev/fun/sum.hpp>
+#include <stan/math/rev/fun/sylvester_solver.hpp>
 #include <stan/math/rev/fun/tan.hpp>
 #include <stan/math/rev/fun/tanh.hpp>
 #include <stan/math/rev/fun/tcrossprod.hpp>