#ifndef MODEL_HPP

#define MODEL_HPP

#include <map>

#include "Experiment.hpp"

namespace PoissonFactorization {

const double local_phi_scaling_factor = 50;

const int EXPERIMENT_NUM_DIGITS = 4;

template <typename Type>

struct Model {

using features_t = typename Type::features_t;

using weights_t = typename Type::weights_t;

using experiment_t = Experiment<Type>;

// TODO consider const

/** number of genes */

size_t G;

/** number of factors */

size_t T;

/** number of experiments */

size_t E;

std::vector<experiment_t> experiments;

Parameters parameters;

/** hidden contributions to the count data due to the different factors */

Matrix contributions_gene_type;

Vector contributions_gene;

/** factor loading matrix */

features_t features;

struct CoordinateSystem {

// std::vector<Matrix> coords;

std::vector<size_t> members;

};

std::vector<CoordinateSystem> coordinate_systems;

std::map<std::pair<size_t, size_t>, Matrix> kernels;

template <typename Fnc1, typename Fnc2>

Vector get_low_high(size_t coord_sys_idx, Float init, Fnc1 fnc1,

Fnc2 fnc2) const;

Vector get_low(size_t coord_sys_idx) const;

Vector get_high(size_t coord_sys_idx) const;

void predict_field(std::ofstream &ofs, size_t coord_sys_idx) const;

typename weights_t::prior_type mix_prior;

Model(const std::vector<Counts> &data, size_t T,

const Parameters &parameters, bool same_coord_sys);

void store(const std::string &prefix, bool reorder = true) const;

void restore(const std::string &prefix);

void perform_pairwise_dge(const std::string &prefix) const;

void perform_local_dge(const std::string &prefix) const;

/** sample each of the variables from their conditional posterior */

void gibbs_sample(bool report_likelihood);

void sample_contributions();

void sample_global_theta_priors();

void sample_fields();

double log_likelihood() const;

double log_likelihood_poisson_counts() const;

inline Float &phi(size_t g, size_t t) { return features.matrix(g, t); };

inline Float phi(size_t g, size_t t) const { return features.matrix(g, t); };

// computes a matrix M(g,t)

// with M(g,t) = sum_e local_baseline_phi(e,g) local_phi(e,g,t) sum_s theta(e,s,t) sigma(e,s)

Matrix explained_gene_type() const;

// computes a matrix M(g,t)

// with M(g,t) = phi(g,t) sum_e local_baseline_phi(e,g) local_phi(e,g,t) sum_s theta(e,s,t) sigma(e,s)

Matrix expected_gene_type() const;

void update_contributions();

void update_kernels();

void identity_kernels();

void add_experiment(const Counts &data, size_t coord_sys);

};

template <typename Type>

std::ostream &operator<<(std::ostream &os, const Model<Type> &pfa);

template <typename Type>

Model<Type>::Model(const std::vector<Counts> &c, size_t T_,

const Parameters &parameters_, bool same_coord_sys)

: G(max_row_number(c)),

T(T_),

E(0),

experiments(),

parameters(parameters_),

contributions_gene_type(G, T, arma::fill::zeros),

contributions_gene(G, arma::fill::zeros),

features(G, T, parameters),

mix_prior(sum_rows(c), T, parameters) {

LOG(debug) << "Model G = " << G << " T = " << T << " E = " << E;

size_t coord_sys = 0;

for (auto &counts : c)

add_experiment(counts, same_coord_sys ? 0 : coord_sys++);

update_contributions();

// TODO move this code into the classes for prior and features

if (not parameters.targeted(Target::phi_prior_local))

features.prior.set_unit();

if (not parameters.targeted(Target::phi_local))

features.matrix.ones();

if (parameters.targeted(Target::field)) {

if (parameters.identity_kernels)

identity_kernels();

else

update_kernels();

}

}

template <typename Type>

void Model<Type>::identity_kernels() {

LOG(debug) << "Updating kernels: using identity kernels";

for (auto &coordinate_system : coordinate_systems)

for (auto e1 : coordinate_system.members)

for (auto e2 : coordinate_system.members) {

Matrix m(experiments[e1].S, experiments[e2].S, arma::fill::zeros);

if (e1 == e2)

m.eye();

kernels[{e1, e2}] = m;

}

}

template <typename Type>

void Model<Type>::update_kernels() {

LOG(debug) << "Updating kernels";

for (auto &coordinate_system : coordinate_systems)

for (auto e1 : coordinate_system.members) {

for (auto e2 : coordinate_system.members)

kernels[{e1, e2}]

= apply_kernel(compute_sq_distances(experiments[e1].coords,

experiments[e2].coords),

parameters.hyperparameters.sigma);

}

// row normalize

// TODO check should we do column normalization?

for (auto &coordinate_system : coordinate_systems)

if (true)

for (auto e1 : coordinate_system.members) {

Vector z(experiments[e1].S, arma::fill::zeros);

for (auto e2 : coordinate_system.members)

z += rowSums<Vector>(kernels[{e1, e2}]);

for (auto e2 : coordinate_system.members)

kernels[{e1, e2}].each_col() /= z;

}

else

for (auto e2 : coordinate_system.members) {

Vector z(experiments[e2].S, arma::fill::zeros);

for (auto e1 : coordinate_system.members)

z += colSums<Vector>(kernels[{e1, e2}]);

for (auto e1 : coordinate_system.members)

kernels[{e1, e2}].each_row() /= z.t();

}

for (auto &kernel : kernels)

LOG(debug) << "Kernel " << kernel.first.first << " " << kernel.first.second

<< std::endl

<< kernel.second;

}

template <typename V>

std::vector<size_t> get_order(const V &v) {

size_t N = v.size();

std::vector<size_t> order(N);

std::iota(begin(order), end(order), 0);

std::sort(begin(order), end(order), [&v] (size_t a, size_t b) { return v[a] > v[b]; });

return order;

}

template <typename Type>

void Model<Type>::store(const std::string &prefix, bool reorder) const {

auto factor_names = form_factor_names(T);

auto &gene_names = experiments.begin()->data.row_names;

auto exp_gene_type = expected_gene_type();

std::vector<size_t> order;

if (reorder) {

auto cs = colSums<Vector>(exp_gene_type);

order = get_order(cs);

}

#pragma omp parallel sections if (DO_PARALLEL)

{

#pragma omp section

write_matrix(exp_gene_type, prefix + "expected-features" + FILENAME_ENDING,

parameters.compression_mode, gene_names, factor_names, order);

#pragma omp section

features.store(prefix, gene_names, factor_names, order);

#pragma omp section

write_matrix(contributions_gene_type,

prefix + "contributions_gene_type" + FILENAME_ENDING,

parameters.compression_mode, gene_names, factor_names, order);

#pragma omp section

write_vector(contributions_gene,

prefix + "contributions_gene" + FILENAME_ENDING,

parameters.compression_mode, gene_names);

}

for (size_t e = 0; e < E; ++e) {

std::string exp_prefix = prefix + "experiment"

+ to_string_embedded(e, EXPERIMENT_NUM_DIGITS)

+ "-";

experiments[e].store(exp_prefix, features, order);

}

}

template <typename Type>

void Model<Type>::restore(const std::string &prefix) {

contributions_gene_type = parse_file<Matrix>(prefix + "contributions_gene_type" + FILENAME_ENDING, read_matrix, "\t");

contributions_gene = parse_file<Vector>(prefix + "contributions_gene" + FILENAME_ENDING, read_vector<Vector>, "\t");

features.restore(prefix);

for (size_t e = 0; e < E; ++e) {

std::string exp_prefix = prefix + "experiment"

+ to_string_embedded(e, EXPERIMENT_NUM_DIGITS)

+ "-";

experiments[e].restore(exp_prefix);

}

}

template <typename Type>

void Model<Type>::perform_pairwise_dge(const std::string &prefix) const {

for (size_t e = 0; e < E; ++e) {

std::string exp_prefix

= prefix + "experiment" + to_string_embedded(e, EXPERIMENT_NUM_DIGITS) + "-";

experiments[e].perform_pairwise_dge(exp_prefix, features);

}

}

template <typename Type>

void Model<Type>::perform_local_dge(const std::string &prefix) const {

for (size_t e = 0; e < E; ++e) {

std::string exp_prefix

= prefix + "experiment" + to_string_embedded(e, EXPERIMENT_NUM_DIGITS) + "-";

experiments[e].perform_local_dge(exp_prefix, features);

}

}

template <typename Type>

void Model<Type>::gibbs_sample(bool report_likelihood) {

LOG(verbose) << "perform Gibbs step for " << parameters.targets;

if (parameters.targeted(Target::contributions))

sample_contributions();

if (report_likelihood) {

if (verbosity >= Verbosity::verbose)

LOG(info) << "Log-likelihood = " << log_likelihood();

else

LOG(info) << "Observed Log-likelihood = " << log_likelihood_poisson_counts();

}

// TODO add CLI switch

auto order = random_order(4);

for (auto &o : order)

switch (o) {

case 0:

if (parameters.targeted(Target::theta_prior)

and not parameters.theta_local_priors)

sample_global_theta_priors();

break;

case 1:

if (parameters.targeted(Target::phi_prior))

features.prior.sample(*this);

break;

case 2:

if (parameters.targeted(Target::phi))

features.sample(*this);

break;

case 3:

if (parameters.targeted(Target::field))

sample_fields();

for (auto &experiment : experiments)

experiment.gibbs_sample(features.matrix);

break;

default:

break;

}

}

template <typename Type>

void Model<Type>::sample_contributions() {

for (auto &experiment: experiments)

experiment.sample_contributions(features.matrix);

update_contributions();

}

template <Partial::Kind feat_kind>

void do_sample_fields(

Model<ModelType<feat_kind, Partial::Kind::HierGamma>> &model) {

LOG(verbose) << "Sampling fields";

std::vector<Matrix> observed;

std::vector<Matrix> explained;

for (auto &experiment : model.experiments) {

observed.push_back(Matrix(experiment.S, experiment.T, arma::fill::zeros));

explained.push_back(Matrix(experiment.S, experiment.T, arma::fill::zeros));

}

for (auto &coordinate_system : model.coordinate_systems)

for (auto e2 : coordinate_system.members) {

const auto intensities

= model.experiments[e2].marginalize_genes(model.features.matrix);

#pragma omp parallel for if (DO_PARALLEL)

for (size_t t = 0; t < model.T; ++t)

for (auto e1 : coordinate_system.members) {

const auto &kernel = model.kernels.find({e2, e1})->second;

for (size_t s2 = 0; s2 < model.experiments[e2].S; ++s2) {

for (size_t s1 = 0; s1 < model.experiments[e1].S; ++s1) {

const Float w = kernel(s2, s1);

observed[e1](s1, t)

+= w

* (model.experiments[e2].weights.prior.r(t)

+ model.experiments[e2].contributions_spot_type(s2, t));

explained[e1](s1, t)

+= w * (model.experiments[e2].weights.prior.p(t)

+ intensities[t] * model.experiments[e2].spot[s2]

* (e1 == e2 and s1 == s2

? 1

: model.experiments[e2].field(s2, t)));

}

}

}

}

for (size_t e = 0; e < model.E; ++e) {

LOG(verbose) << "Sampling field for experiment " << e;

Partial::perform_sampling(observed[e], explained[e],

model.experiments[e].field,

model.parameters.over_relax);

}

}

template <Partial::Kind feat_kind>

void do_sample_fields(

Model<ModelType<feat_kind, Partial::Kind::Dirichlet>> &model

__attribute__((unused))) {}

template <typename Type>

void Model<Type>::sample_fields() {

do_sample_fields(*this);

}

template <typename Type>

void Model<Type>::sample_global_theta_priors() {

// TODO refactor

if (std::is_same<typename Type::weights_t::prior_type,

PRIOR::THETA::Dirichlet>::value)

return;

Matrix observed(0, T);

for (auto &experiment : experiments)

observed = arma::join_vert(observed, experiment.contributions_spot_type);

Matrix explained(0, T);

for (auto &experiment : experiments)

explained = arma::join_vert(

explained, experiment.explained_spot_type(features.matrix));

assert(observed.n_rows == explained.n_rows);

if (parameters.normalize_spot_stats) {

observed.each_col() /= rowSums<Vector>(observed);

explained.each_col() /= rowSums<Vector>(explained);

}

mix_prior.sample(observed, explained);

for (auto &experiment : experiments)

experiment.weights.prior = mix_prior;

}

template <typename Type>

double Model<Type>::log_likelihood_poisson_counts() const {

double l = 0;

for (auto &experiment : experiments)

l += experiment.log_likelihood_poisson_counts();

return l;

}

template <typename Type>

double Model<Type>::log_likelihood() const {

double l = features.log_likelihood();

LOG(verbose) << "Global feature log likelihood: " << l;

for (auto &experiment : experiments)

l += experiment.log_likelihood();

return l;

}

// computes a matrix M(g,t)

// with M(g,t) = sum_e local_baseline_phi(e,g) local_phi(e,g,t) sum_s theta(e,s,t) sigma(e,s)

template <typename Type>

Matrix Model<Type>::explained_gene_type() const {

Matrix explained(G, T, arma::fill::zeros);

for (auto &experiment : experiments) {

Vector theta_t = experiment.marginalize_spots();

for (size_t t = 0; t < T; ++t)

#pragma omp parallel for if (DO_PARALLEL)

for (size_t g = 0; g < G; ++g)

explained(g, t)

+= experiment.baseline_phi(g) * experiment.phi(g, t) * theta_t(t);

}

return explained;

}

// computes a matrix M(g,t)

// with M(g,t) = phi(g,t) sum_e local_baseline_phi(e,g) local_phi(e,g,t) sum_s

// theta(e,s,t) sigma(e,s)

template <typename Type>

Matrix Model<Type>::expected_gene_type() const {

Matrix m(G, T, arma::fill::zeros);

for (auto &experiment : experiments)

m += experiment.expected_gene_type(features.matrix);

return m;

}

template <typename V>

bool generate_next(V &v, const V &low, const V &high, double step) {

auto l_iter = low.begin();

auto h_iter = high.begin();

for (auto &x : v) {

if ((x += step) > *h_iter)

x = *l_iter;

else

return true;

l_iter++;

h_iter++;

}

return false;

}

template <typename Type>

template <typename Fnc1, typename Fnc2>

Vector Model<Type>::get_low_high(size_t coord_sys_idx, Float init, Fnc1 fnc1,

Fnc2 fnc2) const {

if (coordinate_systems.size() <= coord_sys_idx

or coordinate_systems[coord_sys_idx].members.empty())

return {0};

size_t D = 0;

for (auto &exp_idx : coordinate_systems[coord_sys_idx].members)

D = std::max<size_t>(D, experiments[exp_idx].coords.n_cols);

Vector v(D, arma::fill::ones);

v *= init;

for (size_t d = 0; d < D; ++d)

for (auto &exp_idx : coordinate_systems[coord_sys_idx].members)

v[d] = fnc1(v[d], fnc2(experiments[exp_idx].coords.col(d)));

return v;

}

template <typename Type>

Vector Model<Type>::get_low(size_t coord_sys_idx) const {

return get_low_high(coord_sys_idx, std::numeric_limits<Float>::infinity(),

[](double a, double b) { return std::min(a, b); },

[](const Vector &x) { return arma::min(x); });

}

template <typename Type>

Vector Model<Type>::get_high(size_t coord_sys_idx) const {

return get_low_high(coord_sys_idx, -std::numeric_limits<Float>::infinity(),

[](double a, double b) { return std::max(a, b); },

[](const Vector &x) { return arma::max(x); });

}

template <typename Type>

void Model<Type>::predict_field(std::ofstream &ofs,

size_t coord_sys_idx) const {

const double sigma = parameters.hyperparameters.sigma;

// Matrix feature_marginal(E, T, arma::fill::zeros); TODO

double N = 8e8;

Vector low = get_low(coord_sys_idx);

Vector high = get_high(coord_sys_idx);

double alpha = 0.05;

low = low - (high - low) * alpha;

high = high + (high - low) * alpha;

// one over N of the total volume

double step = arma::prod((high - low) % (high - low)) / N;

// n-th root of step

step = exp(log(step) / low.n_elem);

Vector coord = low;

do {

bool first = true;

for (auto &c : coord) {

if (first)

first = false;

else

ofs << ",";

ofs << c;

}

Vector observed(T, arma::fill::ones);

Vector explained(T, arma::fill::ones);

for (auto exp_idx : coordinate_systems[coord_sys_idx].members) {

for (size_t s = 0; s < experiments[exp_idx].S; ++s) {

const Vector diff = coord - experiments[exp_idx].coords.row(s).t();

const double d = arma::accu(diff % diff);

const double w

= 1 / sqrt(2 * M_PI) / sigma * exp(-d / 2 / sigma / sigma);

for (size_t t = 0; t < T; ++t) {

observed[t] += w * experiments[exp_idx].contributions_spot_type(s, t);

explained[t] += w * 1 // TODO feature_marginal(exp_idx, t)

// * experiments[exp_idx].theta(s, t)

// / experiments[exp_idx].field(s, t)

* experiments[exp_idx].spot(s);

}

}

}

double z = 0;

for (size_t t = 0; t < T; ++t)

z += observed[t] / explained[t];

for (size_t t = 0; t < T; ++t) {

double predicted = observed[t] / explained[t] / z;

ofs << "," << predicted;

}

ofs << std::endl;

} while (generate_next(coord, low, high, step));

}

template <typename Type>

void Model<Type>::update_contributions() {

contributions_gene_type.zeros();

contributions_gene.zeros();

for (auto &experiment : experiments)

#pragma omp parallel for if (DO_PARALLEL)

for (size_t g = 0; g < G; ++g) {

contributions_gene(g) += experiment.contributions_gene(g);

for (size_t t = 0; t < T; ++t)

contributions_gene_type(g, t)

+= experiment.contributions_gene_type(g, t);

}

}

template <typename Type>

void Model<Type>::add_experiment(const Counts &counts, size_t coord_sys) {

Parameters experiment_parameters = parameters;

experiment_parameters.hyperparameters.phi_p_1 *= local_phi_scaling_factor;

experiment_parameters.hyperparameters.phi_r_1 *= local_phi_scaling_factor;

experiment_parameters.hyperparameters.phi_p_2 *= local_phi_scaling_factor;

experiment_parameters.hyperparameters.phi_r_2 *= local_phi_scaling_factor;

experiments.push_back({counts, T, experiment_parameters});

E++;

// TODO check redundancy with Experiment constructor

experiments.rbegin()->features.matrix.ones();

experiments.rbegin()->features.prior.set_unit(local_phi_scaling_factor);

while(coordinate_systems.size() <= coord_sys)

coordinate_systems.push_back({});

coordinate_systems[coord_sys].members.push_back(E-1);

}

template <typename Type>

std::ostream &operator<<(std::ostream &os, const Model<Type> &model) {

os << "Poisson Factorization "

<< "G = " << model.G << " "

<< "T = " << model.T << " "

<< "E = " << model.E << std::endl;

if (verbosity >= Verbosity::debug) {

print_matrix_head(os, model.features.matrix, "Φ");

os << model.features.prior;

}

for (auto &experiment : model.experiments)

os << experiment;

return os;

}

template <typename Type>

Model<Type> operator*(const Model<Type> &a, const Model<Type> &b) {

Model<Type> model = a;

model.contributions_gene_type %= b.contributions_gene_type;

model.contributions_gene %= b.contributions_gene;

model.features.matrix %= b.features.matrix;

for (size_t e = 0; e < model.E; ++e)

model.experiments[e] = model.experiments[e] * b.experiments[e];

return model;

}

template <typename Type>

Model<Type> operator+(const Model<Type> &a, const Model<Type> &b) {

Model<Type> model = a;

model.contributions_gene_type += b.contributions_gene_type;

model.contributions_gene += b.contributions_gene;

model.features.matrix += b.features.matrix;

for (size_t e = 0; e < model.E; ++e)

model.experiments[e] = model.experiments[e] + b.experiments[e];

return model;

}

template <typename Type>

Model<Type> operator-(const Model<Type> &a, const Model<Type> &b) {

Model<Type> model = a;

model.contributions_gene_type -= b.contributions_gene_type;

model.contributions_gene -= b.contributions_gene;

model.features.matrix -= b.features.matrix;

for (size_t e = 0; e < model.E; ++e)

model.experiments[e] = model.experiments[e] - b.experiments[e];

return model;

}

template <typename Type>

Model<Type> operator*(const Model<Type> &a, double x) {

Model<Type> model = a;

model.contributions_gene_type *= x;

model.contributions_gene *= x;

model.features.matrix *= x;

for (auto &experiment : model.experiments)

experiment = experiment * x;

return model;

}

template <typename Type>

Model<Type> operator/(const Model<Type> &a, double x) {

Model<Type> model = a;

model.contributions_gene_type /= x;

model.contributions_gene /= x;

model.features.matrix /= x;

for (auto &experiment : model.experiments)

experiment = experiment / x;

return model;

}

}

#endif