From d65ca3aff98c07d7422b617da64b248c4c47bffe Mon Sep 17 00:00:00 2001
From: Huangzizhou <zizhou@nyu.edu>
Date: Fri, 8 Dec 2023 14:01:15 -0500
Subject: [PATCH] add solver_info to PostStepData

---
 src/polysolve/nonlinear/PostStepData.cpp |  3 ++-
 src/polysolve/nonlinear/PostStepData.hpp |  2 ++
 src/polysolve/nonlinear/Solver.cpp       | 10 ++++------
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/polysolve/nonlinear/PostStepData.cpp b/src/polysolve/nonlinear/PostStepData.cpp
index a411dce..da33e43 100644
--- a/src/polysolve/nonlinear/PostStepData.cpp
+++ b/src/polysolve/nonlinear/PostStepData.cpp
@@ -3,9 +3,10 @@
 namespace polysolve::nonlinear
 {
     PostStepData::PostStepData(const int iter_num,
+                               const json &solver_info,
                                const Eigen::VectorXd &x,
                                const Eigen::VectorXd &grad)
-        : iter_num(iter_num), x(x), grad(grad)
+        : iter_num(iter_num), solver_info(solver_info), x(x), grad(grad)
     {
     }
 } // namespace polysolve::nonlinear
diff --git a/src/polysolve/nonlinear/PostStepData.hpp b/src/polysolve/nonlinear/PostStepData.hpp
index dc5ccdb..a7d155a 100644
--- a/src/polysolve/nonlinear/PostStepData.hpp
+++ b/src/polysolve/nonlinear/PostStepData.hpp
@@ -9,10 +9,12 @@ namespace polysolve::nonlinear
     {
     public:
         PostStepData(const int iter_num,
+                     const json &solver_info,
                      const Eigen::VectorXd &x,
                      const Eigen::VectorXd &grad);
 
         const int iter_num;
+        const json &solver_info;
         const Eigen::VectorXd &x;
         const Eigen::VectorXd &grad;
     };
diff --git a/src/polysolve/nonlinear/Solver.cpp b/src/polysolve/nonlinear/Solver.cpp
index aa2a632..7430e4d 100644
--- a/src/polysolve/nonlinear/Solver.cpp
+++ b/src/polysolve/nonlinear/Solver.cpp
@@ -213,8 +213,6 @@ namespace polysolve::nonlinear
             objFunc.solution_changed(x);
         }
 
-        objFunc.post_step(PostStepData(this->m_current.iterations, x, grad));
-
         const auto g_norm_tol = this->m_stop.gradNorm;
         this->m_stop.gradNorm = first_grad_norm_tol;
 
@@ -229,6 +227,7 @@ namespace polysolve::nonlinear
             this->m_stop.fDelta, this->m_stop.gradNorm, this->m_stop.xDelta);
 
         update_solver_info(objFunc.value(x));
+        objFunc.post_step(PostStepData(this->m_current.iterations, solver_info, x, grad));
 
         int f_delta_step_cnt = 0;
         double f_delta = 0;
@@ -386,6 +385,9 @@ namespace polysolve::nonlinear
             // -----------
             const double step = (rate * delta_x).norm();
 
+            update_solver_info(energy);
+            objFunc.post_step(PostStepData(this->m_current.iterations, solver_info, x, grad));
+
             if (objFunc.stop(x))
             {
                 this->m_status = cppoptlib::Status::UserDefined;
@@ -393,8 +395,6 @@ namespace polysolve::nonlinear
                 m_logger.debug("[{}][{}] Objective decided to stop", name(), m_line_search->name());
             }
 
-            objFunc.post_step(PostStepData(this->m_current.iterations, x, grad));
-
             if (f_delta < this->m_stop.fDelta)
                 f_delta_step_cnt++;
             else
@@ -411,8 +411,6 @@ namespace polysolve::nonlinear
             if (++this->m_current.iterations >= this->m_stop.iterations)
                 this->m_status = cppoptlib::Status::IterationLimit;
 
-            update_solver_info(energy);
-
             // reset the tolerance, since in the first iter it might be smaller
             this->m_stop.gradNorm = g_norm_tol;
         } while (objFunc.callback(this->m_current, x) && (this->m_status == cppoptlib::Status::Continue));