Add variants of predict which allow predicting only the mean

or only the diagonal variance which can often be done much
faster.

albatross/core/model.h

@@ -152,6 +152,27 @@ class RegressionModel : public ParameterHandlingMixin {
     return predict(features);
   }
 
+  DiagonalDistribution
+  predict_marginal(const std::vector<FeatureType> &features) const {
+    assert(has_been_fit());
+    DiagonalDistribution preds = predict_marginal_(features);
+    assert(static_cast<s32>(preds.mean.size()) ==
+           static_cast<s32>(features.size()));
+    return preds;
+  }
+
+  Eigen::VectorXd predict_mean(const std::vector<FeatureType> &features) const {
+    assert(has_been_fit());
+    Eigen::VectorXd preds = predict_mean_(features);
+    assert(static_cast<s32>(preds.size()) == static_cast<s32>(features.size()));
+    return preds;
+  }
+
+  double predict_mean(const FeatureType &feature) const {
+    std::vector<FeatureType> features = {feature};
+    return predict_mean(features)[0];
+  }
+
   /*
    * Computes predictions for the test features given set of training
    * features and targets. In the general case this is simply a call to fit,
@@ -217,6 +238,24 @@ class RegressionModel : public ParameterHandlingMixin {
   virtual PredictDistribution
   predict_(const std::vector<FeatureType> &features) const = 0;
 
+  virtual DiagonalDistribution
+  predict_marginal_(const std::vector<FeatureType> &features) const {
+    std::cout << "WARNING: A marginal prediction is being made, but in a "
+                 "horribly inefficient way.";
+    const auto full_distribution = predict_(features);
+    return DiagonalDistribution(
+        full_distribution.mean,
+        full_distribution.covariance.diagonal().asDiagonal());
+  }
+
+  virtual Eigen::VectorXd
+  predict_mean_(const std::vector<FeatureType> &features) const {
+    std::cout << "WARNING: A mean prediction is being made, but in a horribly "
+                 "inefficient way.";
+    const auto full_distribution = predict_(features);
+    return full_distribution.mean;
+  }
+
   bool has_been_fit_;
 };
 

albatross/core/model_adapter.h

@@ -142,6 +142,16 @@ class AdaptedRegressionModel
     return sub_model_.predict(convert_features(features));
   }
 
+  virtual DiagonalDistribution
+  predict_marginal_(const std::vector<FeatureType> &features) const override {
+    return sub_model_.predict_marginal(convert_features(features));
+  }
+
+  virtual Eigen::VectorXd
+  predict_mean_(const std::vector<FeatureType> &features) const override {
+    return sub_model_.predict_mean(convert_features(features));
+  }
+
   const std::vector<SubFeature>
   convert_features(const std::vector<FeatureType> &parent_features) const {
     std::vector<SubFeature> converted;

albatross/models/gp.h

@@ -156,6 +156,33 @@ class GaussianProcessRegression
     return PredictDistribution(pred, pred_cov);
   }
 
+  virtual DiagonalDistribution
+  predict_marginal_(const std::vector<FeatureType> &features) const override {
+    const auto cross_cov = asymmetric_covariance(
+        covariance_function_, features, this->model_fit_.train_features);
+    // Then we can use the information vector to determine the posterior
+    const Eigen::VectorXd pred = cross_cov * this->model_fit_.information;
+
+    auto ldlt = this->model_fit_.train_ldlt;
+    Eigen::MatrixXd explained = ldlt.solve(cross_cov.transpose());
+    Eigen::VectorXd marginal_variance =
+        -explained.cwiseProduct(cross_cov.transpose()).array().colwise().sum();
+    for (Eigen::Index i = 0; i < pred.size(); i++) {
+      marginal_variance[i] += covariance_function_(features[i], features[i]);
+    }
+
+    return DiagonalDistribution(pred, marginal_variance.asDiagonal());
+  }
+
+  virtual Eigen::VectorXd
+  predict_mean_(const std::vector<FeatureType> &features) const override {
+    const auto cross_cov = asymmetric_covariance(
+        covariance_function_, features, this->model_fit_.train_features);
+    // Then we can use the information vector to determine the posterior
+    const Eigen::VectorXd pred = cross_cov * this->model_fit_.information;
+    return pred;
+  }
+
 private:
   CovarianceFunction covariance_function_;
   std::string model_name_;