norms.h

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


template <typename T>
auto SquaredL2(const T &x) {
  constexpr bool IsMatrix = traits::is_matrix_or_array_v<T> || traits::is_sparse_matrix_v<T>;
  if constexpr (traits::is_pair_v<T>) {  // pair
    return SquaredL2(x.first, x.second);
  } else if constexpr (IsMatrix) {  // scalar
    return x.squaredNorm();
  } else {  // scalar
    return x * x;
  }
}

template <typename T, typename ExportJ>
auto SquaredL2(const T &x, const ExportJ &Jx_or_bool, bool add_scale = true) {
  constexpr bool IsMatrix = traits::is_matrix_or_array_v<T> || traits::is_sparse_matrix_v<T>;
  const auto l = SquaredL2(x);
  if constexpr (std::is_null_pointer_v<ExportJ>) {
    return l;
  } else if constexpr (IsMatrix) {
    using Scalar = typename traits::params_trait<T>::Scalar;
    RowVector<Scalar, traits::params_trait<T>::Dims> J = x.transpose();
    if (add_scale) J *= Scalar(2.0);
    if constexpr (traits::is_bool_v<ExportJ>)
      return std::make_pair(l, J);
    else
      return std::make_pair(l, (J * Jx_or_bool).eval());
  } else {  // scalar
    return std::make_pair(l, add_scale ? T(2.0) * x : x);
  }
}

template <typename T>
auto L2(const T &x) {
  constexpr bool IsMatrix = traits::is_matrix_or_array_v<T> || traits::is_sparse_matrix_v<T>;
  if constexpr (traits::is_pair_v<T>) {  // pair
    return L2(x.first, x.second);
  } else if constexpr (IsMatrix) {  // scalar
    return x.norm();
  } else {  // scalar
    using std::abs;
    return abs(x);  // same as sqrt(x*x);
  }
}

template <typename T, typename ExportJ>
auto L2(const T &x, const ExportJ &Jx_or_bool) {
  constexpr bool IsMatrix = traits::is_matrix_or_array_v<T> || traits::is_sparse_matrix_v<T>;
  const auto l = L2(x);
  if constexpr (std::is_null_pointer_v<ExportJ>) {
    return l;
  } else if constexpr (IsMatrix) {
    const auto J = l > FloatEpsilon() ? (x / l).eval() : x;
    if constexpr (traits::is_bool_v<ExportJ>)
      return std::make_pair(l, J.transpose().eval());
    else
      return std::make_pair(l, (J.transpose() * Jx_or_bool).eval());
  } else {  // scalar
    return std::make_pair(l, 1);
  }
}

template <typename T>
auto L1(const T &x) {
  constexpr bool IsMatrix = traits::is_matrix_or_array_v<T> || traits::is_sparse_matrix_v<T>;
  if constexpr (traits::is_pair_v<T>) {  // pair
    return L1(x.first, x.second);
  } else if constexpr (IsMatrix) {  // scalar
    return x.template lpNorm<1>();
  } else {  // scalar
    using std::abs;
    return abs(x);
  }
}

template <typename T, typename ExportJ>
auto L1(const T &x, const ExportJ &Jx_or_bool) {
  constexpr bool IsMatrix = traits::is_matrix_or_array_v<T> || traits::is_sparse_matrix_v<T>;
  const auto l = L1(x);
  if constexpr (std::is_null_pointer_v<ExportJ>) {
    return l;
  } else if constexpr (IsMatrix) {
    using Scalar = typename traits::params_trait<T>::Scalar;
    constexpr int DimsJ = traits::params_trait<T>::Dims;
    if constexpr (traits::is_bool_v<ExportJ>)
      return std::make_pair(l, RowVector<Scalar, DimsJ>(x.array().sign()));
    else
      return std::make_pair(l, (x.array().sign() * Jx_or_bool).eval());
  } else {  // scalar
    return std::make_pair(l, 1);
  }
}

template <typename T>
auto Linf(const T &x) {
  constexpr bool IsMatrix = traits::is_matrix_or_array_v<T> || traits::is_sparse_matrix_v<T>;
  if constexpr (traits::is_pair_v<T>) {  // pair
    return Linf(x.first, x.second);
  } else if constexpr (IsMatrix) {
    return x.template lpNorm<Infinity>();
  } else {  // scalar
    using std::abs;
    return abs(x);
  }
}

template <typename T, typename ExportJ>
auto Linf(const T &x, const ExportJ &Jx_or_bool) {
  constexpr bool IsMatrix = traits::is_matrix_or_array_v<T> || traits::is_sparse_matrix_v<T>;
  if constexpr (std::is_null_pointer_v<ExportJ>) {
    return Linf(x);
  } else if constexpr (IsMatrix) {
    using Scalar = typename traits::params_trait<T>::Scalar;
    constexpr int DimsJ = traits::params_trait<T>::Dims;
    RowVector<Scalar, DimsJ> J(1, x.rows());
    J.setZero();
    int max_idx;
    const auto l = x.cwiseAbs().maxCoeff(&max_idx);
    J[max_idx] = x[max_idx] >= 0 ? 1 : -1;
    if constexpr (traits::is_bool_v<ExportJ>)
      return std::make_pair(l, J);
    else
      return std::make_pair(l, (J * Jx_or_bool).eval());
  } else {  // scalar
    if (x >= 0)
      return std::make_pair(x, T(1.0));
    else
      return std::make_pair(-x, T(-1.0));
  }
}

 
}  // namespace tinyopt::losses