num_diff.h

// Copyright 2026 Julien Michot.
// SPDX-License-Identifier: Apache-2.0
 
#pragma once
 
#include <tinyopt/cost.h>
#include <tinyopt/math.h>  // Defines Matrix and Vector
#include <tinyopt/traits.h>
#include <cstddef>
 
namespace tinyopt::diff {
enum Method {
  kForward = 0,

  kCentral,

  kFastCentral
};

template <typename X_t, typename Func>
auto NumEval(const X_t &x, const Func &f, const diff::Method &method = diff::Method::kCentral,
             typename traits::params_trait<X_t>::Scalar h =
                 FloatEpsilon<typename traits::params_trait<X_t>::Scalar>()) {
  using ptrait = traits::params_trait<X_t>;
  using Scalar = typename ptrait::Scalar;
 
  constexpr Index Dims = ptrait::Dims;
  const Index dims = traits::DynDims(x);
 
  // Support different function signatures
  auto fg = [&](const auto &x) {
    std::nullptr_t nul;
    if constexpr (std::is_invocable_v<Func, const X_t &>)
      return f(x);
    else if constexpr (std::is_invocable_v<Func, const X_t &, std::nullptr_t &>)
      return f(x, nul);
    else if constexpr (std::is_invocable_v<Func, const X_t &, std::nullptr_t &, std::nullptr_t &>)
      return f(x, nul, nul);
    else {  // likely a SparseMatrix<Scalar> hessian
      SparseMatrix<Scalar> H;
      H.resize(dims, dims);
      return f(x, nul, H);
    }
  };
 
  // Recover current residuals
  const auto res = fg(x);
  // Declare the jacobian matrix
  using ResType = typename std::decay_t<decltype(res)>;
  constexpr int ResDims = traits::params_trait<ResType>::Dims;
  const Index res_dims = traits::DynDims(res);
 
  using J_t = Matrix<Scalar, ResDims, Dims>;
  J_t J(res_dims, dims);
 
  // Estimate the jacobian using numerical differentiation
  Vector<Scalar, Dims> dx = Vector<Scalar, Dims>::Zero(dims);
  for (Index r = 0; r < dims; ++r) {
    X_t y = x;  // copy
    if (r > 0) dx[r - 1] = 0;
    dx[r] = h;
    ptrait::PlusEq(y, dx);
    const auto res_plus = fg(y);
    using ResType2 = typename std::decay_t<decltype(res_plus)>;
    if (method == Method::kCentral) {
      y = x;  // copy again
      dx[r] = -h;
      ptrait::PlusEq(y, dx);
      const auto res_minus = fg(y);
      if constexpr (std::is_scalar_v<ResType2>)
        J[r] = (res_plus - res_minus) / (2 * h);
      else
        J.col(r) = (res_plus.reshaped() - res_minus.reshaped()) / (2 * h);
    } else if (method == Method::kFastCentral) {
      dx[r] = -2 * h;  // given a small h, one can use this approximation, hopefully
      ptrait::PlusEq(y, dx);
      const auto res_minus = fg(y);
      if constexpr (std::is_scalar_v<ResType2>)
        J[r] = (res_plus - res_minus) / (2 * h);
      else
        J.col(r) = (res_plus.reshaped() - res_minus.reshaped()) / (2 * h);
    } else {
      if constexpr (std::is_scalar_v<ResType2>)
        J[r] = (res_plus - res) / h;
      else
        J.col(r) = (res_plus.reshaped() - res) / h;
    }
  }
  return std::make_pair(res, J);
}

template <typename X_t, typename Func>
auto EstimateNumJac(const X_t &x, const Func &f,
                    const diff::Method &method = diff::Method::kCentral,
                    typename traits::params_trait<X_t>::Scalar h =
                        FloatEpsilon<typename traits::params_trait<X_t>::Scalar>()) {
  const auto &[res, J] = NumEval(x, f, method, h);
  return J;
}

template <typename X_t, typename ResidualsFunc>
auto CreateNumDiffFunc1(X_t &, const ResidualsFunc &residuals,
                        const Method &method = Method::kCentral,
                        typename traits::params_trait<X_t>::Scalar h =
                            FloatEpsilon<typename traits::params_trait<X_t>::Scalar>()) {
  auto loss = [&residuals, method, h](const auto &x, auto &grad) {
    constexpr bool HasGrad = !traits::is_nullptr_v<decltype(grad)>;
    // Recover current residuals
    const auto res = residuals(x);
 
    if constexpr (HasGrad) {
      const auto J = EstimateNumJac(x, residuals, method, h);
      grad = J.transpose() * res;
    }
 
    using ResType = typename std::decay_t<decltype(res)>;
    if constexpr (std::is_scalar_v<ResType>) {
      return std::abs(res);
    } else {
      // Returns the norm + number of residuals
      return Cost(res.matrix().norm(), res.size());
    }
  };
  return loss;
}

template <typename X_t, typename ResidualsFunc>
auto CreateNumDiffFunc2(X_t &, const ResidualsFunc &residuals,
                        const Method &method = Method::kCentral,
                        typename traits::params_trait<X_t>::Scalar h =
                            FloatEpsilon<typename traits::params_trait<X_t>::Scalar>()) {
  auto loss = [&residuals, method, h](const auto &x, auto &grad, auto &H) {
    constexpr bool HasGrad = !traits::is_nullptr_v<decltype(grad)>;
    constexpr bool HasH = !traits::is_nullptr_v<decltype(H)>;
    // Recover current residuals
    const auto res = residuals(x);
 
    if constexpr (HasGrad) {
      const auto J = EstimateNumJac(x, residuals, method, h);
      grad = J.transpose() * res;
      if constexpr (HasH) H = J.transpose() * J;
    }
 
    using ResType = typename std::decay_t<decltype(res)>;
    if constexpr (std::is_scalar_v<ResType>) {
      return std::abs(res);
    } else {
      // Returns the norm + number of residuals
      return Cost(res.matrix().norm(), res.size());
    }
  };
  return loss;
}
 
}  // namespace tinyopt::diff