graph_transform.hpp

#ifndef HAVE_GRAPH_TRANSFORM_HPP
#define HAVE_GRAPH_TRANSFORM_HPP
// Autogenerated - do not edit by hand !
#include <cstring>
#include <list>
#include <map>
#include "checkpoint.hpp"
#include "global.hpp"
#include "integrate.hpp"
#include "radix.hpp"

namespace TMBad {

template <class T>
std::vector<bool> lmatch(const std::vector<T> &x, const std::vector<T> &y) {
  std::vector<bool> ans(x.size(), false);
  for (size_t i = 0; i < x.size(); i++)
    for (size_t j = 0; j < y.size(); j++) ans[i] = ans[i] || (x[i] == y[j]);
  return ans;
}

template <class I>
std::vector<I> which(const std::vector<bool> &x) {
  std::vector<I> y;
  for (size_t i = 0; i < x.size(); i++)
    if (x[i]) y.push_back(i);
  return y;
}

std::vector<size_t> which(const std::vector<bool> &x);

template <class T>
std::vector<T> subset(const std::vector<T> &x, const std::vector<bool> &y) {
  TMBAD_ASSERT(x.size() == y.size());
  std::vector<T> ans;
  for (size_t i = 0; i < x.size(); i++)
    if (y[i]) ans.push_back(x[i]);
  return ans;
}

template <class T, class I>
std::vector<T> subset(const std::vector<T> &x, const std::vector<I> &ind) {
  std::vector<T> ans(ind.size());
  for (size_t i = 0; i < ind.size(); i++) ans[i] = x[ind[i]];
  return ans;
}

template <class T, class I>
void make_space_inplace(std::vector<T> &x, std::vector<I> &i, T space = T(0)) {
  std::vector<bool> mark(x.size(), false);
  for (size_t k = 0; k < i.size(); k++) {
    TMBAD_ASSERT(!mark[i[k]]);
    mark[i[k]] = true;
  }
  std::vector<T> x_new;
  std::vector<I> i_new;
  for (size_t k = 0; k < x.size(); k++) {
    if (mark[k]) {
      x_new.push_back(space);
      i_new.push_back(x_new.size());
    }
    x_new.push_back(x[k]);
  }
  std::swap(x, x_new);
  std::swap(i, i_new);
}

template <class T>
std::vector<T> invperm(const std::vector<T> &perm) {
  std::vector<T> iperm(perm.size());
  for (size_t i = 0; i < perm.size(); i++) iperm[perm[i]] = i;
  return iperm;
}

template <class T>
std::vector<size_t> match(const std::vector<T> &x, const std::vector<T> &y) {
  return which(lmatch(x, y));
}

size_t prod_int(const std::vector<size_t> &x);

template <class T>
std::vector<size_t> order(std::vector<T> x) {
  std::vector<std::pair<T, size_t> > y(x.size());
  for (size_t i = 0; i < x.size(); i++) {
    y[i].first = x[i];
    y[i].second = i;
  }
  sort_inplace(y);
  std::vector<size_t> z(x.size());
  for (size_t i = 0; i < x.size(); i++) {
    z[i] = y[i].second;
  }
  return z;
}

std::vector<bool> reverse_boundary(global &glob, const std::vector<bool> &vars);

std::vector<Index> get_accumulation_tree(global &glob, bool boundary = false);

std::vector<Index> find_op_by_name(global &glob, const char *name);

std::vector<Index> substitute(global &glob, const std::vector<Index> &seq,
                              bool inv_tags = true, bool dep_tags = true);

std::vector<Index> substitute(global &glob, const char *name,
                              bool inv_tags = true, bool dep_tags = true);

global accumulation_tree_split(global glob, bool sum_ = false);

void aggregate(global &glob, int sign = 1);

struct old_state {
  std::vector<Index> dep_index;
  size_t opstack_size;
  global &glob;
  old_state(global &glob);
  void restore();
};

std::vector<Index> remap_identical_sub_expressions(
    global &glob, std::vector<Index> inv_remap);
struct term_info {
  global &glob;
  std::vector<Index> id;
  std::vector<size_t> count;
  term_info(global &glob, bool do_init = true);
  void initialize(std::vector<Index> inv_remap = std::vector<Index>(0));
};

struct gk_config {
  bool debug;
  bool adaptive;
  bool nan2zero; 
  double ytol;
  double dx;
  gk_config();
};

template <class Float = ad_adapt>
struct logIntegrate_t {
  typedef Float Scalar;
  global glob;
  double mu, sigma, f_mu;
  gk_config cfg;
  double f(double x) {
    Index k = glob.inv_index.size();
    glob.value_inv(k - 1) = x;
    glob.forward();
    return glob.value_dep(0);
  }
  double g(double x) {
    return (f(x + .5 * cfg.dx) - f(x - .5 * cfg.dx)) / cfg.dx;
  }
  double h(double x) {
    return (g(x + .5 * cfg.dx) - g(x - .5 * cfg.dx)) / cfg.dx;
  }
  void rescale_integrand(const std::vector<ad_aug> &x) {
    TMBAD_ASSERT(x.size() + 1 == glob.inv_index.size());
    if (cfg.debug) Rcout << "rescale integrand:\n";
    for (size_t i = 0; i < x.size(); i++) glob.value_inv(i) = x[i].Value();
    mu = glob.value_inv(x.size());
    f_mu = f(mu);
    int i = 0;
    for (; i < 100; i++) {
      double g_mu = g(mu);
      double h_mu = h(mu);
      if (std::isfinite(f_mu) && !std::isfinite(h_mu)) {
        cfg.dx = cfg.dx * .5;
        continue;
      }
      double mu_new;
      if (h_mu < 0)
        mu_new = mu - g_mu / h_mu;
      else
        mu_new = mu + (g_mu > 0 ? cfg.dx : -cfg.dx);
      double f_mu_new = f(mu_new);
      if (cfg.debug) {
        Rcout << "mu=" << mu << " mu_new=" << mu_new << " g_mu=" << g_mu
              << " h_mu=" << h_mu << " f_mu=" << f_mu
              << " f_mu_new=" << f_mu_new << "\n";
      }
      if (f_mu_new > f_mu + cfg.ytol) {
        mu = mu_new;
        f_mu = f_mu_new;
      } else {
        break;
      }
    }
    sigma = 1. / sqrt(-h(mu));
    if (!std::isfinite(sigma)) sigma = 10000;
    if (cfg.debug)
      Rcout << "==>  i=" << i << " mu=" << mu << " f_mu=" << f_mu
            << " sigma=" << sigma << "\n";
  }

  logIntegrate_t(global &glob, gk_config cfg)
      : glob(glob), mu(0), sigma(1), f_mu(0), cfg(cfg) {
    TMBAD_ASSERT(glob.inv_index.size() >= 1);
    TMBAD_ASSERT(glob.dep_index.size() == 1);
  }
  logIntegrate_t() {}
  global::replay *p_replay;

  Float operator()(Float u) {
    Index k = glob.inv_index.size();
    p_replay->value_inv(k - 1) = sigma * u + mu;
    p_replay->forward(false, false);
    Float ans = exp(p_replay->value_dep(0) - f_mu);
    if (cfg.nan2zero && ans != ans) ans = 0;
    return ans;
  }

  std::vector<ad_aug> operator()(const std::vector<ad_aug> &x) {
    rescale_integrand(x);
    global::replay replay(this->glob, *get_glob());
    p_replay = &replay;
    replay.start();
    Index k = glob.inv_index.size();
    for (Index i = 0; i < k - 1; i++) replay.value_inv(i) = x[i];
    Float I = integrate(*this);
    Float ans = log(I) + log(sigma) + f_mu;
    replay.stop();
    return std::vector<ad_aug>(1, ans);
  }
};

template <class ADFun>
struct integrate_subgraph {
  global &glob;
  std::vector<Index> random;
  graph forward_graph;
  graph reverse_graph;
  std::vector<Index> var_remap;
  std::vector<bool> mark;
  gk_config cfg;
  integrate_subgraph(global &glob, std::vector<Index> random,
                     gk_config cfg = gk_config())
      : glob(glob),
        random(random),
        forward_graph(glob.forward_graph()),
        reverse_graph(glob.reverse_graph()),
        cfg(cfg) {
    glob.subgraph_cache_ptr();
    mark.resize(glob.opstack.size(), false);
  }
  global &try_integrate_variable(Index i) {
    const std::vector<Index> &inv2op = forward_graph.inv2op;

    Index start_node = inv2op[i];
    glob.subgraph_seq.resize(0);
    glob.subgraph_seq.push_back(start_node);
    forward_graph.search(glob.subgraph_seq);

    if (glob.subgraph_seq.size() == 1) return glob;

    bool any_marked = false;
    for (Index i = 0; i < glob.subgraph_seq.size(); i++) {
      any_marked |= mark[glob.subgraph_seq[i]];
      if (any_marked) {
        return glob;
      }
    }

    for (Index i = 0; i < glob.subgraph_seq.size(); i++) {
      mark[glob.subgraph_seq[i]] = true;
    }

    std::vector<Index> boundary = reverse_graph.boundary(glob.subgraph_seq);

    global new_glob;
    var_remap.resize(glob.values.size());
    new_glob.ad_start();
    Index total_boundary_vars = 0;
    std::vector<ad_plain> boundary_vars;
    OperatorPure *constant = glob.getOperator<global::ConstOp>();
    for (Index i = 0; i < boundary.size(); i++) {
      Index m = glob.opstack[boundary[i]]->output_size();
      for (Index j = 0; j < m; j++) {
        Index boundary_var = glob.subgraph_ptr[boundary[i]].second + j;
        var_remap[boundary_var] = total_boundary_vars;
        total_boundary_vars++;
        if (glob.opstack[boundary[i]] != constant) {
          ad_plain().Independent();
          ad_plain tmp;
          tmp.index = boundary_var;
          boundary_vars.push_back(tmp);
        } else {
          ad_plain(glob.values[boundary_var]);
        }
      }
    }
    new_glob.ad_stop();

    new_glob = glob.extract_sub(var_remap, new_glob);

    aggregate(new_glob);

    logIntegrate_t<> taped_integral(new_glob, cfg);

    glob.ad_start();
    std::vector<ad_aug> boundary_vars2(boundary_vars.begin(),
                                       boundary_vars.end());
    if (cfg.adaptive) {
      typedef retaping_derivative_table<logIntegrate_t<>, ADFun> DTab;
      global::Complete<AtomOp<DTab> > taped_integral_operator(taped_integral,
                                                              boundary_vars2);
      taped_integral_operator(boundary_vars)[0].Dependent();
    } else {
      taped_integral(boundary_vars2)[0].Dependent();
    }
    glob.ad_stop();
    return glob;
  }
  global &gk() {
    for (Index i = 0; i < random.size(); i++) {
      try_integrate_variable(random[i]);
    }

    std::vector<bool> keep_node = mark;
    keep_node.flip();

    keep_node.resize(glob.opstack.size(), true);

    std::vector<Index> v2o = glob.var2op();
    for (Index i = 0; i < glob.inv_index.size(); i++) {
      keep_node[v2o[glob.inv_index[i]]] = true;
    }

    glob.subgraph_seq.resize(0);
    for (Index i = 0; i < keep_node.size(); i++) {
      if (keep_node[i]) glob.subgraph_seq.push_back(i);
    }

    glob = glob.extract_sub();
    return glob;
  }
};

struct multivariate_index {
 private:
  std::vector<size_t> x;
  std::vector<bool> mask_;
  size_t pointer;
  std::vector<size_t> bound;

 public:
  size_t count();
  multivariate_index(size_t bound_, size_t dim, bool flag = true);
  multivariate_index(std::vector<size_t> bound, bool flag = true);
  void flip();
  multivariate_index &operator++();
  operator size_t();
  size_t index(size_t i);
  std::vector<size_t> index();
  std::vector<bool>::reference mask(size_t i);
  void set_mask(const std::vector<bool> &mask);
};

struct clique {
  std::vector<Index> indices;
  std::vector<ad_aug> logsum;
  std::vector<size_t> dim;
  size_t clique_size();
  clique();
  void subset_inplace(const std::vector<bool> &mask);
  void logsum_init();
  bool empty() const;
  bool contains(Index i);
  void get_stride(const clique &super, Index ind, std::vector<ad_plain> &offset,
                  Index &stride);
};

struct sr_grid {
  std::vector<Scalar> x;
  std::vector<Scalar> w;
  sr_grid();

  sr_grid(Scalar a, Scalar b, size_t n);

  sr_grid(size_t n);
  size_t size();

  std::vector<ad_plain> logw;
  ad_plain logw_offset();
};

struct sequential_reduction {
  std::list<clique> cliques;
  std::vector<sr_grid> grid;
  std::vector<Index> inv2grid;
  global &glob;
  global new_glob;
  std::vector<Index> random;
  global::replay replay;
  std::vector<bool> mark;
  graph forward_graph;
  graph reverse_graph;
  std::vector<Index> var_remap;
  const static Index NA = -1;
  std::vector<Index> op2inv_idx;
  std::vector<Index> op2dep_idx;
  std::vector<bool> terms_done;
  term_info tinfo;
  std::map<size_t, std::vector<ad_aug> > cache;
  sequential_reduction(global &glob, std::vector<Index> random,
                       std::vector<sr_grid> grid =
                           std::vector<sr_grid>(1, sr_grid(-20, 20, 200)),
                       std::vector<Index> random2grid = std::vector<Index>(0),
                       bool perm = true);
  void reorder_random();

  std::vector<size_t> get_grid_bounds(std::vector<Index> inv_index);

  std::vector<sr_grid *> get_grid(std::vector<Index> inv_index);
  std::vector<ad_aug> tabulate(std::vector<Index> inv_index, Index dep_index);

  void merge(Index i);

  void update(Index i);
  void show_cliques();
  void update_all();
  ad_aug get_result();
  global marginal();
};

struct autopar {
  global &glob;
  graph reverse_graph;
  size_t num_threads;
  bool do_aggregate;
  bool keep_all_inv;
  std::vector<std::vector<Index> > node_split;
  std::vector<std::vector<Index> > inv_idx;
  std::vector<std::vector<Index> > dep_idx;
  std::vector<global> vglob;
  autopar(global &glob, size_t num_threads);
  std::vector<size_t> max_tree_depth();

  template <class T>
  size_t which_min(const std::vector<T> &x) {
    return std::min_element(x.begin(), x.end()) - x.begin();
  }

  void run();
  void extract();
  size_t input_size() const;
  size_t output_size() const;
};

struct ParalOp : global::DynamicOperator<-1, -1> {
  static const bool have_input_size_output_size = true;
  std::vector<global> vglob;
  std::vector<std::vector<Index> > inv_idx;
  std::vector<std::vector<Index> > dep_idx;

  Index n, m;
  Index input_size() const;
  Index output_size() const;
  ParalOp(const autopar &ap);

  template <class T>
  void reverse(ReverseArgs<T> &args) {
    bool parallel_operator_used_with_valid_type = false;
    TMBAD_ASSERT(parallel_operator_used_with_valid_type);
  }
  static const bool add_forward_replay_copy = true;
  template <class T>
  void forward(ForwardArgs<T> &args) {
    bool parallel_operator_used_with_valid_type = false;
    TMBAD_ASSERT(parallel_operator_used_with_valid_type);
  }

  void forward(ForwardArgs<Scalar> &args);
  void reverse(ReverseArgs<Scalar> &args);
  const char *op_name();
  void print(global::print_config cfg);
};

std::vector<Index> get_likely_expression_duplicates(
    const global &glob, std::vector<Index> inv_remap);

bool all_allow_remap(const global &glob);

template <class T>
struct forbid_remap {
  T &remap;
  forbid_remap(T &remap) : remap(remap) {}
  void operator()(Index a, Index b) {
    bool ok = true;
    for (Index i = a + 1; i <= b; i++) {
      ok &= (remap[i] - remap[i - 1] == 1);
    }
    if (ok) {
      return;
    } else {
      for (Index i = a; i <= b; i++) {
        remap[i] = i;
      }
    }
  }
};

std::vector<Index> remap_identical_sub_expressions(
    global &glob, std::vector<Index> inv_remap);

void remap_identical_sub_expressions(global &glob);

std::vector<Position> inv_positions(global &glob);

void reorder_graph(global &glob, std::vector<Index> inv_idx);

}  // namespace TMBad
#endif  // HAVE_GRAPH_TRANSFORM_HPP