Merge branch 'main' into adam

polyfem · Dec 8, 2023 · cb8c59d · cb8c59d
2 parents d1e6d9b + e9bb718
commit cb8c59d
Show file tree

Hide file tree

Showing 12 changed files with 147 additions and 59 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -157,8 +157,6 @@ add_subdirectory(src/polysolve/nonlinear)
 target_link_libraries(polysolve_linear PUBLIC polysolve_coverage_config)
 target_link_libraries(polysolve PUBLIC polysolve_coverage_config)
 
-
-
 target_compile_features(polysolve_linear PUBLIC cxx_std_17)
 target_compile_features(polysolve PUBLIC cxx_std_17)
 
@@ -183,33 +181,11 @@ target_compile_definitions(polysolve_linear PUBLIC POLYSOLVE_JSON_SPEC_DIR="${PR
 # Dependencies
 ################################################################################
 
-# polysolve_linear
-target_link_libraries(polysolve PUBLIC polysolve::linear)
-
-# CppNumericalSolvers
-include(cppoptlib)
-target_link_libraries(polysolve PUBLIC cppoptlib)
-
-# LBFGSpp
-include(LBFGSpp)
-target_link_libraries(polysolve PUBLIC LBFGSpp::LBFGSpp)
-
-
-
-
-# JSON Specification Engine library
-include(jse)
-target_link_libraries(polysolve_linear PUBLIC jse::jse)
-
-# spdlog
-include(spdlog)
-target_link_libraries(polysolve_linear PUBLIC spdlog::spdlog)
-
-# finite-diff
-include(finite-diff)
-target_link_libraries(polysolve PUBLIC finitediff::finitediff)
+# ------
+# Linear
+# ------
 
-# Accelerate solver
+# Accelerate solver (Include before Eigen)
 if(POLYSOLVE_WITH_ACCELERATE)
     include(CPM)
     CPMAddPackage(
@@ -222,23 +198,24 @@ if(POLYSOLVE_WITH_ACCELERATE)
     find_package(BLAS REQUIRED)
     find_package(LAPACK REQUIRED)
     target_link_libraries(polysolve_linear PRIVATE BLAS::BLAS LAPACK::LAPACK)
-endif()
-
-# Extra warnings
-include(polysolve_warnings)
-target_link_libraries(polysolve_linear PRIVATE polysolve::warnings)
-target_link_libraries(polysolve PRIVATE polysolve::warnings)
-
-# Sanitizers
-if(POLYSOLVE_WITH_SANITIZERS)
-    include(sanitizers)
-    add_sanitizers(polysolve_linear)
-    add_sanitizers(polysolve)
+    target_compile_definitions(polysolve_linear PRIVATE -DPOLYSOLVE_WITH_ACCELERATE)
 endif()
 
 include(eigen)
 target_link_libraries(polysolve_linear PUBLIC Eigen3::Eigen)
 
+# spdlog
+include(spdlog)
+target_link_libraries(polysolve_linear PUBLIC spdlog::spdlog)
+
+# JSON (MIT)
+include(json)
+target_link_libraries(polysolve_linear PUBLIC nlohmann_json::nlohmann_json)
+
+# JSON Specification Engine library
+include(jse)
+target_link_libraries(polysolve_linear PUBLIC jse::jse)
+
 # Hypre (GNU Lesser General Public License)
 if(POLYSOLVE_WITH_HYPRE)
     include(hypre)
@@ -249,15 +226,6 @@ if(POLYSOLVE_WITH_HYPRE)
     endif()
 endif()
 
-# Json (MIT)
-include(json)
-target_link_libraries(polysolve_linear PUBLIC nlohmann_json::nlohmann_json)
-
-# Accelerate solvers
-if(POLYSOLVE_WITH_ACCELERATE)
-    target_compile_definitions(polysolve_linear PRIVATE -DPOLYSOLVE_WITH_ACCELERATE)
-endif()
-
 # CHOLMOD solver
 if(POLYSOLVE_WITH_CHOLMOD)
     include(suitesparse)
@@ -326,6 +294,45 @@ if(POLYSOLVE_WITH_CUSOLVER)
     endif()
 endif()
 
+# Sanitizers
+if(POLYSOLVE_WITH_SANITIZERS)
+    include(sanitizers)
+    add_sanitizers(polysolve_linear)
+endif()
+
+# Extra warnings (include last for highest priority)
+include(polysolve_warnings)
+target_link_libraries(polysolve_linear PRIVATE polysolve::warnings)
+
+# ---------
+# Nonlinear
+# ---------
+
+# polysolve::linear
+target_link_libraries(polysolve PUBLIC polysolve::linear)
+
+# CppNumericalSolvers
+include(cppoptlib)
+target_link_libraries(polysolve PUBLIC cppoptlib)
+
+# LBFGSpp
+include(LBFGSpp)
+target_link_libraries(polysolve PUBLIC LBFGSpp::LBFGSpp)
+
+# finite-diff (include this after eigen)
+include(finite-diff)
+target_link_libraries(polysolve PUBLIC finitediff::finitediff)
+
+# Sanitizers
+if(POLYSOLVE_WITH_SANITIZERS)
+    include(sanitizers)
+    add_sanitizers(polysolve)
+endif()
+
+# Extra warnings (include last for highest priority)
+include(polysolve_warnings)
+target_link_libraries(polysolve PRIVATE polysolve::warnings)
+
 
 ################################################################################
 # Compiler options

diff --git a/cmake/recipes/jse.cmake b/cmake/recipes/jse.cmake
@@ -8,4 +8,4 @@ endif()
 message(STATUS "Third-party: creating target 'jse::jse'")
 
 include(CPM)
-CPMAddPackage("gh:geometryprocessing/json-spec-engine#1261dc89478c7646ff99cbed8bc5357c2813565d")
+CPMAddPackage("gh:geometryprocessing/json-spec-engine#49f1a30f8c2912814916ec3d6108a649b23cb243")
diff --git a/non-linear-solver-spec.json b/non-linear-solver-spec.json
@@ -16,6 +16,7 @@
             "LBFGSB",
             "Newton",
             "ADAM",
+            "StochasticGradientDescent",
             "box_constraints",
             "advanced"
         ],
@@ -30,6 +31,7 @@
             "DenseNewton",
             "GradientDescent",
             "ADAM",
+            "StochasticGradientDescent",
             "L-BFGS",
             "BFGS",
             "L-BFGS-B",
@@ -202,6 +204,21 @@
         "type": "float",
         "doc": "Parameter epsilon for ADAM."
     },
+    {
+        "pointer": "/StochasticGradientDescent",
+        "default": null,
+        "type": "object",
+        "optional": [
+            "erase_component_probability"
+        ],
+        "doc": "Options for Stochastic Gradient Descent."
+    },
+    {
+        "pointer": "/StochasticGradientDescent/erase_component_probability",
+        "default": 0.3,
+        "type": "float",
+        "doc": "Probability of erasing a component on the gradient for StochasticGradientDescent."
+    },
     {
         "pointer": "/line_search",
         "default": null,

diff --git a/src/polysolve/nonlinear/CMakeLists.txt b/src/polysolve/nonlinear/CMakeLists.txt
@@ -5,6 +5,8 @@ set(SOURCES
 	Solver.cpp
 	Problem.hpp
 	Problem.cpp
+	PostStepData.hpp
+	PostStepData.cpp
 )
 
 source_group(TREE "${CMAKE_CURRENT_SOURCE_DIR}" PREFIX "Source Files" FILES ${SOURCES})

diff --git a/src/polysolve/nonlinear/PostStepData.cpp b/src/polysolve/nonlinear/PostStepData.cpp
@@ -0,0 +1,11 @@
+#include "PostStepData.hpp"
+
+namespace polysolve::nonlinear
+{
+    PostStepData::PostStepData(const int iter_num,
+                               const Eigen::VectorXd &x,
+                               const Eigen::VectorXd &grad)
+        : iter_num(iter_num), x(x), grad(grad)
+    {
+    }
+} // namespace polysolve::nonlinear
diff --git a/src/polysolve/nonlinear/PostStepData.hpp b/src/polysolve/nonlinear/PostStepData.hpp
@@ -0,0 +1,19 @@
+#pragma once
+
+#include <polysolve/Types.hpp>
+
+namespace polysolve::nonlinear
+{
+
+    class PostStepData
+    {
+    public:
+        PostStepData(const int iter_num,
+                     const Eigen::VectorXd &x,
+                     const Eigen::VectorXd &grad);
+
+        const int iter_num;
+        const Eigen::VectorXd &x;
+        const Eigen::VectorXd &grad;
+    };
+} // namespace polysolve::nonlinear
diff --git a/src/polysolve/nonlinear/Problem.hpp b/src/polysolve/nonlinear/Problem.hpp
@@ -2,6 +2,8 @@
 
 #include <polysolve/Types.hpp>
 
+#include "PostStepData.hpp"
+
 #include <cppoptlib/problem.h>
 
 #include <memory>
@@ -33,7 +35,7 @@ namespace polysolve::nonlinear
 
         virtual void line_search_begin(const TVector &x0, const TVector &x1) {}
         virtual void line_search_end() {}
-        virtual void post_step(const int iter_num, const TVector &x) {}
+        virtual void post_step(const PostStepData &data) {}
 
         virtual void set_project_to_psd(bool val) {}
 

diff --git a/src/polysolve/nonlinear/Solver.cpp b/src/polysolve/nonlinear/Solver.cpp
@@ -1,6 +1,8 @@
 
 #include "Solver.hpp"
 
+#include "PostStepData.hpp"
+
 #include "descent_strategies/BFGS.hpp"
 #include "descent_strategies/Newton.hpp"
 #include "descent_strategies/ADAM.hpp"
@@ -88,6 +90,10 @@ namespace polysolve::nonlinear
         {
             solver->add_strategy(std::make_unique<ADAM>(
                 solver_params, characteristic_length, logger));
+        else if (solver_name == "StochasticGradientDescent" || solver_name == "stochastic_gradient_descent")
+        {
+            solver->add_strategy(std::make_unique<GradientDescent>(
+                solver_params, true, characteristic_length, logger));
         }
         else if (solver_name == "GradientDescent" || solver_name == "gradient_descent")
         {
@@ -97,7 +103,7 @@ namespace polysolve::nonlinear
             throw std::runtime_error("Unrecognized solver type: " + solver_name);
 
         solver->add_strategy(std::make_unique<GradientDescent>(
-            solver_params, characteristic_length, logger));
+            solver_params, false, characteristic_length, logger));
 
         solver->set_strategies_iterations(solver_params);
         return solver;
@@ -110,6 +116,7 @@ namespace polysolve::nonlinear
                 "Newton",
                 "ADAM",
                 "GradientDescent",
+                "StochasticGradientDescent",
                 "L-BFGS"};
     }
 
@@ -133,8 +140,8 @@ namespace polysolve::nonlinear
         this->setStopCriteria(criteria);
 
         use_grad_norm_tol = solver_params["line_search"]["use_grad_norm_tol"];
-
         first_grad_norm_tol = solver_params["first_grad_norm_tol"];
+        allow_out_of_iterations = solver_params["allow_out_of_iterations"];
 
         use_grad_norm_tol *= characteristic_length;
         first_grad_norm_tol *= characteristic_length;
@@ -198,7 +205,7 @@ namespace polysolve::nonlinear
             objFunc.solution_changed(x);
         }
 
-        objFunc.post_step(this->m_current.iterations, 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;
@@ -378,7 +385,7 @@ namespace polysolve::nonlinear
                 m_logger.debug("[{}][{}] Objective decided to stop", name(), m_line_search->name());
             }
 
-            objFunc.post_step(this->m_current.iterations, x);
+            objFunc.post_step(PostStepData(this->m_current.iterations, x, grad));
 
             if (f_delta < this->m_stop.fDelta)
                 f_delta_step_cnt++;

diff --git a/src/polysolve/nonlinear/descent_strategies/GradientDescent.cpp b/src/polysolve/nonlinear/descent_strategies/GradientDescent.cpp
@@ -4,10 +4,13 @@ namespace polysolve::nonlinear
 {
 
     GradientDescent::GradientDescent(const json &solver_params_,
+                                     const bool is_stochastic,
                                      const double characteristic_length,
                                      spdlog::logger &logger)
-        : Superclass(solver_params_, characteristic_length, logger)
+        : Superclass(solver_params_, characteristic_length, logger), is_stochastic_(is_stochastic)
     {
+        if (is_stochastic_)
+            erase_component_probability_ = solver_params_["StochasticGradientDescent"]["erase_component_probability"];
     }
 
     bool GradientDescent::compute_update_direction(
@@ -18,6 +21,13 @@ namespace polysolve::nonlinear
     {
         direction = -grad;
 
+        if (is_stochastic_)
+        {
+            Eigen::VectorXd mask = (Eigen::VectorXd::Random(direction.size()).array() + 1.) / 2.;
+            for (int i = 0; i < direction.size(); ++i)
+                direction(i) *= (mask(i) < erase_component_probability_) ? 0. : 1.;
+        }
+
         return true;
     }
 

diff --git a/src/polysolve/nonlinear/descent_strategies/GradientDescent.hpp b/src/polysolve/nonlinear/descent_strategies/GradientDescent.hpp
@@ -13,15 +13,20 @@ namespace polysolve::nonlinear
         using Superclass = DescentStrategy;
 
         GradientDescent(const json &solver_params_,
+                        const bool is_stochastic,
                         const double characteristic_length,
                         spdlog::logger &logger);
 
-        std::string name() const override { return "GradientDescent"; }
+        std::string name() const override { return is_stochastic_ ? "StochasticGradientDescent" : "GradientDescent"; }
 
         bool compute_update_direction(
             Problem &objFunc,
             const TVector &x,
             const TVector &grad,
             TVector &direction) override;
+
+    private:
+        bool is_stochastic_ = false;
+        double erase_component_probability_ = 0;
     };
 } // namespace polysolve::nonlinear
diff --git a/src/polysolve/nonlinear/descent_strategies/Newton.cpp b/src/polysolve/nonlinear/descent_strategies/Newton.cpp
@@ -243,8 +243,13 @@ namespace polysolve::nonlinear
                                             polysolve::StiffnessMatrix &hessian)
 
     {
-        objFunc.set_project_to_psd(true);
-        objFunc.hessian(x, hessian);
+        if (x.size() != x_cache.size() || x != x_cache)
+        {
+            objFunc.set_project_to_psd(true);
+            objFunc.hessian(x, hessian_cache);
+            x_cache = x;
+        }
+        hessian = hessian_cache;
         if (reg_weight > 0)
         {
             hessian += reg_weight * sparse_identity(hessian.rows(), hessian.cols());

diff --git a/src/polysolve/nonlinear/descent_strategies/Newton.hpp b/src/polysolve/nonlinear/descent_strategies/Newton.hpp
@@ -114,6 +114,9 @@ namespace polysolve::nonlinear
         double reg_weight_max;
         double reg_weight_inc;
 
+        TVector x_cache;
+        polysolve::StiffnessMatrix hessian_cache;
+
         double reg_weight; ///< Regularization Coefficients
     protected:
         void compute_hessian(Problem &objFunc,