stage-model.hpp

Go to the documentation of this file.

#pragma once
 
#include "aligator/core/function-abstract.hpp"
#include "aligator/core/constraint.hpp"
#include <fmt/format.h>
 
namespace aligator {
using xyz::polymorphic;
 
#define ALIGATOR_CHECK_DERIVED_CLASS(Base, Derived)                            \
  static_assert((std::is_base_of_v<Base, Derived>),                            \
                "Failed check for derived class.")

template <typename _Scalar> struct StageModelTpl {
public:
  using Scalar = _Scalar;
  ALIGATOR_DYNAMIC_TYPEDEFS(Scalar);
 
  using Manifold = ManifoldAbstractTpl<Scalar>;
  using PolyManifold = polymorphic<Manifold>;
  using Dynamics = ExplicitDynamicsModelTpl<Scalar>;
  using PolyDynamics = polymorphic<Dynamics>;
  using PolyFunction = polymorphic<StageFunctionTpl<Scalar>>;
  using PolyConstraintSet = polymorphic<ConstraintSetTpl<Scalar>>;
  using Cost = CostAbstractTpl<Scalar>;
  using PolyCost = polymorphic<Cost>;
  using Data = StageDataTpl<Scalar>;

  PolyManifold xspace_;
  PolyManifold xspace_next_;
  PolyManifold uspace_;
  ConstraintStackTpl<Scalar> constraints_;
  PolyCost cost_;
  PolyDynamics dynamics_;

  template <typename U> U *getCost() {
    ALIGATOR_CHECK_DERIVED_CLASS(Cost, U);
    return dynamic_cast<U *>(&*cost_);
  }

  template <typename U> const U *getCost() const {
    ALIGATOR_CHECK_DERIVED_CLASS(Cost, U);
    return dynamic_cast<const U *>(&*cost_);
  }

  template <typename U> U *getDynamics() {
    ALIGATOR_CHECK_DERIVED_CLASS(Dynamics, U);
    return dynamic_cast<U *>(&*dynamics_);
  }

  template <typename U> const U *getDynamics() const {
    ALIGATOR_CHECK_DERIVED_CLASS(Dynamics, U);
    return dynamic_cast<const U *>(&*dynamics_);
  }

  StageModelTpl(const PolyCost &cost, const PolyDynamics &dynamics);
  virtual ~StageModelTpl() = default;
 
  const Manifold &xspace() const { return *xspace_; }
  const Manifold &uspace() const { return *uspace_; }
  const Manifold &xspace_next() const { return *xspace_next_; }
 
  const Cost &cost() const { return *cost_; }
  virtual bool hasDynModel() const { return true; }
 
  int nx1() const { return xspace_->nx(); }
  int ndx1() const { return xspace_->ndx(); }
  int nu() const { return uspace_->ndx(); }
  int nx2() const { return xspace_next_->nx(); }
  int ndx2() const { return xspace_next_->ndx(); }
  int nc() const { return (int)constraints_.totalDim(); }

  std::size_t numConstraints() const { return constraints_.size(); }

  int numDual() const { return ndx2() + nc(); }

  void addConstraint(const PolyFunction &func,
                     const PolyConstraintSet &cstr_set);
 
  /* Evaluate costs, constraints, ... */

  virtual void evaluate(const ConstVectorRef &x, const ConstVectorRef &u,
                        Data &data) const;

  virtual void computeFirstOrderDerivatives(const ConstVectorRef &x,
                                            const ConstVectorRef &u,
                                            Data &data) const;

  virtual void computeSecondOrderDerivatives(const ConstVectorRef &x,
                                             const ConstVectorRef &u,
                                             Data &data) const;

  virtual shared_ptr<Data> createData() const;
};
 
#ifdef ALIGATOR_ENABLE_TEMPLATE_INSTANTIATION
extern template struct StageModelTpl<context::Scalar>;
#endif
 
} // namespace aligator
 
template <typename Scalar>
struct fmt::formatter<aligator::StageModelTpl<Scalar>> {
  constexpr auto parse(format_parse_context &ctx) const
      -> decltype(ctx.begin()) {
    return ctx.end();
  }
 
  auto format(const aligator::StageModelTpl<Scalar> &stage,
              format_context &ctx) const -> decltype(ctx.out()) {
    if (stage.ndx1() == stage.ndx2()) {
      return fmt::format_to(ctx.out(),
                            "StageModel {{"
                            "\n  ndx:   {:d},"
                            "\n  nu:    {:d},"
                            "\n  nc:    {:d}, [{:d} constraints]"
                            "\n}}",
                            stage.ndx1(), stage.nu(), stage.nc(),
                            stage.numConstraints());
    } else {
      return fmt::format_to(ctx.out(),
                            "StageModel {{"
                            "\n  ndx1:  {:d},"
                            "\n  nu:    {:d},"
                            "\n  nc:    {:d}, [{:d} constraints]"
                            "\n  ndx2:  {:d},"
                            "\n}}",
                            stage.ndx1(), stage.nu(), stage.nc(),
                            stage.numConstraints(), stage.ndx2());
    }
  }
};
 
namespace aligator {
template <typename Scalar>
std::ostream &operator<<(std::ostream &oss, const StageModelTpl<Scalar> &sm) {
  return oss << fmt::format("{}", sm);
}
} // namespace aligator