checkpoint.hpp

#ifndef HAVE_CHECKPOINT_HPP
#define HAVE_CHECKPOINT_HPP
// Autogenerated - do not edit by hand !
#include <memory>
#include "global.hpp"
#include "vectorize.hpp"

namespace TMBad {

template <class ADFun, bool packed_ = false>
struct standard_derivative_table : std::vector<ADFun> {
  static const bool packed = packed_;
  void requireOrder(size_t n) {
    while ((*this).size() <= n) {
      (*this).push_back((*this).back().WgtJacFun());
    }
  }
  void retape(ForwardArgs<Scalar> &args) {}
  standard_derivative_table(const ADFun &F) : std::vector<ADFun>(1, F) {}
};

struct ParametersChanged {
  static const bool packed = false;
  std::vector<Scalar> x_prev;
  bool operator()(const std::vector<Scalar> &x);
};

template <class Functor, class ADFun, class Test = ParametersChanged,
          bool packed_ = false>
struct retaping_derivative_table : standard_derivative_table<ADFun, packed_> {
  Functor F;
  Test test;
  void retape(ForwardArgs<Scalar> &args) {
    size_t n = (*this)[0].Domain();
    std::vector<Scalar> x = args.x_segment(0, n);
    bool change = test(x);
    if (change) {
      (*this).resize(1);
      (*this)[0] = ADFun(F, x);
    }
  }
  template <class V>
  retaping_derivative_table(const Functor &F, const V &x, Test test = Test())
      : standard_derivative_table<ADFun, packed_>(ADFun(F, x)),
        F(F),
        test(test) {}
};

template <class T>
struct omp_shared_ptr {
  typedef std::shared_ptr<T> Base;
  Base sp;
  std::shared_ptr<std::vector<std::weak_ptr<T> > > weak_refs;

  omp_shared_ptr(const Base &x)
      : sp(x), weak_refs(std::make_shared<std::vector<std::weak_ptr<T> > >()) {
    (*weak_refs).resize(TMBAD_MAX_NUM_THREADS);
    (*weak_refs)[TMBAD_THREAD_NUM] = x;
  }
  omp_shared_ptr(const omp_shared_ptr &other) : weak_refs(other.weak_refs) {
    if ((*weak_refs)[TMBAD_THREAD_NUM].expired()) {
      sp = std::make_shared<T>(*other);

      (*weak_refs)[TMBAD_THREAD_NUM] = sp;
    } else {
      sp = (*weak_refs)[TMBAD_THREAD_NUM].lock();
    }
  }
  omp_shared_ptr() {}
  T &operator*() const { return *sp; }
  T *operator->() const { return sp.get(); }
  explicit operator bool() const { return (bool)sp; }
};

template <class DerivativeTable>
struct AtomOp : global::DynamicOperator<-1, -1> {
  static const bool have_input_size_output_size = true;
  static const bool add_forward_replay_copy = true;

  TMBAD_SHARED_PTR<DerivativeTable> dtab;

  int order;

  template <class T1>
  AtomOp(const T1 &F) : dtab(std::make_shared<DerivativeTable>(F)), order(0) {}
  template <class T1, class T2>
  AtomOp(const T1 &F, const T2 &x)
      : dtab(std::make_shared<DerivativeTable>(F, x)), order(0) {}
  template <class T1, class T2, class T3>
  AtomOp(const T1 &F, const T2 &x, const T3 &t)
      : dtab(std::make_shared<DerivativeTable>(F, x, t)), order(0) {}

  Index input_size() const { return (*dtab)[order].Domain(); }
  Index output_size() const { return (*dtab)[order].Range(); }

  void forward(ForwardArgs<Scalar> &args) {
    (*dtab).retape(args);

    (*dtab).requireOrder(order);

    size_t n = input_size();
    size_t m = output_size();

    auto x = args.x_segment(0, n);

    args.y_segment(0, m) = (*dtab)[order](x);
  }

  void reverse(ReverseArgs<Scalar> &args) {
    size_t n = input_size();
    size_t m = output_size();

    auto x = args.x_segment(0, n);
    auto w = args.dy_segment(0, m);

    args.dx_segment(0, n) += (*dtab)[order].Jacobian(x, w);
  }

  void reverse(ReverseArgs<global::Replay> &args) {
    size_t n = input_size();
    size_t m = output_size();

    std::vector<global::Replay> x = args.x_segment(0, n);
    if (DerivativeTable::packed) x = repack(x);
    std::vector<global::Replay> w = args.dy_segment(0, m);
    std::vector<global::Replay> xw;
    xw.insert(xw.end(), x.begin(), x.end());
    xw.insert(xw.end(), w.begin(), w.end());

    (*dtab).requireOrder(order + 1);
    AtomOp cpy(*this);
    cpy.order++;
    args.dx_segment(0, n) += global::Complete<AtomOp>(cpy)(xw);
  }

  template <class T>
  void forward(ForwardArgs<T> &args) {
    TMBAD_ASSERT(false);
  }
  void reverse(ReverseArgs<Writer> &args) { TMBAD_ASSERT(false); }

  const char *op_name() { return "AtomOp"; }

  void print(global::print_config cfg) {
    Rcout << cfg.prefix;
    Rcout << "order=" << order << " ";
    Rcout << "(*dtab).size()=" << (*dtab).size() << " ";
    Rcout << "dtab=" << &(*dtab) << "\n";
    (*dtab)[order].print(cfg);
  }
};

template <class Functor>
struct PackWrap {
  Functor F;
  PackWrap(const Functor &F) : F(F) {}
  template <class T>
  std::vector<T> operator()(const std::vector<T> &xp) {
    Index K = ScalarPack<SegmentRef>::size;
    size_t n = xp.size() / K;
    TMBAD_ASSERT2(n * K == xp.size(), "Invalid packed arguments");
    std::vector<ad_segment> x(n);
    for (size_t i = 0; i < n; i++) x[i] = unpack(xp, i);
    ad_segment y = F(x);
    ad_segment yp = pack(y);
    std::vector<T> ans = concat(std::vector<ad_segment>(1, yp));
    return ans;
  }
  bool operator()(const std::vector<Scalar> &xp) {
    Index K = ScalarPack<SegmentRef>::size;
    size_t n = xp.size() / K;
    TMBAD_ASSERT2(n * K == xp.size(), "Invalid packed arguments");
    std::vector<Scalar *> x(n);
    for (size_t i = 0; i < n; i++) x[i] = unpack(xp, i);
    return F(x);
  }
};

}  // namespace TMBad
#endif  // HAVE_CHECKPOINT_HPP