Feature/lyapunov

stan/math/fwd/fun/solve_sylvester.hpp

@@ -103,6 +103,87 @@ inline auto solve_sylvester(const MatA& A, const MatB& B, const MatC& C) {
   return res;
 }
 
+/**
+ * Solve the continuous Lyapunov equation A X + X A^H = C for X, forward mode.
+ *
+ * The value X solves the equation on the values of A and C. Differentiating
+ * the equation gives the tangent as the solution of another Lyapunov equation
+ * with the same coefficient,
+ *   A X' + X' A^H = C' - A' X - X A'^H.
+ * 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 MatC type of C (real or complex, fvar or arithmetic scalars)
+ * @param[in] A n x n matrix
+ * @param[in] C n x n matrix
+ * @return complex n x n solution X with std::complex<fvar> entries
+ * @throw std::invalid_argument if A is not square, or if the dimensions of C
+ *   do not match
+ */
+template <typename MatA, typename MatC,
+          require_all_matrix_t<MatA, MatC>* = nullptr,
+          require_any_bt_fvar<MatA, MatC>* = nullptr>
+inline auto solve_lyapunov(const MatA& A, const MatC& C) {
+  using fvar_t = return_type_t<base_type_t<MatA>, base_type_t<MatC>>;
+  using ret_type
+      = Eigen::Matrix<std::complex<fvar_t>, Eigen::Dynamic, Eigen::Dynamic>;
+
+  check_square("solve_lyapunov", "A", A);
+  check_size_match("solve_lyapunov", "rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("solve_lyapunov", "rows of A", A.rows(), "cols of C",
+                   C.cols());
+  if (unlikely(A.rows() == 0)) {
+    return ret_type(0, 0);
+  }
+
+  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& C_ref = to_ref(C);
+  const auto& A_val = to_ref(value_of(A_ref));
+  using inner_mat_t
+      = Eigen::Matrix<complex_inner_t, Eigen::Dynamic, Eigen::Dynamic>;
+
+  auto X = solve_lyapunov(A_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) {
+    const auto A_d = tangent(A_ref);
+    rhs -= A_d * X + X * A_d.adjoint();
+  }
+  inner_mat_t X_d = solve_lyapunov(A_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/solve_sylvester.hpp

@@ -70,6 +70,62 @@ inline auto solve_sylvester(MatA&& A, MatB&& B, MatC&& C) {
   return complex_mat_t(U * Y * V.adjoint());
 }
 
+/**
+ * Solve the continuous Lyapunov equation A X + X A^H = C for X.
+ *
+ * Special case of solve_sylvester with B = A^H. The complex Schur
+ * decomposition A = U T U^H also gives A^H = U T^H U^H, so only one
+ * decomposition is needed. The equation reduces to T Y + Y T^H = U^H C U,
+ * solved column by column from last to first, and X = U Y U^H. A unique
+ * solution requires lambda_i(A) + conj(lambda_j(A)) != 0 for all
+ * eigenvalues of A.
+ *
+ * @tparam MatA type of A (real or complex arithmetic scalars)
+ * @tparam MatC type of C (real or complex arithmetic scalars)
+ * @param[in] A n x n matrix
+ * @param[in] C n x n matrix
+ * @return complex n x n solution X
+ * @throw std::invalid_argument if A is not square, or if the dimensions of C
+ *   do not match
+ */
+template <typename MatA, typename MatC,
+          require_all_matrix_t<MatA, MatC>* = nullptr,
+          require_all_bt_arithmetic<MatA, MatC>* = nullptr>
+inline auto solve_lyapunov(MatA&& A, MatC&& C) {
+  using complex_mat_t
+      = Eigen::Matrix<std::complex<double>, Eigen::Dynamic, Eigen::Dynamic>;
+
+  check_square("solve_lyapunov", "A", A);
+  check_size_match("solve_lyapunov", "rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("solve_lyapunov", "rows of A", A.rows(), "cols of C",
+                   C.cols());
+
+  if (unlikely(A.rows() == 0)) {
+    return complex_mat_t(0, 0);
+  }
+
+  const Eigen::Index n = A.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&& T = schur_A.matrixT();
+
+  // Y holds D = U^H C U initially; T^H is lower triangular, so Y is
+  // overwritten column by column from last to first
+  auto&& C_ref = to_ref(std::forward<MatC>(C));
+  auto Y = (U.adjoint() * C_ref * U).eval();
+  plain_type_t<decltype(T)> Tshift = T;
+  for (Eigen::Index j = n - 1; j >= 0; --j) {
+    const Eigen::Index k = n - 1 - j;  // already-solved columns
+    Y.col(j).noalias() -= Y.rightCols(k) * T.row(j).tail(k).adjoint();
+    Tshift.diagonal() = T.diagonal().array() + std::conj(T(j, j));
+    Tshift.template triangularView<Eigen::Upper>().solveInPlace(Y.col(j));
+  }
+  return complex_mat_t(U * Y * U.adjoint());
+}
+
 }  // namespace math
 }  // namespace stan
 #endif

stan/math/rev/fun/solve_sylvester.hpp

@@ -138,6 +138,123 @@ inline auto solve_sylvester(MatA&& A, MatB&& B, MatC&& C) {
   return ret_type(res);
 }
 
+/**
+ * Solve the continuous Lyapunov equation A X + X A^H = C for X, reverse mode.
+ *
+ * Forward pass as in prim, with a single complex Schur decomposition
+ * A = U T U^H. With G the adjoint of X (real part: adjoint of Re X,
+ * imaginary part: adjoint of Im X) and Z the solution of A^H Z + Z A = G,
+ * the gradients are
+ *   grad_A = -(Z X^H + Z^H X),  grad_C = Z,
+ * where the second term of grad_A comes from A^H in place of B. The adjoint
+ * equation reuses the same Schur factors as T^H W + W T = U^H G U.
+ * Complex inputs receive Re and Im of their gradient; real inputs receive
+ * only the real part.
+ *
+ * @tparam MatA type of A (real or complex, var or double scalars)
+ * @tparam MatC type of C (real or complex, var or double scalars)
+ * @param[in] A n x n matrix
+ * @param[in] C n x n matrix
+ * @return complex n x n solution X with std::complex<var> entries
+ * @throw std::invalid_argument if A is not square, or if the dimensions of C
+ *   do not match
+ */
+template <typename MatA, typename MatC,
+          require_all_matrix_t<MatA, MatC>* = nullptr,
+          require_any_bt_var<MatA, MatC>* = nullptr>
+inline auto solve_lyapunov(MatA&& A, MatC&& C) {
+  using cdouble = std::complex<double>;
+  using ret_type
+      = Eigen::Matrix<std::complex<var>, Eigen::Dynamic, Eigen::Dynamic>;
+
+  check_square("solve_lyapunov", "A", A);
+  check_size_match("solve_lyapunov", "rows of A", A.rows(), "rows of C",
+                   C.rows());
+  check_size_match("solve_lyapunov", "rows of A", A.rows(), "cols of C",
+                   C.cols());
+  if (unlikely(A.rows() == 0)) {
+    return ret_type(0, 0);
+  }
+
+  auto arena_A = to_arena(std::forward<MatA>(A));
+  auto arena_C = to_arena(std::forward<MatC>(C));
+
+  const Eigen::Index n = A.rows();
+
+  // ---- forward pass: identical to prim, on values ----
+  Eigen::ComplexSchur<plain_type_t<decltype(value_of(arena_A))>> schur_A(
+      value_of(arena_A));
+  auto arena_U = to_arena(schur_A.matrixU());
+  auto arena_T = to_arena(schur_A.matrixT());
+
+  // Y holds D = U^H C U initially, overwritten column by column from last
+  // to first
+  auto Y = (arena_U.adjoint() * value_of(arena_C) * arena_U).eval();
+  plain_type_t<decltype(arena_T)> Tshift = arena_T;
+  for (Eigen::Index j = n - 1; j >= 0; --j) {
+    const Eigen::Index k = n - 1 - j;  // already-solved columns
+    Y.col(j).noalias() -= Y.rightCols(k) * arena_T.row(j).tail(k).adjoint();
+    Tshift.diagonal() = arena_T.diagonal().array() + std::conj(arena_T(j, j));
+    Tshift.template triangularView<Eigen::Upper>().solveInPlace(Y.col(j));
+  }
+  arena_t<Eigen::MatrixXcd> arena_X = arena_U * Y * arena_U.adjoint();
+
+  auto res = to_arena(arena_X.cast<std::complex<var>>());
+
+  // ---- reverse pass ----
+  reverse_pass_callback([arena_A, arena_C, arena_U, arena_T, arena_X,
+                         res]() mutable {
+    const Eigen::Index n = arena_T.rows();
+
+    // G = adj(Re X) + i * adj(Im X)
+    Eigen::MatrixXcd G(n, n);
+    for (Eigen::Index c = 0; c < n; ++c) {
+      for (Eigen::Index r = 0; r < n; ++r) {
+        G(r, c) = cdouble(res(r, c).real().adj(), res(r, c).imag().adj());
+      }
+    }
+
+    // W holds E = U^H G U initially; solve T^H W + W T = E,
+    // overwritten column by column from first to last
+    auto W = (arena_U.adjoint() * G * arena_U).eval();
+    auto Tshift_r = arena_T.eval();
+    for (Eigen::Index j = 0; j < n; ++j) {
+      W.col(j).noalias() -= W.leftCols(j) * arena_T.col(j).head(j);
+      Tshift_r.diagonal()
+          = arena_T.diagonal().array() + std::conj(arena_T(j, j));
+      Tshift_r.template triangularView<Eigen::Upper>().adjoint().solveInPlace(
+          W.col(j));
+    }
+    const auto Z = (arena_U * W * arena_U.adjoint()).eval();
+
+    // Re/Im of the gradient go to the adjoints of Re/Im of the input;
+    // real inputs only receive Re(grad)
+    auto accumulate = [](auto& in, const Eigen::MatrixXcd& g) {
+      using in_t = std::decay_t<decltype(in)>;
+      if constexpr (is_complex<value_type_t<in_t>>::value) {
+        const Eigen::Index rows = g.rows();
+        const Eigen::Index cols = g.cols();
+        for (Eigen::Index c = 0; c < cols; ++c) {
+          for (Eigen::Index r = 0; r < rows; ++r) {
+            in(r, c).real().adj() += g(r, c).real();
+            in(r, c).imag().adj() += g(r, c).imag();
+          }
+        }
+      } else {
+        in.adj() += g.real();
+      }
+    };
+
+    if constexpr (is_var<base_type_t<MatA>>::value) {
+      accumulate(arena_A, -(Z * arena_X.adjoint() + Z.adjoint() * arena_X));
+    }
+    if constexpr (is_var<base_type_t<MatC>>::value) {
+      accumulate(arena_C, Z);
+    }
+  });
+  return ret_type(res);
+}
+
 }  // namespace math
 }  // namespace stan
 #endif

test/unit/math/mix/fun/solve_sylvester_test.cpp

@@ -33,14 +33,6 @@ TEST(mathMixMatFun, solve_sylvester) {
   // expect_ad_matvar does not apply: the result has complex<var> scalars,
   // which have no var_value<Matrix> form. See solve_sylvester_varmat.
 
-  // Lyapunov: B = A^T shares A's autodiff variables
-  auto g = [](const auto& a, const auto& c) {
-    return stan::math::solve_sylvester(a, a.transpose(), c);
-  };
-  Eigen::MatrixXd Q(2, 2);
-  Q << 1, 2, 3, 4;
-  stan::test::expect_ad(g, A, Q);
-
   // size mismatch must throw for every autodiff type, as it does for double
   Eigen::MatrixXd C_bad(3, 3);
   C_bad.setOnes();
@@ -164,3 +156,57 @@ TEST(mathMixMatFun, solve_sylvester_result_is_independent) {
                               Eigen::MatrixXd(C2.adj()));
   stan::math::recover_memory();
 }
+
+TEST(mathMixMatFun, solve_lyapunov) {
+  auto f = [](const auto& a, const auto& c) {
+    return stan::math::solve_lyapunov(a, c);
+  };
+
+  // A has a complex eigenvalue pair, so the complex Schur path is used
+  Eigen::MatrixXd A(3, 3);
+  A << -1, -1, 0, 1, -1, 0.5, 0.3, -0.2, -2;
+  Eigen::MatrixXd C(3, 3);
+  C << 1, 2, 3, 4, 5, 6, 7, 8, 9;
+  stan::test::expect_ad(f, A, C);
+
+  // size mismatch must throw for every autodiff type, as it does for double
+  Eigen::MatrixXd C_bad(3, 2);
+  C_bad.setOnes();
+  EXPECT_THROW(stan::math::solve_lyapunov(A, C_bad), std::invalid_argument);
+
+  auto A_v = stan::math::to_var(A);
+  auto C_bad_v = stan::math::to_var(C_bad);
+  EXPECT_THROW(stan::math::solve_lyapunov(A_v, C_bad_v), std::invalid_argument);
+  stan::math::recover_memory();
+}
+
+TEST(mathMixMatFun, solve_lyapunov_complex) {
+  auto f = [](const auto& a, const auto& c) {
+    return stan::math::solve_lyapunov(a, c);
+  };
+  using cdouble = std::complex<double>;
+
+  // eigenvalue sums lambda_i + conj(lambda_j) have real part around -7
+  Eigen::MatrixXcd A(2, 2);
+  A << cdouble(-4, 1), cdouble(1, -0.5), cdouble(-1, 0.3), cdouble(-3, -0.2);
+  Eigen::MatrixXcd C(2, 2);
+  C << cdouble(1, 2), cdouble(-0.5, 1), cdouble(3, -1), cdouble(0.7, 0.2);
+  stan::test::expect_ad(f, A, C);
+
+  // real system, complex right-hand side
+  Eigen::MatrixXd A_r = A.real();
+  stan::test::expect_ad(f, A_r, C);
+}
+
+TEST(mathMixMatFun, solve_lyapunov_triangular) {
+  // Schur of a triangular matrix deflates on exact zeros, which must not
+  // drop the tangents of the zero entries
+  auto f = [](const auto& a, const auto& c) {
+    return stan::math::solve_lyapunov(a, c);
+  };
+  Eigen::MatrixXd A(3, 3);
+  A << -5, 2, -1, 0, -6, 1, 0, 0, -4;
+  Eigen::MatrixXd C(3, 3);
+  C << 1, 2, 3, 4, 5, 6, 7, 8, 9;
+  stan::test::expect_ad(f, A, C);
+}