src/frequencymatrix.h

#ifndef FREQUENCYMATRIX_H
#define FREQUENCYMATRIX_H

#include <cassert>
#include <vector>
#include "main.h"

template <class T>
class FrequencyMatrix {
  std::vector<T> _array;
  std::vector<T> _rowsums;
  size_t _M;
  size_t _N;
  bool _logged;
  bool _fixed;

public:
  FrequencyMatrix(){};
  FrequencyMatrix(size_t m, size_t n, T alpha, bool logged = true);
  T operator()(size_t i, size_t j, bool normalized=true) const;
  T operator()(size_t k, bool normalized=true) const;
  void increment(size_t i, size_t j, T incr_amt);
  void increment(size_t k, T incr_amt);
  T sum(size_t i) const { return _rowsums[i]; }
  size_t argmax(size_t i) const;
  void set_logged(bool logged);
  void fix();
  bool is_fixed() const { return _fixed; }
};

template <class T>
FrequencyMatrix<T>::FrequencyMatrix(size_t m, size_t n, T alpha, bool logged)
    : _array(m*n, logged ? log(alpha):alpha),
      _rowsums(m, logged ? log(n*alpha):n*alpha),
      _M(m),
      _N(n),
      _logged(logged),
      _fixed(false){
}

template <class T>
T FrequencyMatrix<T>::operator()(size_t i, size_t j, bool normalized) const {
  assert(i*_N+j < _M*_N);
  if (_fixed || !normalized) {
      return _array[i*_N+j];
  }
  if (_logged) {
    return _array[i*_N+j]-_rowsums[i];
  } else {
    return _array[i*_N+j]/_rowsums[i];
  }
}

template <class T>
T FrequencyMatrix<T>::operator()(size_t k, bool normalized) const {
  return operator()(0, k, normalized);
}

template <class T>
void FrequencyMatrix<T>::increment(size_t i, size_t j, T incr_amt) {
  if (_fixed) {
    return;
  }

  assert(i < _M && j < _N);
  size_t k = i*_N+j;
  if (_logged) {
    _array[k] = log_add(_array[k], incr_amt);
    _rowsums[i] = log_add(_rowsums[i], incr_amt);
  } else {
    _array[k] += incr_amt;
    _rowsums[i] += incr_amt;
  }
//    assert(!std::isnan(_rowsums[i]) && !std::isinf(_rowsums[i]));
}

template <class T>
void FrequencyMatrix<T>::increment(size_t k, T incr_amt) {
  increment(0, k, incr_amt);
}

template <class T>
void FrequencyMatrix<T>::set_logged(bool logged) {
  if (logged == _logged || _fixed) {
    return;
  }
  if (logged) {
    for (size_t i = 0; i < _M*_N; ++i) {
      _array[i] = log(_array[i]);
    }
    for(size_t i = 0; i < _M; ++i) {
      _rowsums[i] = log(_rowsums[i]);
    }
  } else {
    for (size_t i = 0; i < _M*_N; ++i) {
       _array[i] = sexp(_array[i]);
    }
    for (size_t i = 0; i < _M; ++i) {
      _rowsums[i] = sexp(_rowsums[i]);
    }
  }
  _logged = logged;
}

template <class T>
size_t FrequencyMatrix<T>::argmax(size_t i) const {
  assert(i < _M);
  size_t k = i*_N;
  size_t arg = 0;
  T val = _array[k];
  for (size_t j = 1; j < _N; j++) {
    if (_array[k+j] > val) {
      val = _array[k+j];
      arg = j;
    }
  }
  return arg;
}

template <class T>
void FrequencyMatrix<T>::fix() {
  if (_fixed) {
    return;
  }
  for (size_t i = 0; i < _M; ++i) {
    for (size_t j = 0; j < _N; ++j) {
      _array[i*_N+j] = operator()(i,j);
    }
    _rowsums[i] = (_logged) ? 0 : 1;
  }
  _fixed = true;
}

#endif