parallel-solver.hpp

Go to the documentation of this file.

#pragma once
 
#include "aligator/gar/riccati-base.hpp"
#include "aligator/gar/riccati-kernel.hpp"
#include "aligator/gar/lqr-problem.hpp"
#include "aligator/tracy.hpp"
 
#ifdef ALIGATOR_MULTITHREADING
namespace aligator {
namespace gar {

template <typename _Scalar>
class ParallelRiccatiSolver : public RiccatiSolverBase<_Scalar> {
public:
  using Scalar = _Scalar;
  ALIGATOR_DYNAMIC_TYPEDEFS_WITH_ROW_TYPES(Scalar);
  using Base = RiccatiSolverBase<Scalar>;
  using Kernel = ProximalRiccatiKernel<Scalar>;
  using KnotType = LqrKnotTpl<Scalar>;
  using BlkView = BlkMatrix<VectorRef, -1, 1>;
  using allocator_type = ::aligator::polymorphic_allocator;
 
  explicit ParallelRiccatiSolver(LqrProblemTpl<Scalar> &problem,
                                 const uint num_threads);
 
  bool backward(const Scalar mueq) override;
 
  inline void collapseFeedback() override {
    StageFactor<Scalar> &d = datas[0];
    RowMatrixRef K = d.fb.blockRow(0);
    RowMatrixRef Kth = d.fth.blockRow(0);
 
    // condensedSystem.subdiagonal contains the 'U' factors in the
    // block-tridiag UDUt decomposition
    // and ∂Xi+1 = -Ui+1.t ∂Xi
    auto &Up1t = condensedKktSystem.subdiagonal[1];
    K.noalias() -= Kth * Up1t;
  }
 
  struct CondensedKkt {
    using ArMat = ArenaMatrix<MatrixXs>;
    std::pmr::vector<ArMat> subdiagonal;
    std::pmr::vector<ArMat> diagonal;
    std::pmr::vector<ArMat> superdiagonal;
    // factors
    std::pmr::vector<ArMat> diagonalFacs; //< diagonal factors
    std::pmr::vector<ArMat> upFacs;       //< transposed U factors
    std::pmr::vector<BunchKaufman<ArMat>> ldlt;
    explicit CondensedKkt(const allocator_type &alloc)
        : subdiagonal{alloc}
        , diagonal{alloc}
        , superdiagonal{alloc}
        , diagonalFacs{alloc}
        , upFacs{alloc}
        , ldlt{alloc} {}
  };

  void assembleCondensedSystem(const Scalar mudyn);
 
  bool
  forward(VectorOfVectors &xs, VectorOfVectors &us, VectorOfVectors &vs,
          VectorOfVectors &lbdas,
          const std::optional<ConstVectorRef> & = std::nullopt) const override;
 
  void cycleAppend(const KnotType &knot) override;
  VectorRef getFeedforward(size_t i) override { return datas[i].ff.matrix(); }
  RowMatrixRef getFeedback(size_t i) override { return datas[i].fb.matrix(); }
 
  allocator_type get_allocator() const { return problem_->get_allocator(); }
 
  std::pmr::vector<StageFactor<Scalar>> datas;

  CondensedKkt condensedKktSystem;
  ArenaMatrix<VectorXs> condensedKktRhs, condensedKktSolution, condensedErr;
  Scalar condensedThreshold = 1e-10;
  uint maxRefinementSteps = 5u;

  uint getNumThreads() const noexcept { return numThreads; }

  void initializeTridiagSystem();
 
protected:
  uint numThreads;
  LqrProblemTpl<Scalar> *problem_;
  std::vector<long> rhsDims_;
 
  void initialize();
};
 
template <typename Scalar>
ParallelRiccatiSolver(LqrProblemTpl<Scalar> &, const uint)
    -> ParallelRiccatiSolver<Scalar>;
 
#ifdef ALIGATOR_ENABLE_TEMPLATE_INSTANTIATION
extern template class ParallelRiccatiSolver<context::Scalar>;
#endif
} // namespace gar
} // namespace aligator
#endif