Add ADAM optimizer

non-linear-solver-spec.json

@@ -15,6 +15,7 @@
             "LBFGS",
             "LBFGSB",
             "Newton",
+            "ADAM",
             "box_constraints",
             "advanced"
         ],
@@ -28,6 +29,7 @@
             "Newton",
             "DenseNewton",
             "GradientDescent",
+            "ADAM",
             "L-BFGS",
             "BFGS",
             "L-BFGS-B",
@@ -164,6 +166,42 @@
         "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": "/line_search",
         "default": null,

src/polysolve/nonlinear/Solver.cpp

@@ -3,6 +3,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"
 
@@ -83,6 +84,11 @@ namespace polysolve::nonlinear
             solver->add_strategy(std::make_unique<LBFGS>(
                 solver_params, characteristic_length, logger));
         }
+        else if (solver_name == "ADAM" || solver_name == "adam")
+        {
+            solver->add_strategy(std::make_unique<ADAM>(
+                solver_params, characteristic_length, logger));
+        }
         else if (solver_name == "GradientDescent" || solver_name == "gradient_descent")
         {
             // grad descent always there
@@ -102,6 +108,7 @@ namespace polysolve::nonlinear
         return {"BFGS",
                 "DenseNewton",
                 "Newton",
+                "ADAM",
                 "GradientDescent",
                 "L-BFGS"};
     }

src/polysolve/nonlinear/descent_strategies/ADAM.cpp

@@ -0,0 +1,54 @@
+// ADAM from "ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION"
+
+#include "ADAM.hpp"
+
+namespace polysolve::nonlinear
+{
+
+    ADAM::ADAM(const json &solver_params,
+               const double characteristic_length,
+               spdlog::logger &logger)
+        : Superclass(solver_params, characteristic_length, logger)
+    {
+        alpha = solver_params["ADAM"]["alpha"];
+        beta_1 = solver_params["ADAM"]["beta_1"];
+        beta_2 = solver_params["ADAM"]["beta_2"];
+        epsilon = solver_params["ADAM"]["epsilon"];
+    }
+
+    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 m = (beta_1 * m_prev) + ((1 - beta_1) * grad);
+        TVector v = beta_2 * v_prev;
+        for (int i = 0; i < v.size(); ++i)
+            v(i) += (1 - beta_2) * grad(i) * grad(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

src/polysolve/nonlinear/descent_strategies/ADAM.hpp

@@ -0,0 +1,39 @@
+#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 double characteristic_length,
+             spdlog::logger &logger);
+
+        std::string name() const override { return "ADAM"; }
+
+        void reset(const int ndof) override;
+
+        virtual bool compute_update_direction(
+            Problem &objFunc,
+            const TVector &x,
+            const TVector &grad,
+            TVector &direction) override;
+
+    private:
+        TVector m_prev;
+        TVector v_prev;
+
+        double beta_1, beta_2;
+        double alpha;
+
+        int t = 0;
+        double epsilon;
+    };
+} // namespace polysolve::nonlinear

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
 )