fix(nelder-mead): argsort before early exit (#984)

Nixtla · Feb 11, 2025 · 788d170 · 788d170
1 parent 52a5c9a
commit 788d170
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 36 deletions.
diff --git a/include/statsforecast/nelder_mead.h b/include/statsforecast/nelder_mead.h
@@ -7,6 +7,7 @@ namespace nm {
 using Eigen::VectorXd;
 using RowMajorMatrixXd =
     Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
+using OptimResult = std::tuple<VectorXd, double, int, RowMajorMatrixXd>;
 
 inline auto Clamp(const VectorXd &x, const VectorXd &lower,
                   const VectorXd &upper) {
@@ -26,11 +27,11 @@ inline Eigen::VectorX<Eigen::Index> ArgSort(const VectorXd &v) {
 }
 
 template <typename Func, typename... Args>
-std::tuple<VectorXd, double, int>
-NelderMead(Func F, const VectorXd &x, const VectorXd &lower,
-           const VectorXd upper, double init_step, double zero_pert,
-           double alpha, double gamma, double rho, double sigma, int max_iter,
-           double tol_std, bool adaptive, Args &&...args) {
+OptimResult NelderMead(Func F, const VectorXd &x, const VectorXd &lower,
+                       const VectorXd upper, double init_step, double zero_pert,
+                       double alpha, double gamma, double rho, double sigma,
+                       int max_iter, double tol_std, bool adaptive,
+                       Args &&...args) {
   auto x0 = Clamp(x, lower, upper);
   auto n = x0.size();
   if (adaptive) {
@@ -57,17 +58,17 @@ NelderMead(Func F, const VectorXd &x, const VectorXd &lower,
   int i;
   Eigen::Index best_idx = 0;
   for (i = 0; i < max_iter; ++i) {
-    // check whether method should stop
-    if (StandardDeviation(f_simplex) < tol_std) {
-      break;
-    }
-
     // Step1: order of f_simplex
     Eigen::VectorX<Eigen::Index> order_f = ArgSort(f_simplex);
     best_idx = order_f(0);
     auto worst_idx = order_f(n);
     auto second_worst_idx = order_f(n - 1);
 
+    // check whether method should stop
+    if (StandardDeviation(f_simplex) < tol_std) {
+      break;
+    }
+
     // calculate centroid as the col means removing the row with the max fval
     VectorXd x_o = (simplex.colwise().sum() - simplex.row(worst_idx)) / n;
 
@@ -131,6 +132,6 @@ NelderMead(Func F, const VectorXd &x, const VectorXd &lower,
       f_simplex(j) = F(simplex.row(j), std::forward<Args>(args)...);
     }
   }
-  return {simplex.row(best_idx), f_simplex(best_idx), i + 1};
+  return {simplex.row(best_idx), f_simplex(best_idx), i + 1, simplex};
 }
 } // namespace nm
diff --git a/nbs/src/ets.ipynb b/nbs/src/ets.ipynb
@@ -755,7 +755,7 @@
     "        1_000,\n",
     "        True,\n",
     "    )\n",
-    "    return results(*opt_res, None)"
+    "    return results(*opt_res)"
    ]
   },
   {

diff --git a/nbs/src/theta.ipynb b/nbs/src/theta.ipynb
@@ -479,7 +479,7 @@
     "        modeltype,\n",
     "        nmse\n",
     "    )\n",
-    "    return results(*opt_res, None)  # type: ignore"
+    "    return results(*opt_res)"
    ]
   },
   {

diff --git a/python/statsforecast/ets.py b/python/statsforecast/ets.py
@@ -501,7 +501,7 @@ def optimize_ets_target_fn(
         1_000,
         True,
     )
-    return results(*opt_res, None)
+    return results(*opt_res)
 
 # %% ../../nbs/src/ets.ipynb 26
 def etsmodel(

diff --git a/python/statsforecast/theta.py b/python/statsforecast/theta.py
@@ -127,7 +127,7 @@ def optimize_theta_target_fn(
         modeltype,
         nmse,
     )
-    return results(*opt_res, None)  # type: ignore
+    return results(*opt_res)
 
 # %% ../../nbs/src/theta.ipynb 14
 def thetamodel(

diff --git a/src/ets.cpp b/src/ets.cpp
@@ -279,15 +279,15 @@ double ObjectiveFunction(const VectorXd &params, const VectorXd &y, int n_state,
   }
   return obj_val;
 }
-std::tuple<VectorXd, double, int>
-Optimize(const Eigen::Ref<const VectorXd> &x0,
-         const Eigen::Ref<const VectorXd> &y, int n_state, Component error,
-         Component trend, Component season, Criterion opt_crit, int n_mse,
-         int m, bool opt_alpha, bool opt_beta, bool opt_gamma, bool opt_phi,
-         double alpha, double beta, double gamma, double phi,
-         const Eigen::Ref<const VectorXd> &lower,
-         const Eigen::Ref<const VectorXd> &upper, double tol_std, int max_iter,
-         bool adaptive) {
+nm::OptimResult Optimize(const Eigen::Ref<const VectorXd> &x0,
+                         const Eigen::Ref<const VectorXd> &y, int n_state,
+                         Component error, Component trend, Component season,
+                         Criterion opt_crit, int n_mse, int m, bool opt_alpha,
+                         bool opt_beta, bool opt_gamma, bool opt_phi,
+                         double alpha, double beta, double gamma, double phi,
+                         const Eigen::Ref<const VectorXd> &lower,
+                         const Eigen::Ref<const VectorXd> &upper,
+                         double tol_std, int max_iter, bool adaptive) {
   double init_step = 0.05;
   double nm_alpha = 1.0;
   double nm_gamma = 2.0;

diff --git a/src/theta.cpp b/src/theta.cpp
@@ -120,10 +120,10 @@ pegels_resid(const Eigen::Ref<const VectorXd> &y, ModelType model_type,
   return {amse, e, states, mse};
 }
 
-double theta_target_fn(const VectorXd &params, double init_level,
-                       double init_alpha, double init_theta, bool opt_level,
-                       bool opt_alpha, bool opt_theta, const VectorXd &y,
-                       ModelType model_type, size_t nmse) {
+double target_fn(const VectorXd &params, double init_level, double init_alpha,
+                 double init_theta, bool opt_level, bool opt_alpha,
+                 bool opt_theta, const VectorXd &y, ModelType model_type,
+                 size_t nmse) {
   RowMajorMatrixXd states = RowMajorMatrixXd::Zero(y.size(), 5);
   size_t j = 0;
   double level, alpha, theta;
@@ -152,13 +152,13 @@ double theta_target_fn(const VectorXd &params, double init_level,
   return mse;
 }
 
-std::tuple<VectorXd, double, int>
-optimize(const Eigen::Ref<const VectorXd> &x0,
-         const Eigen::Ref<const VectorXd> &lower,
-         const Eigen::Ref<const VectorXd> &upper, double init_level,
-         double init_alpha, double init_theta, bool opt_level, bool opt_alpha,
-         bool opt_theta, const Eigen::Ref<const VectorXd> &y,
-         ModelType model_type, size_t nmse) {
+nm::OptimResult optimize(const Eigen::Ref<const VectorXd> &x0,
+                         const Eigen::Ref<const VectorXd> &lower,
+                         const Eigen::Ref<const VectorXd> &upper,
+                         double init_level, double init_alpha,
+                         double init_theta, bool opt_level, bool opt_alpha,
+                         bool opt_theta, const Eigen::Ref<const VectorXd> &y,
+                         ModelType model_type, size_t nmse) {
   double init_step = 0.05;
   double zero_pert = 1e-4;
   double alpha = 1.0;
@@ -168,7 +168,7 @@ optimize(const Eigen::Ref<const VectorXd> &x0,
   int max_iter = 1'000;
   double tol_std = 1e-4;
   bool adaptive = true;
-  return nm::NelderMead(theta_target_fn, x0, lower, upper, init_step, zero_pert,
+  return nm::NelderMead(target_fn, x0, lower, upper, init_step, zero_pert,
                         alpha, gamma, rho, sigma, max_iter, tol_std, adaptive,
                         init_level, init_alpha, init_theta, opt_level,
                         opt_alpha, opt_theta, y, model_type, nmse);