robust_norms.h

// Copyright 2026 Julien Michot.
// SPDX-License-Identifier: Apache-2.0

 
#pragma once
 
#include <tinyopt/diff/jet_traits.h>
#include <tinyopt/math.h>
#include <tinyopt/traits.h>
 
#include <tinyopt/losses/norms.h>
 
#include <tuple>
 
namespace tinyopt::losses {
 
struct HuberTag {};
struct CauchyTag {};
struct TukeyTag {};
 
template <typename ResidualT, typename LossTag>
struct RobustResidual {
  using StorageT = std::decay_t<ResidualT>;
 
  explicit RobustResidual(const ResidualT &residual, double delta)
      : residual(StorageT(residual)), delta(delta), delta2(delta * delta) {}
 
  StorageT residual;
  double delta = 0.0;
  double delta2 = 0.0;
};
 
template <typename T>
const T &UnwrapResidual(const T &residual) {
  return residual;
}
 
template <typename ResidualT, typename LossTag>
const ResidualT &UnwrapResidual(const RobustResidual<ResidualT, LossTag> &residual) {
  return residual.residual;
}
 
inline double HuberWeight(double n2, double delta) {
  const double delta2 = delta * delta;
  if (n2 <= delta2) return 1.0;
  return delta / std::sqrt(std::max(n2, std::numeric_limits<double>::min()));
}
 
inline double HuberLossValue(double n2, double delta) {
  const double delta2 = delta * delta;
  if (n2 <= delta2) return n2;
  return 2.0 * delta * std::sqrt(std::max(n2, std::numeric_limits<double>::min())) - delta2;
}
 
template <typename T, typename = void>
struct has_rows : std::false_type {};
 
template <typename T>
struct has_rows<T, std::void_t<decltype(std::declval<const T &>().rows())>> : std::true_type {};
 
template <typename T>
auto Huber(const T &residual, double delta)
    -> std::enable_if_t<!std::is_arithmetic_v<std::decay_t<T>> && !traits::is_pair_v<T> &&
                            !traits::is_jet_type_v<std::decay_t<T>> &&
                            (traits::is_matrix_or_array_v<std::decay_t<T>> ||
                             traits::is_robust_residual_v<std::decay_t<T>>),
                        RobustResidual<std::decay_t<T>, HuberTag>> {
  return RobustResidual<std::decay_t<T>, HuberTag>(residual, delta);
}
 
template <typename T>
auto Cauchy(const T &residual, double delta)
    -> std::enable_if_t<!std::is_arithmetic_v<std::decay_t<T>> && !traits::is_pair_v<T> &&
                            !traits::is_jet_type_v<std::decay_t<T>> &&
                            (traits::is_matrix_or_array_v<std::decay_t<T>> ||
                             traits::is_robust_residual_v<std::decay_t<T>>),
                        RobustResidual<std::decay_t<T>, CauchyTag>> {
  return RobustResidual<std::decay_t<T>, CauchyTag>(residual, delta);
}
 
template <typename T>
auto Tukey(const T &residual, double delta)
    -> std::enable_if_t<!std::is_arithmetic_v<std::decay_t<T>> && !traits::is_pair_v<T> &&
                            !traits::is_jet_type_v<std::decay_t<T>> &&
                            (traits::is_matrix_or_array_v<std::decay_t<T>> ||
                             traits::is_robust_residual_v<std::decay_t<T>>),
                        RobustResidual<std::decay_t<T>, TukeyTag>> {
  return RobustResidual<std::decay_t<T>, TukeyTag>(residual, delta);
}



template <typename T, typename ExportJ = std::nullptr_t>
auto Truncated(const T &n2, typename traits::params_trait<T>::Scalar th2,
               const ExportJ &Jx_or_bool = nullptr) {
  if constexpr (traits::is_pair_v<T>) {  // pair (loss, jacobian)
    return Truncated(n2.first, th2, n2.second);
  } else {
    static_assert(!traits::is_matrix_or_array_v<T>, "`n` must be a scalar");
    const auto l = n2 <= th2 ? n2 : T(th2);
    if constexpr (std::is_null_pointer_v<ExportJ>) {
      return l;
    } else if constexpr (traits::is_bool_v<ExportJ>) {
      return std::make_pair(l, n2 <= th2 ? T(1) : T(0));
    } else {  // Jx_or_bool is Jx
      assert(Jx_or_bool.rows() == 1);
      if (n2 <= th2) {
        return std::make_pair(l, Jx_or_bool);
      } else {
        return std::make_pair(l, ExportJ::Zero(1, Jx_or_bool.cols()).eval());
      }
    }
  }
}

template <typename T, typename ExportJ = std::nullptr_t>
auto TruncatedLoss(const T &x, typename traits::params_trait<T>::Scalar th2,
                   const ExportJ &Jx_or_bool = nullptr) {
  return Truncated(SquaredL2(x, Jx_or_bool), th2);
}


template <
    typename T, typename ExportJ = std::nullptr_t,
    std::enable_if_t<!traits::is_matrix_or_array_v<T> && !traits::is_robust_residual_v<T>, int> = 0>
auto Huber(const T &n2, typename traits::params_trait<T>::Scalar th2,
           const ExportJ &Jx_or_bool = nullptr) {
  if constexpr (traits::is_pair_v<T>) {  // pair (loss, jacobian)
    return Huber(n2.first, th2, n2.second);
  } else {
    using std::max;
    using std::sqrt;
    static_assert(!traits::is_matrix_or_array_v<T>, "`n` must be a scalar");
 
    if (n2 <= th2) {  // Inlier
      const auto l = n2;
      if constexpr (std::is_null_pointer_v<ExportJ>) {  // no jacobians
        return l;
      } else {
        if constexpr (traits::is_bool_v<ExportJ>)
          return std::make_pair(l, T(1));
        else
          return std::make_pair(l, Jx_or_bool);
      }
    } else {  // Outlier
      const auto th = T(sqrt(th2)), n = sqrt(n2);
      const auto l = T(2.0) * th * n - T(th2);
      if constexpr (std::is_null_pointer_v<ExportJ>) {  // no jacobians
        return l;
      } else {
        const T J_scale = max<T>(std::numeric_limits<T>::min(), th / n);
        if constexpr (traits::is_bool_v<ExportJ>)
          return std::make_pair(l, J_scale);
        else
          return std::make_pair(l, (J_scale * Jx_or_bool).eval());
      }
    }
  }
}

template <typename T, typename ExportJ = std::nullptr_t>
auto HuberLoss(const T &x, typename traits::params_trait<T>::Scalar th2,
               const ExportJ &Jx_or_bool = nullptr) {
  return Huber(SquaredL2(x, Jx_or_bool), th2);
}


template <
    typename T, typename ExportJ = std::nullptr_t,
    std::enable_if_t<!traits::is_matrix_or_array_v<T> && !traits::is_robust_residual_v<T>, int> = 0>
auto Tukey(const T &n2, typename traits::params_trait<T>::Scalar th2,
           const ExportJ &Jx_or_bool = nullptr) {
  if constexpr (traits::is_pair_v<T>) {  // pair (loss, jacobian)
    return Tukey(n2.first, th2, n2.second);
  } else {
    static_assert(!traits::is_matrix_or_array_v<T>, "`n` must be a scalar");
    if (n2 <= th2) {  // Inlier
      const auto s = T(1.0) - n2 / T(th2), s2 = s * s;
      const auto l = T(th2) * (T(1.0) - s2 * s);
      if constexpr (std::is_null_pointer_v<ExportJ>) {  // no jacobians
        return l;
      } else {
        const auto J_scale = T(3.0) * (th2 - n2) * (th2 - n2) / (th2 * th2);
        if constexpr (traits::is_bool_v<ExportJ>)
          return std::make_pair(l, J_scale);
        else
          return std::make_pair(l, (J_scale * Jx_or_bool).eval());
      }
    } else {  // Outlier
      const T l = T(th2);
      if constexpr (std::is_null_pointer_v<ExportJ>) {  // no jacobians
        return l;
      } else {
        if constexpr (traits::is_bool_v<ExportJ>)
          return std::make_pair(l, T(0));
        else
          return std::make_pair(l, ExportJ::Zero(1, Jx_or_bool.cols()).eval());
      }
    }
  }
}

template <typename T, typename ExportJ = std::nullptr_t>
auto TukeyLoss(const T &x, typename traits::params_trait<T>::Scalar th2,
               const ExportJ &Jx_or_bool = nullptr) {
  return Tukey(SquaredL2(x, Jx_or_bool), th2);
}


template <typename T, typename ExportJ = std::nullptr_t>
auto Arctan(const T &n2, typename traits::params_trait<T>::Scalar th2,
            const ExportJ &Jx_or_bool = nullptr) {
  if constexpr (traits::is_pair_v<T>) {  // pair (loss, jacobian)
    return Arctan(n2.first, th2, n2.second);
  } else {
    using std::atan2;
    using std::max;
    using std::sqrt;
    static_assert(!traits::is_matrix_or_array_v<T>, "`n` must be a scalar");
    const T th = T(sqrt(th2));
    const auto l = th * atan2(n2, th);
    if constexpr (std::is_null_pointer_v<ExportJ>) {  // no jacobians
      return l;
    } else {
      const auto tmp = n2 * n2 / th2;
      const T J_scale = max<T>(std::numeric_limits<T>::min(), T(1.0) / (tmp + T(1.0)));
      if constexpr (traits::is_bool_v<ExportJ>)
        return std::make_pair(l, J_scale);
      else
        return std::make_pair(l, (J_scale * Jx_or_bool).eval());
    }
  }
}

template <typename T, typename ExportJ = std::nullptr_t>
auto ArctanLoss(const T &x, typename traits::params_trait<T>::Scalar th2,
                const ExportJ &Jx_or_bool = nullptr) {
  return Arctan(SquaredL2(x, Jx_or_bool), th2);
}


template <
    typename T, typename ExportJ = std::nullptr_t,
    std::enable_if_t<!traits::is_matrix_or_array_v<T> && !traits::is_robust_residual_v<T>, int> = 0>
auto Cauchy(const T &n2, typename traits::params_trait<T>::Scalar th2,
            const ExportJ &Jx_or_bool = nullptr) {
  if constexpr (traits::is_pair_v<T>) {  // pair (loss, jacobian)
    return Cauchy(n2.first, th2, n2.second);
  } else {
    using std::log;
    using std::sqrt;
    static_assert(!traits::is_matrix_or_array_v<T>, "`n` must be a scalar");
    const auto s = T(1.0) + n2 / T(th2);
    const auto l = T(th2) * log(s);
    if constexpr (std::is_null_pointer_v<ExportJ>) {  // no jacobians
      return l;
    } else {
      const T J_scale = std::max<T>(std::numeric_limits<T>::min(), T(1.0) / s);
      if constexpr (traits::is_bool_v<ExportJ>)
        return std::make_pair(l, J_scale);
      else
        return std::make_pair(l, (J_scale * Jx_or_bool).eval());
    }
  }
}

template <typename T, typename ExportJ = std::nullptr_t>
auto CauchyLoss(const T &x, typename traits::params_trait<T>::Scalar th2,
                const ExportJ &Jx_or_bool = nullptr) {
  return Cauchy(SquaredL2(x, Jx_or_bool), th2);
}


template <typename T, typename ExportJ = std::nullptr_t>
auto GemanMcClure(const T &n2, typename traits::params_trait<T>::Scalar th2,
                  const ExportJ &Jx_or_bool = nullptr) {
  if constexpr (traits::is_pair_v<T>) {  // pair (loss, jacobian)
    return GemanMcClure(n2.first, th2, n2.second);
  } else {
    using std::log;
    using std::sqrt;
    static_assert(!traits::is_matrix_or_array_v<T>, "`n` must be a scalar");
    const T e2_th2 = n2 + T(th2);
    const auto l = n2 / (e2_th2);
    if constexpr (std::is_null_pointer_v<ExportJ>) {  // no jacobians
      return l;
    } else {
      const auto J_scale = th2 / (e2_th2 * e2_th2);
      if constexpr (traits::is_bool_v<ExportJ>)
        return std::make_pair(l, J_scale);
      else
        return std::make_pair(l, (J_scale * Jx_or_bool).eval());
    }
  }
}

template <typename T, typename ExportJ = std::nullptr_t>
auto GemanMcClureLoss(const T &x, typename traits::params_trait<T>::Scalar th2,
                      const ExportJ &Jx_or_bool = nullptr) {
  return GemanMcClure(SquaredL2(x, Jx_or_bool), th2);
}


template <typename T, typename ExportJ = std::nullptr_t>
auto BlakeZisserman(const T &n2, typename traits::params_trait<T>::Scalar th2,
                    const ExportJ &Jx_or_bool = nullptr) {
  if constexpr (traits::is_pair_v<T>) {  // pair (loss, jacobian)
    return BlakeZisserman(n2.first, th2, n2.second);
  } else {
    using std::exp;
    using std::log;
    using std::sqrt;
    static_assert(!traits::is_matrix_or_array_v<T>, "`n` must be a scalar");
    const auto epsilon = T(exp(-th2));
    const auto l = -log(exp(-n2) + epsilon);
    if constexpr (std::is_null_pointer_v<ExportJ>) {  // no jacobians
      return l;
    } else {
      const auto J_scale = T(1.0) / (epsilon * exp(n2) + T(1.0));
      if constexpr (traits::is_bool_v<ExportJ>)
        return std::make_pair(l, J_scale);
      else
        return std::make_pair(l, (J_scale * Jx_or_bool).eval());
    }
  }
}

template <typename T, typename ExportJ = std::nullptr_t>
auto BlakeZissermanLoss(const T &x, typename traits::params_trait<T>::Scalar th2,
                        const ExportJ &Jx_or_bool = nullptr) {
  return BlakeZisserman(SquaredL2(x, Jx_or_bool), th2);
}

 
}  // namespace tinyopt::losses
 
namespace tinyopt {
 
template <typename... Args>
auto Residuals(Args &&...args) {
  return std::tuple<std::decay_t<Args>...>(std::forward<Args>(args)...);
}
 
}  // namespace tinyopt