Allow faster predict methods.

albatross/core/distribution.h

@@ -85,10 +85,18 @@ template <typename CovarianceType> struct Distribution {
   }
 };
 
+// A JointDistribution has a dense covariance matrix, which
+// contains the covariance between each variable and all others.
+using JointDistribution = Distribution<Eigen::MatrixXd>;
+
+// We use a wrapper around DiagonalMatrix in order to make
+// the resulting distribution serializable
 using DiagonalMatrixXd =
     Eigen::SerializableDiagonalMatrix<double, Eigen::Dynamic>;
-using DenseDistribution = Distribution<Eigen::MatrixXd>;
-using DiagonalDistribution = Distribution<DiagonalMatrixXd>;
+// A MarginalDistribution has only a digaonal covariance
+// matrix, so in turn only describes the variance of each
+// variable independent of all others.
+using MarginalDistribution = Distribution<DiagonalMatrixXd>;
 
 template <typename CovarianceType, typename SizeType>
 Distribution<CovarianceType> subset(const std::vector<SizeType> &indices,

albatross/core/model.h

@@ -24,9 +24,6 @@
 
 namespace albatross {
 
-using TargetDistribution = DiagonalDistribution;
-using PredictDistribution = DenseDistribution;
-
 /*
  * A RegressionDataset holds two vectors of data, the features
  * where a single feature can be any class that contains the information used
@@ -36,12 +33,12 @@ using PredictDistribution = DenseDistribution;
  */
 template <typename FeatureType> struct RegressionDataset {
   std::vector<FeatureType> features;
-  TargetDistribution targets;
+  MarginalDistribution targets;
 
   RegressionDataset(){};
 
   RegressionDataset(const std::vector<FeatureType> &features_,
-                    const TargetDistribution &targets_)
+                    const MarginalDistribution &targets_)
       : features(features_), targets(targets_) {
     // If the two inputs aren't the same size they clearly aren't
     // consistent.
@@ -51,7 +48,7 @@ template <typename FeatureType> struct RegressionDataset {
 
   RegressionDataset(const std::vector<FeatureType> &features_,
                     const Eigen::VectorXd &targets_)
-      : RegressionDataset(features_, TargetDistribution(targets_)) {}
+      : RegressionDataset(features_, MarginalDistribution(targets_)) {}
 
   template <class Archive>
   typename std::enable_if<valid_in_out_serializer<FeatureType, Archive>::value,
@@ -127,7 +124,7 @@ class RegressionModel : public ParameterHandlingMixin {
    * predict.
    */
   void fit(const std::vector<FeatureType> &features,
-           const TargetDistribution &targets) {
+           const MarginalDistribution &targets) {
     assert(features.size() > 0);
     assert(static_cast<s32>(features.size()) ==
            static_cast<s32>(targets.size()));
@@ -140,7 +137,7 @@ class RegressionModel : public ParameterHandlingMixin {
    */
   void fit(const std::vector<FeatureType> &features,
            const Eigen::VectorXd &targets) {
-    fit(features, TargetDistribution(targets));
+    fit(features, MarginalDistribution(targets));
   }
 
   /*
@@ -155,28 +152,49 @@ class RegressionModel : public ParameterHandlingMixin {
    * and makes simple checks to confirm the implementation is returning
    * properly sized Distribution.
    */
-  PredictDistribution predict(const std::vector<FeatureType> &features) const {
+  JointDistribution predict(const std::vector<FeatureType> &features) const {
     assert(has_been_fit());
-    PredictDistribution preds = predict_(features);
+    JointDistribution preds = predict_(features);
     assert(static_cast<s32>(preds.mean.size()) ==
            static_cast<s32>(features.size()));
     return preds;
   }
 
-  PredictDistribution predict(const FeatureType &feature) const {
+  JointDistribution predict(const FeatureType &feature) const {
     std::vector<FeatureType> features = {feature};
     return predict(features);
   }
 
+  MarginalDistribution
+  predict_marginal(const std::vector<FeatureType> &features) const {
+    assert(has_been_fit());
+    MarginalDistribution 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,
    * follwed by predict but overriding this method may speed up computation for
    * some models.
    */
-  PredictDistribution
+  JointDistribution
   fit_and_predict(const std::vector<FeatureType> &train_features,
-                  const TargetDistribution &train_targets,
+                  const MarginalDistribution &train_targets,
                   const std::vector<FeatureType> &test_features) {
     // Fit using the training data, then predict with the test.
     fit(train_features, train_targets);
@@ -187,7 +205,7 @@ class RegressionModel : public ParameterHandlingMixin {
    * A convenience wrapper around fit_and_predict which uses the entries
    * in a RegressionFold struct
    */
-  PredictDistribution fit_and_predict(const RegressionFold<FeatureType> &fold) {
+  JointDistribution fit_and_predict(const RegressionFold<FeatureType> &fold) {
     return fit_and_predict(fold.train.features, fold.train.targets,
                            fold.test.features);
   }
@@ -228,11 +246,29 @@ class RegressionModel : public ParameterHandlingMixin {
 
 protected:
   virtual void fit_(const std::vector<FeatureType> &features,
-                    const TargetDistribution &targets) = 0;
+                    const MarginalDistribution &targets) = 0;
 
-  virtual PredictDistribution
+  virtual JointDistribution
   predict_(const std::vector<FeatureType> &features) const = 0;
 
+  virtual MarginalDistribution
+  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 MarginalDistribution(
+        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

@@ -119,7 +119,7 @@ class AdaptedRegressionModel
 
 protected:
   void fit_(const std::vector<FeatureType> &features,
-            const TargetDistribution &targets) override {
+            const MarginalDistribution &targets) override {
     this->sub_model_.fit(convert_features(features), targets);
   }
 
@@ -130,18 +130,28 @@ class AdaptedRegressionModel
    */
   fit_type_if_serializable<RegressionModelImplementation>
   serializable_fit_(const std::vector<FeatureType> &features,
-                    const TargetDistribution &targets) const override {
+                    const MarginalDistribution &targets) const override {
     assert(false &&
            "serializable_fit_ for an adapted model should never be called");
     typename fit_type_or_void<RegressionModelImplementation>::type dummy;
     return dummy;
   }
 
-  PredictDistribution
+  JointDistribution
   predict_(const std::vector<FeatureType> &features) const override {
     return sub_model_.predict(convert_features(features));
   }
 
+  virtual MarginalDistribution
+  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/core/serialize.h

@@ -89,13 +89,13 @@ class SerializableRegressionModel : public RegressionModel<FeatureType> {
 
 protected:
   void fit_(const std::vector<FeatureType> &features,
-            const TargetDistribution &targets) {
+            const MarginalDistribution &targets) {
     model_fit_ = serializable_fit_(features, targets);
   }
 
   virtual ModelFit
   serializable_fit_(const std::vector<FeatureType> &features,
-                    const TargetDistribution &targets) const = 0;
+                    const MarginalDistribution &targets) const = 0;
 
   ModelFit model_fit_;
 };

albatross/crossvalidation.h

@@ -26,7 +26,7 @@ namespace albatross {
  * the quality of the prediction.
  */
 using EvaluationMetric = std::function<double(
-    const PredictDistribution &prediction, const TargetDistribution &targets)>;
+    const JointDistribution &prediction, const MarginalDistribution &targets)>;
 
 inline FoldIndices get_train_indices(const FoldIndices &test_indices,
                                      const int n) {
@@ -68,11 +68,11 @@ folds_from_fold_indexer(const RegressionDataset<FeatureType> &dataset,
 
     std::vector<FeatureType> train_features =
         subset(train_indices, dataset.features);
-    TargetDistribution train_targets = subset(train_indices, dataset.targets);
+    MarginalDistribution train_targets = subset(train_indices, dataset.targets);
 
     std::vector<FeatureType> test_features =
         subset(test_indices, dataset.features);
-    TargetDistribution test_targets = subset(test_indices, dataset.targets);
+    MarginalDistribution test_targets = subset(test_indices, dataset.targets);
 
     assert(train_features.size() == train_targets.size());
     assert(test_features.size() == test_targets.size());
@@ -151,17 +151,17 @@ static inline std::vector<RegressionFold<FeatureType>> leave_one_group_out(
 }
 
 /*
- * Computes a PredictDistribution for each fold in set of cross validation
- * folds.  The resulting vector of PredictDistributions can then be used
+ * Computes a JointDistribution for each fold in set of cross validation
+ * folds.  The resulting vector of JointDistributions can then be used
  * for things like computing an EvaluationMetric for each fold, or assembling
  * all the predictions into a single cross validated PredictionDistribution.
  */
 template <typename FeatureType>
-static inline std::vector<PredictDistribution> cross_validated_predictions(
+static inline std::vector<JointDistribution> cross_validated_predictions(
     const std::vector<RegressionFold<FeatureType>> &folds,
     RegressionModel<FeatureType> *model) {
   // Iteratively make predictions and assemble the output vector
-  std::vector<PredictDistribution> predictions;
+  std::vector<JointDistribution> predictions;
   for (std::size_t i = 0; i < folds.size(); i++) {
     predictions.push_back(model->fit_and_predict(
         folds[i].train_dataset.features, folds[i].train_dataset.targets,
@@ -178,7 +178,7 @@ template <class FeatureType>
 static inline Eigen::VectorXd
 compute_scores(const EvaluationMetric &metric,
                const std::vector<RegressionFold<FeatureType>> &folds,
-               const std::vector<PredictDistribution> &predictions) {
+               const std::vector<JointDistribution> &predictions) {
   // Create a vector of metrics, one for each fold.
   Eigen::VectorXd metrics(static_cast<s32>(folds.size()));
   // Loop over each fold, making predictions then evaluating them
@@ -200,7 +200,7 @@ cross_validated_scores(const EvaluationMetric &metric,
                        const std::vector<RegressionFold<FeatureType>> &folds,
                        RegressionModel<FeatureType> *model) {
   // Create a vector of predictions.
-  std::vector<PredictDistribution> predictions =
+  std::vector<JointDistribution> predictions =
       cross_validated_predictions<FeatureType>(folds, model);
   return compute_scores(metric, folds, predictions);
 }
@@ -216,13 +216,13 @@ cross_validated_scores(const EvaluationMetric &metric,
  * unknown.
  */
 template <typename FeatureType>
-static inline PredictDistribution
+static inline JointDistribution
 cross_validated_predict(const std::vector<RegressionFold<FeatureType>> &folds,
                         RegressionModel<FeatureType> *model) {
   // Get the cross validated predictions, note however that
   // depending on the type of folds, these predictions may
   // be shuffled.
-  const std::vector<PredictDistribution> predictions =
+  const std::vector<JointDistribution> predictions =
       cross_validated_predictions<FeatureType>(folds, model);
   // Create a new prediction mean that will eventually contain
   // the ordered concatenation of each fold's predictions.
@@ -240,7 +240,7 @@ cross_validated_predict(const std::vector<RegressionFold<FeatureType>> &folds,
           pred.mean[i];
     }
   }
-  return PredictDistribution(mean);
+  return JointDistribution(mean);
 }
 
 } // namespace albatross

albatross/evaluate.h

@@ -105,9 +105,8 @@ negative_log_likelihood(const Eigen::VectorXd &deviation,
  */
 namespace evaluation_metrics {
 
-static inline double
-root_mean_square_error(const PredictDistribution &prediction,
-                       const TargetDistribution &truth) {
+static inline double root_mean_square_error(const JointDistribution &prediction,
+                                            const MarginalDistribution &truth) {
   const Eigen::VectorXd error = prediction.mean - truth.mean;
   double mse = error.dot(error) / static_cast<double>(error.size());
   return sqrt(mse);
@@ -117,8 +116,8 @@ root_mean_square_error(const PredictDistribution &prediction,
  * Takes output from a model (PredictionDistribution)
  * and the corresponding truth and uses them to compute the stddev.
  */
-static inline double standard_deviation(const PredictDistribution &prediction,
-                                        const TargetDistribution &truth) {
+static inline double standard_deviation(const JointDistribution &prediction,
+                                        const MarginalDistribution &truth) {
   Eigen::VectorXd error = prediction.mean - truth.mean;
   const auto n_elements = static_cast<double>(error.size());
   const double mean_error = error.sum() / n_elements;
@@ -131,8 +130,8 @@ static inline double standard_deviation(const PredictDistribution &prediction,
  * distribution is multivariate normal.
  */
 static inline double
-negative_log_likelihood(const PredictDistribution &prediction,
-                        const TargetDistribution &truth) {
+negative_log_likelihood(const JointDistribution &prediction,
+                        const MarginalDistribution &truth) {
   const Eigen::VectorXd mean = prediction.mean - truth.mean;
   Eigen::MatrixXd covariance(prediction.covariance);
   if (truth.has_covariance()) {