gradient_check.h

// Copyright 2026 Julien Michot.
// SPDX-License-Identifier: Apache-2.0
 
#pragma once
 
#include <tinyopt/diff/num_diff.h>
#include <tinyopt/math.h>
 
#include <tinyopt/log.h>
#include <tinyopt/traits.h>
#include <cstddef>
#include <type_traits>
#include "tinyopt/losses/norms.h"
 
namespace tinyopt::diff {

//                      `Cost(const X_t& x, Gradient_t& grad, Hessian_t &)`.  The function *must*
//                      compute
template <typename X_t, typename AccFunc>
inline bool CheckGradient(
    X_t &x, const AccFunc &acc,
    double eps = std::is_same_v<typename traits::params_trait<X_t>::Scalar, float> ? 1e-2 : 1e-5,
    const diff::Method method = diff::Method::kCentral, bool verbose = true) {
  using Scalar = traits::params_trait<X_t>::Scalar;
  constexpr int Dims = traits::params_trait<X_t>::Dims;
  const int dims = traits::DynDims(x);
  using Gradient_t = Vector<Scalar, Dims>;
  constexpr bool HasH = !std::is_invocable_v<AccFunc, const X_t &, Gradient_t &>;
  constexpr bool DenseH =
      std::is_invocable_v<AccFunc, const X_t &, Gradient_t &, Matrix<Scalar, Dims, Dims> &>;
  using MaybeHessian_t = std::conditional_t<
      HasH, std::conditional_t<DenseH, Matrix<Scalar, Dims, Dims>, SparseMatrix<Scalar>>,
      std::nullptr_t>;
 
  std::nullptr_t nul;
  // Get the gradient provided by 'acc'
  Gradient_t grad(dims);
  grad.setZero();
 
  auto fg = [&](const X_t &x, Gradient_t &g) {
    if constexpr (HasH &&
                  !std::is_invocable_v<AccFunc, const X_t &, Gradient_t &, MaybeHessian_t &>)
      return acc(x, g, nul);
    else if constexpr (HasH) {
      MaybeHessian_t H;
      H.resize(dims, dims);
      H.setZero();
      return acc(x, g, H);
    } else
      return acc(x, g);
  };
  fg(x, grad);
 
  // Estimate numerical jacobian
  auto f = [&](const X_t &x) {
    if constexpr (HasH) {
      MaybeHessian_t H;
      H.resize(dims, dims);
      H.setZero();
      return acc(x, nul, H);
    } else
      return acc(x, nul);
  };
 
  const auto grad_num = EstimateNumJac(x, f, method, eps / 10.0f).transpose().eval();
  const double max_dist = (grad_num - grad).cwiseAbs().maxCoeff();
  const bool is_within_tol = max_dist < eps;
  if (!is_within_tol && verbose)
    TINYOPT_LOG("Wrong gradient {:.3e}>eps. \nGradient:{}, Numerial Gradient:{}", max_dist, grad,
                grad_num);
  return is_within_tol;
}

//         This function helps verify the correctness of the
//                      (lambda, std::function, or functor) It should have one of the signatures
//                      `Residuals(const X_t& x, Gradient_t& grad, Hessian_t &)`.  The function
//                      *must* compute the gradient and store it in the `grad` argument.
template <typename X_t, typename AccWithResidualsFunc>
inline bool CheckResidualsGradient(
    X_t &x, const AccWithResidualsFunc &acc_res_func,
    double eps = std::is_same_v<typename traits::params_trait<X_t>::Scalar, float> ? 1e-2 : 1e-5,
    const diff::Method method = diff::Method::kCentral, bool verbose = true, bool check_H = true,
    bool downscale = true) {
  using Scalar = traits::params_trait<X_t>::Scalar;
  constexpr int Dims = traits::params_trait<X_t>::Dims;
  const int dims = traits::DynDims(x);
  using Gradient_t = Vector<Scalar, Dims>;
  constexpr bool HasH = !std::is_invocable_v<AccWithResidualsFunc, const X_t &, std::nullptr_t &>;
  constexpr bool DenseH = std::is_invocable_v<AccWithResidualsFunc, const X_t &, Gradient_t &,
                                              Matrix<Scalar, Dims, Dims> &>;
  using MaybeHessian_t = std::conditional_t<
      HasH, std::conditional_t<DenseH, Matrix<Scalar, Dims, Dims>, SparseMatrix<Scalar>>,
      std::nullptr_t>;
  const double s = downscale ? 0.5 : 1.0;
 
  // Apply L2
  auto loss = [&](const X_t &x, auto &g, MaybeHessian_t &H) {
    if constexpr (HasH) {
      auto res = acc_res_func(x, g, H);
 
      using ResidualsType = std::decay_t<decltype(res)>;
      static_assert(!std::is_same_v<ResidualsType, Cost>,
                    "The function must return residuals, not a cost");
      return s * losses::SquaredL2(res);
    } else {
      auto res = acc_res_func(x, g);
 
      using ResidualsType = std::decay_t<decltype(res)>;
      static_assert(!std::is_same_v<ResidualsType, Cost>,
                    "The function must return residuals, not a cost");
 
      return s * losses::SquaredL2(res);
    }
  };
  bool success = CheckGradient(x, loss, eps, method, verbose);
 
  // Also check the hessian approx Jt*J
  if constexpr (HasH) {
    if (!check_H) return success;
 
    // Get user defined H
    Gradient_t g;
    g.resize(dims);
    g.setZero();
    MaybeHessian_t H;
    H.resize(dims, dims);
    H.setZero();
    acc_res_func(x, g, H);
 
    // Estimate numerical J and H = Jt*J
    auto f = [&](const X_t &x) {
      std::nullptr_t nul;
      return acc_res_func(x, nul, H);
    };
 
    const auto J_num = EstimateNumJac(x, f, method, eps / 10.0f);
    const auto H_num = (J_num.transpose() * J_num).eval();
 
    double max_dist;
    if constexpr (DenseH)
      max_dist = (H_num - H).cwiseAbs().maxCoeff();
    else {
      MaybeHessian_t H_num_sparse = H_num.sparseView();
      H_num_sparse.makeCompressed();
      max_dist = MaxAbsDiff(H_num_sparse, H);
    }
 
    const bool is_within_tol = max_dist < eps;
    if (!is_within_tol && verbose)
      TINYOPT_LOG("Wrong gradient {:.3e}>eps. \nH:\n{}, Numerial JtJ:\n{}", max_dist, H, H_num);
    return success && is_within_tol;
  }
 
  return success;
}
 
}  // namespace tinyopt::diff