canonical_form_builder.hpp

Go to the documentation of this file.

/*
 * Copyright © 2025-2026 Akreon LLC
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <algorithm>
#include "cone_block.hpp"
 
#include "akreon/linalg/affine/affine_map.hpp"
#include "akreon/linalg/assert.hpp"
#include "akreon/convex/canonical_emission.hpp"
#include "akreon/convex/cone_traits.hpp"
#include "akreon/convex/convex_traits.hpp"
 
namespace conic {
 
template <linalg::RnSpace DecVarDom, affine::AffineMapLike CoordMap,
          conic::ConeBlockLike CB, affine::AffineMapLike EqMap,
          typename CanonicalTrace = convex::EmptyCanonicalTrace>
struct CanonicalFormBuilder {
  using CoordMapT = affine::affine_mat_t<CoordMap>;
  using T = linalg::matrix_scalar_t<CoordMapT>;
 
  static constexpr std::size_t N = linalg::space_dim_v<DecVarDom>;
 
  constexpr explicit CanonicalFormBuilder(const CoordMap& coord_in,
                                          const CB& cb_in, const EqMap& eq_in,
                                          const CanonicalTrace& trace_in = {})
      : coord_map(coord_in), cb(cb_in), eq_map(eq_in),
        canonical_trace(trace_in) {}
 
  template <std::size_t SharedCols, affine::AffineMapLike LiftedMap,
            typename... Cs>
  constexpr auto make(int idx, const LiftedMap& lift_map,
                      const std::tuple<Cs...>& cones) const {
    const auto new_map = this->template update_map<SharedCols>();
    const auto new_eq_map = this->template update_eq_map<SharedCols>();
    static constexpr std::size_t FullDom =
        linalg::mat_cols_v<decltype(new_map.linear_part())>;
    constexpr std::size_t FreshCols =
        linalg::mat_cols_v<CoordMapT> - SharedCols;
    static_assert(FreshCols > 0, "Lifted emission must append an auxiliary");
 
    linalg::Vector<T, FullDom> sel{};
    linalg::Vector<T, 1> bias{lift_map.translation()(0)};
 
    AK_ASSERT(idx < 0 && static_cast<std::size_t>(-idx) <= FreshCols &&
              "Aux variable selector must select a fresh column in "
              "CanonicalFormBuilder::make");
 
    const std::size_t selector_idx = FullDom - static_cast<std::size_t>(-idx);
    sel(selector_idx) = lift_map.linear_part()(0, 0);
    const linalg::RowVector<T, FullDom> sel_row = linalg::as_row_vector(sel);
    const auto new_coord_map =
        affine::AffineMap<decltype(sel_row)>(sel_row, bias);
    const auto new_cb =
        conic::ConeBlock(new_map, std::tuple_cat(cb.cones, cones));
    return make_builder<DecVarDom>(new_coord_map, new_cb, new_eq_map,
                                   canonical_trace);
  }
 
  template <std::size_t SharedCols, typename... Cs>
  constexpr auto make(const std::tuple<Cs...>& cones) const {
    const auto new_map = this->template update_map<SharedCols>();
    const auto new_eq_map = this->template update_eq_map<SharedCols>();
    const auto new_cb =
        conic::ConeBlock(new_map, std::tuple_cat(cb.cones, cones));
    return make_builder<DecVarDom>(new_map, new_cb, new_eq_map,
                                   canonical_trace);
  }
 
  template <convex::AtomType Type, affine::AffineMapLike LiftedMap,
            LocalEmissionLike Emission>
  constexpr auto realize(int idx, const LiftedMap& lift_map,
                         const Emission& emission) const {
    // The child selector defines the shared prefix. AtomType affects tracing,
    // never column ownership or placement.
    constexpr std::size_t ChildCols = linalg::mat_cols_v<CoordMapT>;
    const auto traced = add_cone_trace<Type, Emission>(emission);
    const auto local_builder = make_builder<DecVarDom>(
        emission.local_coord_map, cb, emission.local_eq_map, traced);
    return local_builder.template make<ChildCols>(idx, lift_map,
                                                  emission.local_cones);
  }
 
  template <convex::AtomType Type, LocalEmissionLike Emission>
  constexpr auto realize(const Emission& emission) const {
    const auto traced = add_cone_trace<Type, Emission>(emission);
    const auto local_builder = make_builder<DecVarDom>(
        emission.local_coord_map, cb, emission.local_eq_map, traced);
    constexpr std::size_t SharedCols =
        std::max(N, linalg::mat_cols_v<CoordMapT>);
    return local_builder.template make<SharedCols>(emission.local_cones);
  }
 
  template <LocalEqualityEmissionLike Emission>
  constexpr auto realize(const Emission& emission) const {
    using CurrentEqMat = affine::affine_mat_t<EqMap>;
    using LocalEqMap = emission_local_eq_map_t<Emission>;
    using LocalEqMat = affine::affine_mat_t<LocalEqMap>;
 
    static_assert(std::same_as<T, linalg::matrix_scalar_t<LocalEqMat>>,
                  "type mismatch between equality emission and builder");
 
    static constexpr std::size_t CurrentRows = linalg::mat_rows_v<CurrentEqMat>;
    static constexpr std::size_t CurrentCols = linalg::mat_cols_v<CurrentEqMat>;
    static constexpr std::size_t LocalRows = linalg::mat_rows_v<LocalEqMat>;
    static constexpr std::size_t LocalCols = linalg::mat_cols_v<LocalEqMat>;
    static constexpr std::size_t FullRows = CurrentRows + LocalRows;
    static constexpr std::size_t FullCols =
        std::max({N, CurrentCols, LocalCols, cone_block_map_cols_v<CB>,
                  linalg::mat_cols_v<CoordMapT>});
 
    linalg::Matrix<T, FullRows, FullCols> A{};
    linalg::Vector<T, FullRows> b{};
 
    if constexpr (CurrentRows != 0 && CurrentCols != 0) {
      A = linalg::set_block<0, CurrentRows, 0, CurrentCols>(
          A, eq_map.linear_part());
      for (std::size_t i = 0; i < CurrentRows; ++i)
        b(i) = eq_map.translation()(i);
    }
 
    if constexpr (LocalRows != 0 && LocalCols != 0) {
      A = linalg::set_block<CurrentRows, FullRows, 0, LocalCols>(
          A, emission.local_eq_map.linear_part());
      for (std::size_t i = 0; i < LocalRows; ++i)
        b(i + CurrentRows) = emission.local_eq_map.translation()(i);
    }
 
    const auto new_eq_map = affine::AffineMap<decltype(A)>(A, b);
    const auto traced = add_equality_trace(emission);
    return make_builder<DecVarDom>(coord_map, cb, new_eq_map, traced);
  }
 
  CoordMap coord_map; 
  CB cb;              
  EqMap eq_map;       
  CanonicalTrace canonical_trace; 
 
private:
  template <convex::AtomType Type, LocalEmissionLike Emission>
  constexpr auto add_cone_trace(const Emission& /*emission*/) const {
    using ConeTuple = emission_cones_t<Emission>;
    using EqMapT = affine::affine_mat_t<EqMap>;
 
    // The builder owns global offsets: atom emissions only describe local
    // cones/equalities, while the builder knows how many rows were realized
    // before this emission.
    constexpr std::size_t NodeId =
        convex::canonical_trace_num_nodes_v<CanonicalTrace>;
    constexpr std::size_t GlobalConeStart = conic::cone_block_num_cones_v<CB>;
    constexpr std::size_t ConeOffsetStart = conic::cone_block_map_rows_v<CB>;
    constexpr std::size_t LocalEqRows = emission_local_eq_rows_v<Emission>;
    constexpr std::size_t EqOffsetStart = linalg::mat_rows_v<EqMapT>;
 
    const auto nodes = std::make_tuple(convex::CanonicalNodeRecord{
        NodeId, convex::canonical_no_parent, Type, 0});
    const auto cones = convex::make_canonical_cone_records<ConeTuple>(
        NodeId, GlobalConeStart, ConeOffsetStart);
 
    if constexpr (LocalEqRows == 0) {
      return convex::extend_canonical_trace(canonical_trace, nodes, cones,
                                            std::make_tuple());
    } else {
      const auto equalities = std::make_tuple(convex::CanonicalEqualityRecord{
          NodeId, 0, EqOffsetStart, LocalEqRows});
      return convex::extend_canonical_trace(canonical_trace, nodes, cones,
                                            equalities);
    }
  }
 
  template <LocalEqualityEmissionLike Emission>
  constexpr auto add_equality_trace(const Emission& /*emission*/) const {
    using EqMapT = affine::affine_mat_t<EqMap>;
 
    constexpr std::size_t NodeId =
        convex::canonical_trace_num_nodes_v<CanonicalTrace>;
    constexpr std::size_t LocalEqRows = emission_local_eq_rows_v<Emission>;
    constexpr std::size_t EqOffsetStart = linalg::mat_rows_v<EqMapT>;
 
    const auto nodes = std::make_tuple(convex::CanonicalNodeRecord{
        NodeId, convex::canonical_no_parent, convex::AtomType::SET, 0});
 
    if constexpr (LocalEqRows == 0) {
      return convex::extend_canonical_trace(
          canonical_trace, nodes, std::make_tuple(), std::make_tuple());
    } else {
      const auto equalities = std::make_tuple(convex::CanonicalEqualityRecord{
          NodeId, 0, EqOffsetStart, LocalEqRows});
      return convex::extend_canonical_trace(canonical_trace, nodes,
                                            std::make_tuple(), equalities);
    }
  }
 
  template <std::size_t SharedCols>
  struct ColumnLayout {
    static constexpr std::size_t CoordCols = linalg::mat_cols_v<CoordMapT>;
 
    constexpr ColumnLayout()
        : prefix_cols(SharedCols),
          auxiliary_start(
              std::max({SharedCols, cone_block_map_cols_v<CB>,
                        linalg::mat_cols_v<affine::affine_mat_t<EqMap>>})),
          full_cols(auxiliary_start + CoordCols - SharedCols) {
      static_assert(SharedCols >= N && CoordCols >= SharedCols,
                    "Emission must preserve the decision/shared prefix");
    }
 
    const std::size_t
        prefix_cols; 
    const std::size_t
        auxiliary_start;         
    const std::size_t full_cols; 
  };
 
  template <std::size_t SharedCols>
  constexpr auto update_map() const {
    const auto& cone_map = cb.map;
    using ConeMapT = affine::affine_mat_t<decltype(cone_map)>;
 
    static_assert(std::same_as<T, linalg::matrix_scalar_t<ConeMapT>>,
                  "type mismatch between cone map and coord map");
 
    static constexpr std::size_t CoordDom = linalg::mat_cols_v<CoordMapT>;
    static constexpr std::size_t CoordCodom = linalg::mat_rows_v<CoordMapT>;
 
    static constexpr std::size_t ConeDom = linalg::mat_cols_v<ConeMapT>;
    static constexpr std::size_t ConeCodom = linalg::mat_rows_v<ConeMapT>;
    constexpr ColumnLayout<SharedCols> layout{};
    constexpr std::size_t PrefixCols = layout.prefix_cols;
    constexpr std::size_t CoordAux = CoordDom - PrefixCols;
    constexpr std::size_t FullDom = layout.full_cols;
    constexpr std::size_t FullCodom = ConeCodom + CoordCodom;
    constexpr std::size_t CoordAuxStartCol = layout.auxiliary_start;
    constexpr std::size_t CoordAuxEndCol = CoordAuxStartCol + CoordAux;
 
    linalg::Matrix<T, FullCodom, FullDom> F{};
    linalg::Vector<T, FullCodom> g{};
 
    if constexpr (ConeDom != 0) {
      F = linalg::set_block<0, ConeCodom, 0, ConeDom>(F,
                                                      cone_map.linear_part());
      for (std::size_t i = 0; i < ConeCodom; ++i)
        g(i) = cone_map.translation()(i);
    }
 
    const auto A = coord_map.linear_part();
    const auto dec_spanning_blk =
        linalg::get_block<0, CoordCodom, 0, PrefixCols>(A);
    F = linalg::set_block<ConeCodom, FullCodom, 0, PrefixCols>(
        F, dec_spanning_blk);
 
    if constexpr (CoordAux > 0) {
      static_assert(CoordAuxEndCol > CoordAuxStartCol);
      const auto coord_aux_blk =
          linalg::get_block<0, CoordCodom, PrefixCols, CoordDom>(A);
      F = linalg::set_block<ConeCodom, FullCodom, CoordAuxStartCol,
                            CoordAuxEndCol>(F, coord_aux_blk);
    }
 
    for (std::size_t i = 0; i < CoordCodom; ++i)
      g(i + ConeCodom) = coord_map.translation()(i);
 
    return affine::AffineMap<decltype(F)>(F, g);
  }
 
  template <std::size_t SharedCols>
  constexpr auto update_eq_map() const {
    using EqMapT = affine::affine_mat_t<EqMap>;
 
    static_assert(std::same_as<T, linalg::matrix_scalar_t<EqMapT>>,
                  "type mismatch between equality map and coord map");
 
    static constexpr std::size_t EqDom = linalg::mat_cols_v<EqMapT>;
    static constexpr std::size_t EqCodom = linalg::mat_rows_v<EqMapT>;
    constexpr ColumnLayout<SharedCols> layout{};
    constexpr std::size_t FullDom = layout.full_cols;
 
    linalg::Matrix<T, EqCodom, FullDom> A{};
    linalg::Vector<T, EqCodom> b{};
 
    if constexpr (EqCodom != 0 && EqDom != 0) {
      A = linalg::set_block<0, EqCodom, 0, EqDom>(A, eq_map.linear_part());
      for (std::size_t i = 0; i < EqCodom; ++i) b(i) = eq_map.translation()(i);
    }
 
    // Equality coefficients always reference existing global columns.
    // Fresh cone auxiliaries extend their width without rebasing any entries.
    return affine::AffineMap<decltype(A)>(A, b);
  }
};
 
template <linalg::RnSpace DecVarDom, affine::AffineMapLike CoordMap,
          conic::ConeBlockLike CB, affine::AffineMapLike EqMap,
          typename CanonicalTrace = convex::EmptyCanonicalTrace>
constexpr auto make_builder(const CoordMap& coord_map, const CB& cb,
                            const EqMap& eq_map,
                            const CanonicalTrace& canonical_trace = {}) {
  return CanonicalFormBuilder<DecVarDom, CoordMap, CB, EqMap, CanonicalTrace>(
      coord_map, cb, eq_map, canonical_trace);
}
 
template <linalg::Field T>
constexpr auto create_empty_cblk() {
  constexpr auto empty = std::make_tuple();
  constexpr conic::ConeBlock cb(
      affine::AffineMap(linalg::Matrix<T, 0, 0>{}, linalg::Vector<T, 0>{}),
      empty);
  return cb;
}
 
template <linalg::Field T>
constexpr auto create_empty_eq_map() {
  return affine::AffineMap(linalg::Matrix<T, 0, 0>{}, linalg::Vector<T, 0>{});
}
 
template <linalg::RnSpace DecVarDom>
constexpr auto create_empty_cblk_builder() {
  using T = linalg::space_type_t<DecVarDom>;
  const auto cblk = create_empty_cblk<T>();
  const auto eq_map = create_empty_eq_map<T>();
  return make_builder<DecVarDom>(cblk.map, cblk, eq_map);
}
 
} // namespace conic