gd.h

// Copyright 2026 Julien Michot.
// SPDX-License-Identifier: Apache-2.0
 
#pragma once
 
#include <tinyopt/cost.h>
#include <tinyopt/solvers/base.h>
#include <tinyopt/traits.h>
 
namespace tinyopt::solvers {
 
template <typename Gradient_t = VecX>
class SolverGD
    : public SolverBase<typename Gradient_t::Scalar, traits::params_trait<Gradient_t>::Dims> {
 public:
  static constexpr bool IsNLLS = false;  // by default it is not
  static constexpr bool FirstOrder = true;
  using Base = SolverBase<typename Gradient_t::Scalar, traits::params_trait<Gradient_t>::Dims>;
  using Scalar = typename Gradient_t::Scalar;
  static constexpr Index Dims = traits::params_trait<Gradient_t>::Dims;
  // Gradient Type
  using Grad_t = Gradient_t;
  // Hessian Type (here none, so nullptr_t)
  using H_t = std::nullptr_t;
  // Options
  using Base::options_;
 
  explicit SolverGD(const Options &options = {}) : Base(options) {}

  void InitWith(const Grad_t &g) { grad_ = g; }

  void reset() { clear(); }

  template <int D = Dims, std::enable_if_t<D == Dynamic, int> = 0>
  bool resize(int dims) {
    if (dims == Dynamic) {
      throw std::invalid_argument("Dimensions cannot be Dynamic here");
    }
    if (grad_.rows() != dims) {
      grad_.resize(dims);
      return true;
    } else {
      return false;
    }
  }

  template <int D = Dims, std::enable_if_t<D != Dynamic, int> = 0>
  bool resize(int dims = Dims) {
    if (dims != Dims) {
      throw std::invalid_argument("Error: Static and Dynamic Dimensions must match");
    }
    if constexpr (traits::is_sparse_matrix_v<Grad_t>) {
      grad_.resize(dims);
      return true;
    }
    return false;
  }

  void clear() { grad_.setZero(); }

  template <typename X_t>
  bool ResizeIfNeeded(const X_t &x) {
    if constexpr (Dims == Dynamic) {
      const int dims = traits::params_trait<X_t>::dims(x);
      if (grad_.rows() != dims) {
        if (options_.log.enable) TINYOPT_LOG("Need to resize the system");
        return resize(dims);
      }
    }
    return false;
  }

  template <typename AccFunc>
  inline auto GetAccFunc(const AccFunc &acc) const {
    return [&](const auto &x, auto &grad) -> Cost {
      using X_t = decltype(x);
      if constexpr (std::is_invocable_v<AccFunc, const X_t &, std::nullptr_t &>)
        return acc(x, grad);
      else {
        std::nullptr_t nul;
        return acc(x, grad, nul);
      }
    };
  }

  template <typename X_t, typename AccFunc>
  inline Scalar Evaluate(const X_t &x, const AccFunc &acc, bool save) {
    std::nullptr_t nul;
    const auto acc2 = GetAccFunc(acc);
    const auto &cost = acc2(x, nul);
    this->NormalizeCost(cost);
    if (save) this->cost_ = cost;
    return cost.cost;
  }

  template <typename X_t, typename AccFunc>
  inline bool Accumulate(const X_t &x, const AccFunc &acc) {
    const auto acc2 = GetAccFunc(acc);
    this->cost_ = acc2(x, grad_);
    this->NormalizeCost(this->cost_);
    return this->cost_.isValid();
  }

  template <typename X_t, typename AccFunc>
  inline bool Build(const X_t &x, const AccFunc &acc, bool resize_and_clear = true) {
    // Resize the system if needed and clear gradient
    if (resize_and_clear) {
      ResizeIfNeeded(x);
      clear();
    }
    // Update gradient by accumulating changes
    const bool ok = Accumulate(x, acc);
    // Eventually clip the gradients
    this->Clamp(grad_, options_.opt.grad_clipping);
    return ok;
  }

  inline std::optional<Vector<Scalar, Dims>> Solve() const override {
    if (!this->cost().isValid()) return std::nullopt;
    return -options_.gd.lr * grad_;
  }
 
  Index dims() const override { return grad_.size(); }
 
  const Grad_t &Gradient() const { return grad_; }
  Grad_t &Gradient() { return grad_; }
  Scalar GradientNorm() const { return grad_.norm(); }
  Scalar GradientSquaredNorm() const { return grad_.squaredNorm(); }
 
 protected:
  Grad_t grad_;
};
 
}  // namespace tinyopt::solvers