gn.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>
#include <optional>
#include "tinyopt/math.h"
 
namespace tinyopt::solvers {
 
template <typename Hessian_t = MatX>
class SolverGN
    : public SolverBase<typename Hessian_t::Scalar, SQRT(traits::params_trait<Hessian_t>::Dims)> {
 public:
  static constexpr bool IsNLLS = true;
  static constexpr bool FirstOrder = false;  // this is a pseudo second order algorithm
  using Base = SolverBase<typename Hessian_t::Scalar, SQRT(traits::params_trait<Hessian_t>::Dims)>;
  using Scalar = typename Hessian_t::Scalar;
  static constexpr Index Dims = SQRT(traits::params_trait<Hessian_t>::Dims);
 
  // Hessian Type
  using H_t = Hessian_t;
  // Gradient Type
  using Grad_t = Vector<Scalar, Dims>;
  // Options
  using Base::options_;
 
  explicit SolverGN(const Options &options = {}) : Base(options) {}

  void InitWith(const Grad_t &g, const H_t &h) {
    grad_ = g;
    H_ = h;
  }

  virtual 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 || H_.rows() != dims) {
      H_.resize(dims, 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<H_t>) {
      H_.resize(dims, dims);
      grad_.resize(dims);
      return true;
    }
    return false;
  }

  void clear() {
    // Fill H & grad fill 0s (not needed when using auto-jet)
    H_.setZero();
    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 X_t, typename AccFunc>
  inline Scalar Evaluate(const X_t &x, const AccFunc &acc, bool save) {
    std::nullptr_t nul;
    Hessian_t H;  // dummy;
    Cost cost = acc(x, nul, H);
    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) {
    this->cost_ = acc(x, grad_, H_);
    this->NormalizeCost(this->cost_);
    return this->cost_.isValid();
  }

  template <typename X_t, typename AccFunc>
  inline bool Build(const X_t &x, const AccFunc &acc_func, bool resize_and_clear = true) {
    // Resize the system if needed and clear gradient
    if (resize_and_clear) {
      ResizeIfNeeded(x);
      clear();
    }
 
    // Accumulate residuals and update both gardient and Hessian approx (Jt*J)
    const bool success = Accumulate(x, acc_func);
    if (!success) return false;
 
    // Eventually clip the gradient
    this->Clamp(grad_, options_.opt.grad_clipping);
 
    // Verify Hessian's diagonal
    if (options_.hessian.check_min_H_diag > 0 &&
        (H_.diagonal().cwiseAbs().array() < options_.hessian.check_min_H_diag).any()) {
      if (options_.log.enable) TINYOPT_LOG("❌ Hessian has very low diagonal coefficients");
      return false;
    }
 
    // Fill the lower part if H if needed
    {
      if (!options_.hessian.H_is_full && RequiresFullMatrix(options_.linear_solver))
        CompleteSymmetricMatrix(H_);
    }
#if defined(TINYOPT_ENABLE_SUITESPARSE)
    if constexpr (traits::is_sparse_matrix_v<H_t>) {
      if (options_.linear_solver == LinearSolverMethod::SuiteSparse) H_.makeCompressed();
    }
#endif
    return true;
  }

  inline std::optional<Vector<Scalar, Dims>> Solve() const override {
    if (!this->cost().isValid()) return std::nullopt;
 
    const auto dx_ = tinyopt::SolveLinearSystem(H_, -grad_, options_.linear_solver,
                          options_.svd_relative_threshold);
    if (dx_) return dx_;
    // Log on failure
    if (options_.log.enable && options_.log.print_failure) {
      TINYOPT_LOG("❌ Failed solve linear system");
      TINYOPT_LOG("grad = \n{}", grad_);
      TINYOPT_LOG("H = \n{}", H_);
    }
    return std::nullopt;
  }
 
  virtual Index dims() const override { return grad_.size(); }

  Scalar MaxStdDev(bool use_damped = true) const {
    const auto &H = use_damped ? H_ : Hessian();
    if constexpr (traits::is_sparse_matrix_v<H_t>)
      return sqrt(InvCov(H).value().coeffs().maxCoeff());
    else
      return sqrt(InvCov(H).value().maxCoeff());
  }

  virtual std::optional<H_t> Covariance() const { return InvCov(H_); }

  const H_t &Hessian() const { return H_; }
 
  const H_t &H() const { return H_; }
  H_t &H() { return H_; }
 
  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:
  H_t H_;
  Grad_t grad_;
};
 
}  // namespace tinyopt::solvers