Minor cleaning in models

Author: costashatz

include/blackdrops/gp_model.hpp

@@ -147,20 +147,11 @@ namespace blackdrops {
             }
         }
 
-        std::tuple<Eigen::VectorXd, double> predict(const Eigen::VectorXd& x) const
+        std::tuple<Eigen::VectorXd, Eigen::VectorXd> predict(const Eigen::VectorXd& x, bool compute_variance = true) const
         {
-            Eigen::VectorXd ms;
-            Eigen::VectorXd ss;
-            std::tie(ms, ss) = predictm(x);
-            return std::make_tuple(ms, ss.mean());
-        }
-
-        std::tuple<Eigen::VectorXd, Eigen::VectorXd> predictm(const Eigen::VectorXd& x) const
-        {
-            Eigen::VectorXd ms(_gp_model.dim_out());
-            Eigen::VectorXd ss(_gp_model.dim_out());
-
-            return _gp_model.query(x);
+            if (compute_variance)
+                return _gp_model.query(x);
+            return std::make_tuple(_gp_model.mu(x), Eigen::VectorXd::Zero(_gp_model.dim_out()));
         }
 
         Eigen::MatrixXd samples() const

include/blackdrops/mi_model.hpp

@@ -112,7 +112,7 @@ namespace blackdrops {
             }
         }
 
-        std::tuple<Eigen::VectorXd, Eigen::VectorXd> predictm(const Eigen::VectorXd& x) const
+        std::tuple<Eigen::VectorXd, Eigen::VectorXd> predict(const Eigen::VectorXd& x, bool) const
         {
             Eigen::VectorXd mu = _mean(x, x);
             Eigen::VectorXd ss = Eigen::VectorXd::Zero(mu.size());
@@ -142,6 +142,6 @@ namespace blackdrops {
             return limbo::opt::no_grad(-mse);
         }
     };
-}
+} // namespace blackdrops
 
 #endif

include/blackdrops/system/dart_system.hpp

@@ -189,7 +189,7 @@ namespace blackdrops {
 
                     Eigen::VectorXd mu;
                     Eigen::VectorXd sigma;
-                    std::tie(mu, sigma) = model.predictm(query_vec);
+                    std::tie(mu, sigma) = model.predict(query_vec);
 
                     Eigen::VectorXd final = init_diff + mu;
 
@@ -229,7 +229,7 @@ namespace blackdrops {
 
                     Eigen::VectorXd mu;
                     Eigen::VectorXd sigma;
-                    std::tie(mu, sigma) = model.predictm(query_vec);
+                    std::tie(mu, sigma) = model.predict(query_vec);
 
                     if (Params::blackdrops::stochastic()) {
                         sigma = sigma.array().sqrt();

include/blackdrops/system/ode_system.hpp

@@ -163,7 +163,7 @@ namespace blackdrops {
 
                     Eigen::VectorXd mu;
                     Eigen::VectorXd sigma;
-                    std::tie(mu, sigma) = model.predictm(query_vec);
+                    std::tie(mu, sigma) = model.predict(query_vec);
 
                     Eigen::VectorXd final = init_diff + mu;
 
@@ -203,7 +203,7 @@ namespace blackdrops {
 
                     Eigen::VectorXd mu;
                     Eigen::VectorXd sigma;
-                    std::tie(mu, sigma) = model.predictm(query_vec);
+                    std::tie(mu, sigma) = model.predict(query_vec, Params::blackdrops::stochastic());
 
                     if (Params::blackdrops::stochastic()) {
                         sigma = sigma.array().sqrt();