multidim.h

// -*- C++ -*-

// A simple and sane multidimensional array class.

#ifndef multidim_
#define multidim_

#define MDSTR(X) #X
#define MDSTR1(X) MDSTR(X)
#define MDCHECK(X)                                          \
  while (!(X)) {                                            \
    THROW("FAILED: " __FILE__ ":" MDSTR1(__LINE__) ":" #X); \
  }

#include <stdlib.h>

namespace multidim {
template <class T>
struct mdarray {
  static const int MAXRANK = 8;
  int dims[MAXRANK + 1] = {0, 0};
  int total = 0;
  int fill = 0;
  T *data = 0;
  bool owned = false;

  void take(mdarray<T> &other) {
    for (int i = 0; i < MAXRANK + 1; i++) dims[i] = other.dims[i];
    total = other.total;
    fill = other.fill;
    data = other.data;
    owned = other.owned;
    for (int i = 0; i < MAXRANK + 1; i++) other.dims[i] = 0;
    other.total = 0;
    other.fill = 0;
    other.data = 0;
    other.owned = false;
  }

  // cleared array by default
  mdarray() {}

  // deallocate data on deletion
  ~mdarray() { clear(); }

  // create an array of the given size
  template <typename... Args>
  mdarray(Args... args) {
    resize(args...);
  }

  // alias an array to a given pointer
  template <typename... Args>
  mdarray(T *p, int *dims) {
    alias_(p, dims);
  }

  // no copy constructors or assignment
  mdarray(mdarray<T> &) = delete;
  mdarray(const mdarray<T> &) = delete;

  // clear all allocated data
  void clear() {
    if (owned && data) delete[] data;
    data = 0;
    total = 0;
    fill = 0;
    for (int i = 0; i < MAXRANK; i++) dims[i] = 0;
  }

  // allocate n elements
  void allocate(int n) {
    MDCHECK(!data);
    data = new T[n];
    total = n;
    fill = 0;
    owned = true;
  }

  // copy another array
  void copy(mdarray<T> &other) {
    clear();
    allocate(other.total);
    for (int i = 0; i < MAXRANK + 1; i++) dims[i] = other.dims[i];
    fill = other.fill;
    for (int i = 0; i < fill; i++) data[i] = other.data[i];
  }

  // get the extent of dimension i
  int dim(int i) {
    MDCHECK(unsigned(i) < MAXRANK);
    MDCHECK(dims[i] > 0);
    return dims[i];
  }

  // get the rank of the array
  int rank() {
    for (int i = 0; i < MAXRANK + 1; i++)
      if (!dims[i]) return i;
    THROW("bad rank");
  }

  // total number of elements in linearized array
  int size() { return fill; }

  // multidimensional subscripting
  template <typename... Args>
  T &operator()(Args... args) {
    int indexes[] = {args...};
    int rank = sizeof indexes / sizeof indexes[0];
    MDCHECK(rank > 0 && rank <= MAXRANK);
    MDCHECK(dims[rank] == 0);
    MDCHECK(unsigned(indexes[0]) < unsigned(dims[0]));
    int index = indexes[0];
    for (int i = 1; i < rank; i++) {
      MDCHECK(unsigned(indexes[i]) < unsigned(dims[i]));
      index = index * dims[i] + indexes[i];
    }
    return data[index];
  }
  // multidimensional subscripting
  template <typename... Args>
  T &unsafe_at(Args... args) {
    int indexes[] = {args...};
    int rank = sizeof indexes / sizeof indexes[0];
    int index = indexes[0];
    for (int i = 1; i < rank; i++) {
      index = index * dims[i] + indexes[i];
    }
    return data[index];
  }
  T &operator[](int i) {
    MDCHECK(unsigned(i) < unsigned(total));
    return data[i];
  }
  // multidimensional subscripting
  template <typename... Args>
  void getSlice(mdarray<T> &result, Args... args) {
    int indexes[] = {args...};
    int rank = sizeof indexes / sizeof indexes[0];
    MDCHECK(rank > 0 && rank <= MAXRANK);
    MDCHECK(dims[rank] != 0);
    MDCHECK(unsigned(indexes[0]) < unsigned(dims[0]));
    int index = indexes[0];
    for (int i = 1; i < rank; i++) {
      MDCHECK(unsigned(indexes[i]) < unsigned(dims[i]));
      index = index * dims[i] + indexes[i];
    }
    result.alias_(data + index, dims + rank);
  }

  // reshape the array without changing its contents
  template <typename... Args>
  mdarray &reshape(Args... args) {
    int indexes[MAXRANK + 1] = {args...};
    int rank = sizeof indexes / sizeof indexes[0];
    indexes[rank] = 0;
    reshape_(indexes);
    return *this;
  }

  // alias an array to given data; the pointer will not
  // get deleted on destruction
  template <typename... Args>
  mdarray &alias(T *p, Args... args) {
    int indexes[MAXRANK + 1] = {args...};
    int rank = sizeof indexes / sizeof indexes[0];
    indexes[rank] = 0;
    alias_(p, indexes);
    return *this;
  }

  // resize the array, destroying the data it contains
  template <typename... Args>
  mdarray &resize(Args... args) {
    int indexes_[] = {args...};
    int rank = sizeof indexes_ / sizeof indexes_[0];
    int indexes[MAXRANK + 1] = {args...};
    indexes[rank] = 0;
    resize_(indexes);
    return *this;
  }

  // fill the array with the given value
  template <class S>
  mdarray &constant(S value) {
    for (int i = 0; i < total; i++) data[i] = value;
    return *this;
  }

  // equivalent shape-related functions using a pointer
  // to an integral, zero-terminated array

  template <typename INT>
  void reshape_(INT *shape, bool exact = true) {
    int nsize = unsigned(prod_(shape));
    MDCHECK(nsize <= total);
    if (exact) MDCHECK(nsize == prod_(dims));
    int i;
    for (i = 0; i < MAXRANK; i++) {
      if (!shape[i]) break;
      dims[i] = int(shape[i]);
    }
    MDCHECK(i <= MAXRANK);
    for (; i <= MAXRANK; i++) dims[i] = 0;
    fill = prod_(dims);
  }
  template <typename INT>
  void resize_(INT *shape) {
    int nsize = prod_(shape);
    if (nsize < total) {
      // nothing
    } else {
      clear();
      allocate(nsize);
    }
    reshape_(shape, false);
  }
  template <typename INT>
  void alias_(T *p, INT *shape) {
    data = p;
    owned = false;
    total = prod_(shape);
    reshape_(shape);
  }
  template <typename INT>
  static inline int prod_(INT *p) {
    if (p[0] == 0) return 0;
    int total = 1;
    for (int i = 0; p[i]; i++) total *= p[i];
    return total;
  }

  // internal consistency check, invariants
  void check_() {
    MDCHECK(unsigned(rank()) <= MAXRANK);
    MDCHECK(prod_(dims) == total);
  }

  // assignment and arithmetic
  template <class O>
  void operator=(mdarray<O> &other) {
    copy(other);
  }
  template <class O>
  void operator=(const mdarray<O> &other) {
    copy(other);
  }

  template <class O, class P>
  void clip(const O &lo, const P &hi) {
    for (int i = 0; i < fill; i++) {
      T value = data[i];
      if (value < lo)
        data[i] = lo;
      else if (value > hi)
        data[i] = hi;
    }
  }
  T min() {
    T result = data[0];
    for (int i = 1; i < fill; i++) {
      T value = data[i];
      if (value < result) result = value;
    }
    return result;
  }
  T max() {
    T result = data[0];
    for (int i = 1; i < fill; i++) {
      T value = data[i];
      if (value > result) result = value;
    }
    return result;
  }
  double normsq() {
    double result = 0.0;
    for (int i = 0; i < fill; i++) {
      T value = data[i];
      result += value * value;
    }
    return result;
  }
  double norm() { return sqrt(normsq()); }
  void randomize() {
    for (int i = 0; i < fill; i++) data[i] = drand48();
  }

  template <class O>
  void operator+=(const O &other) {
    assert(fill == other.fill);
    for (int i = 0; i < fill; i++) data[i] += other.data[i];
  }
  template <class O>
  void operator-=(const O &other) {
    assert(fill == other.fill);
    for (int i = 0; i < fill; i++) data[i] += other.data[i];
  }
  template <class O>
  void operator*=(const O &other) {
    assert(fill == other.fill);
    for (int i = 0; i < fill; i++) data[i] *= other.data[i];
  }
  template <class O>
  void operator/=(const O &other) {
    assert(fill == other.fill);
    for (int i = 0; i < fill; i++) data[i] /= other.data[i];
  }

  template <class O>
  void operator+=(const mdarray<O> &other) {
    assert(fill == other.fill);
    for (int i = 0; i < fill; i++) data[i] += other.data[i];
  }
  template <class O>
  void operator-=(const mdarray<O> &other) {
    assert(fill == other.fill);
    for (int i = 0; i < fill; i++) data[i] += other.data[i];
  }
  template <class O>
  void operator*=(const mdarray<O> &other) {
    assert(fill == other.fill);
    for (int i = 0; i < fill; i++) data[i] *= other.data[i];
  }
  template <class O>
  void operator/=(const mdarray<O> &other) {
    assert(fill == other.fill);
    for (int i = 0; i < fill; i++) data[i] /= other.data[i];
  }
};

template <class T, class S>
inline void to_eigen_matrix(T result, mdarray<S> &a) {
  MDCHECK(a.rank() == 2);
  result.resize(a.dim(0), a.dim(1));
  for (int i = 0; i < a.dim(0); i++)
    for (int j = 0; j < a.dim(1); j++) result(i, j) = a(i, j);
}

template <class T, class S>
inline void to_eigen_vector(T result, mdarray<S> &a) {
  MDCHECK(a.rank() == 1);
  result.resize(a.dim(0));
  for (int i = 0; i < a.dim(0); i++) result(i) = a(i);
}

template <class T, class S>
inline void from_eigen_matrix(mdarray<S> &result, T a) {
  result.resize(ROWS(a), COLS(a));
  for (int i = 0; i < ROWS(a); i++)
    for (int j = 0; j < COLS(a); j++) result(i, j) = a(i, j);
}
}

#endif