trace.hpp

Go to the documentation of this file.

/*
 * Copyright © 2025-2026 Akreon LLC
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include "akreon/linalg/operator/operator_traits.hpp"
#include "weighted_sum.hpp"
 
namespace convex {
 
template <typename Space>
struct TraceWeights;
 
template <linalg::SquareMatrix Mat>
  requires(!linalg::is_symmetric_matrix_space_v<Mat>)
struct TraceWeights<Mat> {
  using T = linalg::space_type_t<Mat>;
  static constexpr std::size_t N = linalg::mat_rows_v<Mat>;
  static constexpr std::size_t Dim = linalg::space_dim_v<Mat>;
  using Vec = linalg::Vector<T, Dim>;
 
  static constexpr auto make() {
    Vec weights{};
    for (std::size_t i = 0; i < N; ++i) weights(i * N + i) = T(1);
    return weights;
  }
};
 
template <linalg::SquareMatrix Mat>
struct TraceWeights<linalg::SymmetricMatrix<Mat>> {
  using Space = linalg::SymmetricMatrix<Mat>;
  using T = linalg::space_type_t<Space>;
  static constexpr std::size_t N = linalg::mat_rows_v<Mat>;
  static constexpr std::size_t Dim = linalg::space_dim_v<Space>;
  using Vec = linalg::Vector<T, Dim>;
 
  static constexpr auto make() {
    Vec weights{};
    for (std::size_t i = 0; i < N; ++i) weights(i * (i + 1) / 2 + i) = T(1);
    return weights;
  }
};
 
template <linalg::RnSpace Space>
struct Trace {
  using T = linalg::space_type_t<Space>;
  static constexpr std::size_t N = linalg::space_dim_v<Space>;
  using Vec = linalg::Vector<T, N>;
 
  static constexpr auto weights() { return TraceWeights<Space>::make(); }
 
  template <AtomLike Var>
  constexpr auto operator()(const Var& X) const {
    static_assert(is_var_type_v<Var>, "Trace must operate on a variable atom");
    static_assert(
        std::same_as<dom_space_t<meta_codom_t<atom_meta_t<Var>>>, Space>,
        "Trace variable codomain must match trace space");
 
    return WeightedSum<Space>()(weights(), X);
  }
};
 
template <AtomLike Var>
constexpr auto make_trace(const Var& X) {
  using Space = dom_space_t<meta_codom_t<atom_meta_t<Var>>>;
  return Trace<Space>()(X);
}
 
template <std::size_t Offset, linalg::RnSpace TraceSpace, AtomLike Var>
constexpr auto make_trace_block(const Var& X) {
  static_assert(is_var_type_v<Var>, "Trace block must operate on a variable");
  using VarSpace = dom_space_t<meta_codom_t<atom_meta_t<Var>>>;
  using T = linalg::space_type_t<VarSpace>;
  static constexpr std::size_t VarDim = linalg::space_dim_v<VarSpace>;
  static constexpr std::size_t TraceDim = linalg::space_dim_v<TraceSpace>;
  static_assert(Offset + TraceDim <= VarDim,
                "Trace block exceeds variable dimension");
  static_assert(std::same_as<T, linalg::space_type_t<TraceSpace>>,
                "Trace block scalar type must match variable scalar type");
 
  linalg::Vector<T, VarDim> weights{};
  const auto trace_weights = Trace<TraceSpace>::weights();
  for (std::size_t i = 0; i < TraceDim; ++i)
    weights(Offset + i) = trace_weights(i);
 
  return WeightedSum<VarSpace>()(weights, X);
}
 
} // namespace convex