equality_coupling.hpp

Go to the documentation of this file.

/*
 * Copyright © 2025-2026 Akreon LLC
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include "akreon/linalg/linalg_util.hpp"
 
namespace convex {
 
template <linalg::VectorContainer Var, linalg::VectorContainer NuVar>
struct NoEqualityCoupling {
  constexpr auto solve_step(const auto& M_inv, const Var& pz_rhs, const Var& rd,
                            const NuVar& /* req */) const {
    const auto [pz, event] = M_inv.solve(pz_rhs + rd);
    return std::tuple(pz, NuVar{}, event);
  }
};
 
template <linalg::MatrixContainer EqMat, typename SchurFactor,
          linalg::VectorContainer Var, linalg::VectorContainer NuVar>
struct SchurEqualityCoupling {
  constexpr explicit SchurEqualityCoupling(const EqMat& A_in,
                                           const SchurFactor& S_inv_in)
      : A(A_in), S_inv(S_inv_in) {}
 
  constexpr auto solve_step(const auto& M_inv, const Var& pz_rhs, const Var& rd,
                            const NuVar& req) const {
    const auto [M_inv_pz, event_one] = M_inv.solve(pz_rhs + rd);
    if (linalg::is_fatal_event(event_one))
      return std::tuple(Var{}, NuVar{}, event_one);
 
    const auto schur_rhs = A * M_inv_pz - req;
    const auto [pnu, event_two] = S_inv.solve(schur_rhs);
    if (linalg::is_fatal_event(event_two))
      return std::tuple(Var{}, NuVar{}, event_two);
 
    const auto [pz, event_three] =
        M_inv.solve(pz_rhs + rd - linalg::transpose(A) * pnu);
    if (linalg::is_fatal_event(event_three))
      return std::tuple(Var{}, NuVar{}, event_three);
 
    return std::tuple(pz, pnu, linalg::NumericalEvent::NONE);
  }
 
  EqMat A{};         
  SchurFactor S_inv; 
};
 
template <linalg::MatrixContainer EqMat, linalg::Field Scalar,
          linalg::VectorContainer Var, linalg::VectorContainer NuVar>
struct ScalarEqualityCoupling {
  constexpr explicit ScalarEqualityCoupling(const EqMat& A_in,
                                            const Scalar& s_in)
      : A(A_in), s(s_in) {}
 
  constexpr auto solve_step(const auto& M_inv, const Var& pz_rhs, const Var& rd,
                            const NuVar& req) const {
    const auto [M_inv_pz, event_one] = M_inv.solve(pz_rhs + rd);
    if (linalg::is_fatal_event(event_one))
      return std::tuple(Var{}, NuVar{}, event_one);
 
    const auto schur_rhs = A * M_inv_pz - req;
    NuVar pnu{};
    pnu(0) = schur_rhs(0) / s;
 
    const auto [pz, event_two] =
        M_inv.solve(pz_rhs + rd - linalg::transpose(A) * pnu);
    if (linalg::is_fatal_event(event_two))
      return std::tuple(Var{}, NuVar{}, event_two);
 
    return std::tuple(pz, pnu, linalg::NumericalEvent::NONE);
  }
 
  EqMat A{};  
  Scalar s{}; 
};
 
} // namespace convex