Merge branch 'main' into custom-fallback

Teseo Schneider · Teseo Schneider · commit 05775a6f6589 · 2023-12-08T11:50:18.000-08:00
diff --git a/non-linear-solver-spec.json b/non-linear-solver-spec.json
@@ -15,6 +15,8 @@
             "L-BFGS",
             "L-BFGS-B",
             "Newton",
+            "ADAM",
+            "StochasticADAM",
             "StochasticGradientDescent",
             "box_constraints",
             "advanced"
@@ -29,6 +31,8 @@
             "Newton",
             "DenseNewton",
             "GradientDescent",
+            "ADAM",
+            "StochasticADAM",
             "StochasticGradientDescent",
             "L-BFGS",
             "BFGS",
@@ -166,6 +170,85 @@
         "type": "bool",
         "doc": "Use PSD as fallback using second order solvers (i.e., Newton's method)."
     },
+    {
+        "pointer": "/ADAM",
+        "default": null,
+        "type": "object",
+        "optional": [
+            "alpha",
+            "beta_1",
+            "beta_2",
+            "epsilon"
+        ],
+        "doc": "Options for ADAM."
+    },
+    {
+        "pointer": "/ADAM/alpha",
+        "default": 0.001,
+        "type": "float",
+        "doc": "Parameter alpha for ADAM."
+    },
+    {
+        "pointer": "/ADAM/beta_1",
+        "default": 0.9,
+        "type": "float",
+        "doc": "Parameter beta_1 for ADAM."
+    },
+    {
+        "pointer": "/ADAM/beta_2",
+        "default": 0.999,
+        "type": "float",
+        "doc": "Parameter beta_2 for ADAM."
+    },
+    {
+        "pointer": "/ADAM/epsilon",
+        "default": 1e-8,
+        "type": "float",
+        "doc": "Parameter epsilon for ADAM."
+    },
+    {
+        "pointer": "/StochasticADAM",
+        "default": null,
+        "type": "object",
+        "optional": [
+            "alpha",
+            "beta_1",
+            "beta_2",
+            "epsilon",
+            "erase_component_probability"
+        ],
+        "doc": "Options for ADAM."
+    },
+    {
+        "pointer": "/StochasticADAM/alpha",
+        "default": 0.001,
+        "type": "float",
+        "doc": "Parameter alpha for ADAM."
+    },
+    {
+        "pointer": "/StochasticADAM/beta_1",
+        "default": 0.9,
+        "type": "float",
+        "doc": "Parameter beta_1 for ADAM."
+    },
+    {
+        "pointer": "/StochasticADAM/beta_2",
+        "default": 0.999,
+        "type": "float",
+        "doc": "Parameter beta_2 for ADAM."
+    },
+    {
+        "pointer": "/StochasticADAM/epsilon",
+        "default": 1e-8,
+        "type": "float",
+        "doc": "Parameter epsilon for ADAM."
+    },
+    {
+        "pointer": "/StochasticADAM/erase_component_probability",
+        "default": 0.3,
+        "type": "float",
+        "doc": "Probability of erasing a component on the gradient for ADAM."
+    },
     {
         "pointer": "/StochasticGradientDescent",
         "default": null,
@@ -306,6 +389,37 @@
         ],
         "doc": "Options for BFGS."
     },
+    {
+        "pointer": "/solver/*",
+        "type": "object",
+        "type_name": "ADAM",
+        "required": [
+            "type"
+        ],
+        "optional": [
+            "alpha",
+            "beta_1",
+            "beta_2",
+            "epsilon"
+        ],
+        "doc": "Options for ADAM."
+    },
+    {
+        "pointer": "/solver/*",
+        "type": "object",
+        "type_name": "StochasticADAM",
+        "required": [
+            "type"
+        ],
+        "optional": [
+            "alpha",
+            "beta_1",
+            "beta_2",
+            "epsilon",
+            "erase_component_probability"
+        ],
+        "doc": "Options for ADAM."
+    },
     {
         "pointer": "/solver/*/type",
         "type": "string",
@@ -318,6 +432,8 @@
             "DenseRegularizedNewton",
             "GradientDescent",
             "StochasticGradientDescent",
+            "ADAM",
+            "StochasticADAM",
             "L-BFGS",
             "BFGS"
         ],
@@ -359,6 +475,30 @@
         "type": "int",
         "doc": "The number of corrections to approximate the inverse Hessian matrix."
     },
+    {
+        "pointer": "/solver/*/alpha",
+        "default": 0.001,
+        "type": "float",
+        "doc": "Parameter alpha for ADAM."
+    },
+    {
+        "pointer": "/solver/*/beta_1",
+        "default": 0.9,
+        "type": "float",
+        "doc": "Parameter beta_1 for ADAM."
+    },
+    {
+        "pointer": "/solver/*/beta_2",
+        "default": 0.999,
+        "type": "float",
+        "doc": "Parameter beta_2 for ADAM."
+    },
+    {
+        "pointer": "/solver/*/epsilon",
+        "default": 1e-8,
+        "type": "float",
+        "doc": "Parameter epsilon for ADAM."
+    },
     {
         "pointer": "/line_search",
         "default": null,
diff --git a/src/polysolve/nonlinear/Solver.cpp b/src/polysolve/nonlinear/Solver.cpp
@@ -5,6 +5,7 @@
 
 #include "descent_strategies/BFGS.hpp"
 #include "descent_strategies/Newton.hpp"
+#include "descent_strategies/ADAM.hpp"
 #include "descent_strategies/GradientDescent.hpp"
 #include "descent_strategies/LBFGS.hpp"
 
@@ -77,6 +78,15 @@ namespace polysolve::nonlinear
             {
                 return std::make_shared<GradientDescent>(solver_params, false, characteristic_length, logger);
             }
+
+            else if (solver_name == "ADAM" || solver_name == "adam")
+            {
+                return std::make_shared<ADAM>(solver_params, false, characteristic_length, logger);
+            }
+            else if (solver_name == "StochasticADAM" || solver_name == "stochastic_adam")
+            {
+                return std::make_shared<ADAM>(solver_params, true, characteristic_length, logger);
+            }
             else
                 throw std::runtime_error("Unrecognized solver type: " + solver_name);
         }
@@ -166,6 +176,8 @@ namespace polysolve::nonlinear
         return {"BFGS",
                 "DenseNewton",
                 "Newton",
+                "ADAM",
+                "StochasticADAM",
                 "GradientDescent",
                 "StochasticGradientDescent",
                 "L-BFGS"};
diff --git a/src/polysolve/nonlinear/descent_strategies/ADAM.cpp b/src/polysolve/nonlinear/descent_strategies/ADAM.cpp
@@ -0,0 +1,69 @@
+// ADAM from "ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION"
+
+#include "ADAM.hpp"
+
+#include <polysolve/Utils.hpp>
+
+namespace polysolve::nonlinear
+{
+
+    ADAM::ADAM(const json &solver_params,
+               const bool is_stochastic,
+               const double characteristic_length,
+               spdlog::logger &logger)
+        : Superclass(solver_params, characteristic_length, logger), is_stochastic_(is_stochastic)
+    {
+        std::string param_name = is_stochastic ? "StochasticADAM" : "ADAM";
+        alpha_ = extract_param(param_name, "alpha", solver_params);
+        beta_1_ = extract_param(param_name, "beta_1", solver_params);
+        beta_2_ = extract_param(param_name, "beta_2", solver_params);
+        epsilon_ = extract_param(param_name, "epsilon", solver_params);
+        if (is_stochastic)
+            erase_component_probability_ = extract_param("StochasticADAM", "erase_component_probability", solver_params);
+    }
+
+    void ADAM::reset(const int ndof)
+    {
+        Superclass::reset(ndof);
+        m_prev_ = Eigen::VectorXd::Zero(ndof);
+        v_prev_ = Eigen::VectorXd::Zero(ndof);
+        t_ = 0;
+    }
+
+    bool ADAM::compute_update_direction(
+        Problem &objFunc,
+        const TVector &x,
+        const TVector &grad,
+        TVector &direction)
+    {
+        if (m_prev_.size() == 0)
+            m_prev_ = Eigen::VectorXd::Zero(x.size());
+        if (v_prev_.size() == 0)
+            v_prev_ = Eigen::VectorXd::Zero(x.size());
+
+        TVector grad_modified = grad;
+
+        if (is_stochastic_)
+        {
+            Eigen::VectorXd mask = (Eigen::VectorXd::Random(direction.size()).array() + 1.) / 2.;
+            for (int i = 0; i < direction.size(); ++i)
+                grad_modified(i) *= (mask(i) < erase_component_probability_) ? 0. : 1.;
+        }
+
+        TVector m = (beta_1_ * m_prev_) + ((1 - beta_1_) * grad_modified);
+        TVector v = beta_2_ * v_prev_;
+        for (int i = 0; i < v.size(); ++i)
+            v(i) += (1 - beta_2_) * grad_modified(i) * grad_modified(i);
+
+        m = m.array() / (1 - pow(beta_1_, t_));
+        v = v.array() / (1 - pow(beta_2_, t_));
+
+        direction = -alpha_ * m;
+        for (int i = 0; i < v.size(); ++i)
+            direction(i) /= sqrt(v(i) + epsilon_);
+
+        ++t_;
+
+        return true;
+    }
+} // namespace polysolve::nonlinear
diff --git a/src/polysolve/nonlinear/descent_strategies/ADAM.hpp b/src/polysolve/nonlinear/descent_strategies/ADAM.hpp
@@ -0,0 +1,45 @@
+#pragma once
+
+#include "DescentStrategy.hpp"
+#include <polysolve/Utils.hpp>
+
+#include <polysolve/linear/Solver.hpp>
+
+namespace polysolve::nonlinear
+{
+    class ADAM : public DescentStrategy
+    {
+    public:
+        using Superclass = DescentStrategy;
+
+        ADAM(const json &solver_params,
+             const bool is_stochastic,
+             const double characteristic_length,
+             spdlog::logger &logger);
+
+        std::string name() const override { return is_stochastic_ ? "StochasticADAM" : "ADAM"; }
+
+        void reset(const int ndof) override;
+
+        virtual bool compute_update_direction(
+            Problem &objFunc,
+            const TVector &x,
+            const TVector &grad,
+            TVector &direction) override;
+
+        bool is_direction_descent() override { return false; }
+
+    private:
+        TVector m_prev_;
+        TVector v_prev_;
+
+        double beta_1_, beta_2_;
+        double alpha_;
+
+        int t_ = 0;
+        double epsilon_;
+
+        bool is_stochastic_;
+        double erase_component_probability_ = 0;
+    };
+} // namespace polysolve::nonlinear
diff --git a/src/polysolve/nonlinear/descent_strategies/CMakeLists.txt b/src/polysolve/nonlinear/descent_strategies/CMakeLists.txt
@@ -6,6 +6,8 @@ set(SOURCES
 	BFGS.hpp
 	GradientDescent.cpp
 	GradientDescent.hpp
+	ADAM.cpp
+	ADAM.hpp
 	Newton.hpp
 	Newton.cpp
 )

Original file line number	Diff line number	Diff line change
`@@ -6,6 +6,8 @@ set(SOURCES`
`6`	`6`	`BFGS.hpp`
`7`	`7`	`GradientDescent.cpp`
`8`	`8`	`GradientDescent.hpp`
	`9`	`+ ADAM.cpp`
	`10`	`+ ADAM.hpp`
`9`	`11`	`Newton.hpp`
`10`	`12`	`Newton.cpp`
`11`	`13`	`)`