diff --git a/include/albatross/src/cereal/eigen.hpp b/include/albatross/src/cereal/eigen.hpp index 346e190a..188b24ae 100644 --- a/include/albatross/src/cereal/eigen.hpp +++ b/include/albatross/src/cereal/eigen.hpp @@ -91,6 +91,29 @@ inline void load(Archive &archive, v.indices() = indices; } +template +inline void +save(Archive &archive, + const Eigen::PermutationMatrix &v, + const std::uint32_t) { + archive(cereal::make_nvp("indices", v.indices())); +} + +template +inline void +load(Archive &archive, + Eigen::PermutationMatrix &v, + const std::uint32_t) { + typename Eigen::PermutationMatrix::IndicesType indices; + archive(cereal::make_nvp("indices", indices)); + v.indices() = indices; +} + template inline void serialize(Archive &archive, Eigen::DiagonalMatrix<_Scalar, SizeAtCompileTime> &matrix, diff --git a/include/albatross/src/cereal/gp.hpp b/include/albatross/src/cereal/gp.hpp index b3ec473c..b3926397 100644 --- a/include/albatross/src/cereal/gp.hpp +++ b/include/albatross/src/cereal/gp.hpp @@ -41,8 +41,8 @@ inline void serialize(Archive &archive, Fit> &fit, archive(cereal::make_nvp("information", fit.information)); archive(cereal::make_nvp("train_covariance", fit.train_covariance)); archive(cereal::make_nvp("train_features", fit.train_features)); - archive(cereal::make_nvp("sigma_R", fit.sigma_R)); - archive(cereal::make_nvp("permutation_indices", fit.permutation_indices)); + archive(cereal::make_nvp("R", fit.R)); + archive(cereal::make_nvp("P", fit.P)); if (version > 1) { archive(cereal::make_nvp("numerical_rank", fit.numerical_rank)); } else { @@ -53,9 +53,9 @@ inline void serialize(Archive &archive, Fit> &fit, template -void save(Archive &archive, - const GaussianProcessBase &gp, - const std::uint32_t) { +inline void save(Archive &archive, + const GaussianProcessBase &gp, + const std::uint32_t) { archive(cereal::make_nvp("name", gp.get_name())); archive(cereal::make_nvp("params", gp.get_params())); archive(cereal::make_nvp("insights", gp.insights)); @@ -63,9 +63,9 @@ void save(Archive &archive, template -void load(Archive &archive, - GaussianProcessBase &gp, - const std::uint32_t version) { +inline void load(Archive &archive, + GaussianProcessBase &gp, + const std::uint32_t version) { if (version > 0) { std::string model_name; archive(cereal::make_nvp("name", model_name)); diff --git a/include/albatross/src/core/declarations.hpp b/include/albatross/src/core/declarations.hpp index 7cc8bf14..058ac3dd 100644 --- a/include/albatross/src/core/declarations.hpp +++ b/include/albatross/src/core/declarations.hpp @@ -21,6 +21,13 @@ template class variant; using mapbox::util::variant; +/* + * Permutations + */ +namespace Eigen { +using PermutationMatrixX = PermutationMatrix; +} + namespace albatross { /* diff --git a/include/albatross/src/linalg/qr_utils.hpp b/include/albatross/src/linalg/qr_utils.hpp index 6bbb9dac..0770042e 100644 --- a/include/albatross/src/linalg/qr_utils.hpp +++ b/include/albatross/src/linalg/qr_utils.hpp @@ -25,21 +25,23 @@ get_R(const Eigen::ColPivHouseholderQR &qr) { .template triangularView(); } +inline Eigen::PermutationMatrixX +get_P(const Eigen::ColPivHouseholderQR &qr) { + return Eigen::PermutationMatrixX( + qr.colsPermutation().indices().template cast()); +} + /* * Computes R^-T P^T rhs given R and P from a QR decomposition. */ -template +template inline Eigen::MatrixXd sqrt_solve(const Eigen::MatrixXd &R, - const PermutationIndicesType &permutation_indices, + const Eigen::PermutationMatrix &P, const MatrixType &rhs) { - - Eigen::MatrixXd sqrt(rhs.rows(), rhs.cols()); - for (Eigen::Index i = 0; i < permutation_indices.size(); ++i) { - sqrt.row(i) = rhs.row(permutation_indices.coeff(i)); - } - sqrt = R.template triangularView().transpose().solve(sqrt); - return sqrt; + return R.template triangularView().transpose().solve( + P.transpose() * rhs); } template @@ -47,7 +49,7 @@ inline Eigen::MatrixXd sqrt_solve(const Eigen::ColPivHouseholderQR &qr, const MatrixType &rhs) { const Eigen::MatrixXd R = get_R(qr); - return sqrt_solve(R, qr.colsPermutation().indices(), rhs); + return sqrt_solve(R, qr.colsPermutation(), rhs); } } // namespace albatross diff --git a/include/albatross/src/linalg/spqr_utils.hpp b/include/albatross/src/linalg/spqr_utils.hpp index b98e07bd..73aafdb4 100644 --- a/include/albatross/src/linalg/spqr_utils.hpp +++ b/include/albatross/src/linalg/spqr_utils.hpp @@ -19,19 +19,20 @@ using SparseMatrix = Eigen::SparseMatrix; using SPQR = Eigen::SPQR; -using SparsePermutationMatrix = - Eigen::PermutationMatrix; - inline Eigen::MatrixXd get_R(const SPQR &qr) { return qr.matrixR() .topLeftCorner(qr.cols(), qr.cols()) .template triangularView(); } +inline Eigen::PermutationMatrixX get_P(const SPQR &qr) { + return Eigen::PermutationMatrixX( + qr.colsPermutation().indices().template cast()); +} + template inline Eigen::MatrixXd sqrt_solve(const SPQR &qr, const MatrixType &rhs) { - return sqrt_solve(get_R(qr), qr.colsPermutation().indices(), rhs); + return sqrt_solve(get_R(qr), get_P(qr), rhs); } // Matrices with any dimension smaller than this will use a special diff --git a/include/albatross/src/models/sparse_gp.hpp b/include/albatross/src/models/sparse_gp.hpp index d83c969d..36234e4e 100644 --- a/include/albatross/src/models/sparse_gp.hpp +++ b/include/albatross/src/models/sparse_gp.hpp @@ -97,8 +97,8 @@ template struct Fit> { std::vector train_features; Eigen::SerializableLDLT train_covariance; - Eigen::MatrixXd sigma_R; - PermutationIndices permutation_indices; + Eigen::MatrixXd R; + Eigen::PermutationMatrixX P; Eigen::VectorXd information; Eigen::Index numerical_rank; @@ -106,12 +106,10 @@ template struct Fit> { Fit(const std::vector &features_, const Eigen::SerializableLDLT &train_covariance_, - const Eigen::MatrixXd &sigma_R_, - PermutationIndices &&permutation_indices_, + const Eigen::MatrixXd &R_, const Eigen::PermutationMatrixX &P_, const Eigen::VectorXd &information_, Eigen::Index numerical_rank_) - : train_features(features_), train_covariance(train_covariance_), - sigma_R(sigma_R_), permutation_indices(std::move(permutation_indices_)), - information(information_), numerical_rank(numerical_rank_) {} + : train_features(features_), train_covariance(train_covariance_), R(R_), + P(P_), information(information_), numerical_rank(numerical_rank_) {} void shift_mean(const Eigen::VectorXd &mean_shift) { ALBATROSS_ASSERT(mean_shift.size() == information.size()); @@ -120,9 +118,8 @@ template struct Fit> { bool operator==(const Fit> &other) const { return (train_features == other.train_features && - train_covariance == other.train_covariance && - sigma_R == other.sigma_R && - permutation_indices == other.permutation_indices && + train_covariance == other.train_covariance && R == other.R && + P.indices() == other.P.indices() && information == other.information && numerical_rank == other.numerical_rank); } @@ -325,9 +322,9 @@ class SparseGaussianProcessRegression compute_internal_components(old_fit.train_features, features, targets, &A_ldlt, &K_uu_ldlt, &K_fu, &y); - const Eigen::Index n_old = old_fit.sigma_R.rows(); + const Eigen::Index n_old = old_fit.R.rows(); const Eigen::Index n_new = A_ldlt.rows(); - const Eigen::Index k = old_fit.sigma_R.cols(); + const Eigen::Index k = old_fit.R.cols(); Eigen::MatrixXd B = Eigen::MatrixXd::Zero(n_old + n_new, k); ALBATROSS_ASSERT(n_old == k); @@ -335,10 +332,7 @@ class SparseGaussianProcessRegression // Form: // B = |R_old P_old^T| = |Q_1| R P^T // |A^{-1/2} K_fu| |Q_2| - for (Eigen::Index i = 0; i < old_fit.permutation_indices.size(); ++i) { - const Eigen::Index &pi = old_fit.permutation_indices.coeff(i); - B.col(pi).topRows(i + 1) = old_fit.sigma_R.col(i).topRows(i + 1); - } + B.topRows(old_fit.P.rows()) = old_fit.R * old_fit.P.transpose(); B.bottomRows(n_new) = A_ldlt.sqrt_solve(K_fu); const auto B_qr = QRImplementation::compute(B, Base::threads_.get()); @@ -347,13 +341,9 @@ class SparseGaussianProcessRegression // |A^{-1/2} y | ALBATROSS_ASSERT(old_fit.information.size() == n_old); Eigen::VectorXd y_augmented(n_old + n_new); - for (Eigen::Index i = 0; i < old_fit.permutation_indices.size(); ++i) { - y_augmented[i] = - old_fit.information[old_fit.permutation_indices.coeff(i)]; - } y_augmented.topRows(n_old) = - old_fit.sigma_R.template triangularView() * - y_augmented.topRows(n_old); + old_fit.R.template triangularView() * + (old_fit.P.transpose() * old_fit.information); y_augmented.bottomRows(n_new) = A_ldlt.sqrt_solve(y, Base::threads_.get()); const Eigen::VectorXd v = B_qr->solve(y_augmented); @@ -365,10 +355,9 @@ class SparseGaussianProcessRegression Eigen::VectorXd::Constant(B_qr->cols(), details::cSparseRNugget); } using FitType = Fit>; - return FitType( - old_fit.train_features, old_fit.train_covariance, R, - B_qr->colsPermutation().indices().template cast(), v, - B_qr->rank()); + + return FitType(old_fit.train_features, old_fit.train_covariance, R, + get_P(*B_qr), v, B_qr->rank()); } // Here we create the QR decomposition of: @@ -415,10 +404,7 @@ class SparseGaussianProcessRegression using InducingPointFeatureType = typename std::decay::type; using FitType = Fit>; - return FitType( - u, K_uu_ldlt, get_R(*B_qr), - B_qr->colsPermutation().indices().template cast(), v, - B_qr->rank()); + return FitType(u, K_uu_ldlt, get_R(*B_qr), get_P(*B_qr), v, B_qr->rank()); } template @@ -471,9 +457,8 @@ class SparseGaussianProcessRegression const Eigen::MatrixXd sigma_inv_sqrt = C_ldlt.sqrt_solve(K_zz); const auto B_qr = QRImplementation::compute(sigma_inv_sqrt, nullptr); - new_fit.permutation_indices = - B_qr->colsPermutation().indices().template cast(); - new_fit.sigma_R = get_R(*B_qr); + new_fit.P = get_P(*B_qr); + new_fit.R = get_R(*B_qr); new_fit.numerical_rank = B_qr->rank(); return output; @@ -519,8 +504,8 @@ class SparseGaussianProcessRegression Q_sqrt.cwiseProduct(Q_sqrt).array().colwise().sum(); marginal_variance -= Q_diag; - const Eigen::MatrixXd S_sqrt = sqrt_solve( - sparse_gp_fit.sigma_R, sparse_gp_fit.permutation_indices, cross_cov); + const Eigen::MatrixXd S_sqrt = + sqrt_solve(sparse_gp_fit.R, sparse_gp_fit.P, cross_cov); const Eigen::VectorXd S_diag = S_sqrt.cwiseProduct(S_sqrt).array().colwise().sum(); marginal_variance += S_diag; @@ -537,8 +522,8 @@ class SparseGaussianProcessRegression this->covariance_function_(sparse_gp_fit.train_features, features); const Eigen::MatrixXd prior_cov = this->covariance_function_(features); - const Eigen::MatrixXd S_sqrt = sqrt_solve( - sparse_gp_fit.sigma_R, sparse_gp_fit.permutation_indices, cross_cov); + const Eigen::MatrixXd S_sqrt = + sqrt_solve(sparse_gp_fit.R, sparse_gp_fit.P, cross_cov); const Eigen::MatrixXd Q_sqrt = sparse_gp_fit.train_covariance.sqrt_solve(cross_cov); diff --git a/tests/test_sparse_gp.cc b/tests/test_sparse_gp.cc index 693fc013..1dbdc539 100644 --- a/tests/test_sparse_gp.cc +++ b/tests/test_sparse_gp.cc @@ -322,10 +322,10 @@ TYPED_TEST(SparseGaussianProcessTest, test_update) { (updated_in_place_pred.covariance - full_pred.covariance).norm(); auto compute_sigma = [](const auto &fit_model) -> Eigen::MatrixXd { - const Eigen::Index n = fit_model.get_fit().sigma_R.cols(); - Eigen::MatrixXd sigma = sqrt_solve(fit_model.get_fit().sigma_R, - fit_model.get_fit().permutation_indices, - Eigen::MatrixXd::Identity(n, n)); + const Eigen::Index n = fit_model.get_fit().R.cols(); + Eigen::MatrixXd sigma = + sqrt_solve(fit_model.get_fit().R, fit_model.get_fit().P, + Eigen::MatrixXd::Identity(n, n)); return sigma.transpose() * sigma; };