#include "Model.hpp"

#include <map>

#include <memory>

#include <unordered_set>

#include "adagrad.hpp"

#include "adam.hpp"

#include "aux.hpp"

#include "gamma_func.hpp"

#include "io.hpp"

#include "pdist.hpp"

#include "rprop.hpp"

#include "sampling.hpp"

using namespace spec_parser;

using namespace std;

using spec_parser::expression::balance;

using spec_parser::expression::deriv;

using spec_parser::expression::eval;

using spec_parser::expression::simplify;

namespace STD {

namespace {

#include "model_aux.cpp"

size_t iter_cnt = 0;

template <typename T>

using ExprPtr = std::shared_ptr<spec_parser::expression::Exp<T>>;

template <typename T>

void compile_expression_and_derivs(const ExprPtr<T> &expr,

const std::string &tag) {

spec_parser::expression::codegen(simplify(balance(simplify(expr))), tag,

collect_variables(expr));

for (auto variable : collect_variables(expr)) {

auto deriv_expr = simplify(balance(simplify(deriv(variable, expr))));

spec_parser::expression::codegen(

deriv_expr, tag + "-" + to_string(*variable), collect_variables(expr));

}

}

bool initialized_jit = false;

Model::Model(const vector<Counts> &c, size_t T_, const Design::Design &design_,

const ModelSpec &model_spec_, const Parameters &parameters_,

bool initialize, bool construct_gp)

: G(max_row_number(c)),

T(T_),

E(0),

S(0),

model_spec(model_spec_),

design(design_),

module_name("std-module"), // TODO use unique module names

rate_fnc(),

odds_fnc(),

rate_derivs(),

odds_derivs(),

experiments(),

parameters(parameters_),

contributions_gene_type(Matrix::Zero(G, T)),

contributions_gene(Vector::Zero(G)) {

if (not(initialized_jit)) {

JIT::init_runtime(module_name, parameters.output_directory + "/");

compile_expression_and_derivs(model_spec.rate_expr, "rate");

compile_expression_and_derivs(model_spec.odds_expr, "odds");

JIT::finalize_module(module_name);

initialized_jit = true;

}

rate_fnc = JIT::get_function("rate");

odds_fnc = JIT::get_function("odds");

for (auto variable : collect_variables(model_spec.rate_expr))

rate_derivs.push_back(JIT::get_function("rate-" + to_string(*variable)));

for (auto variable : collect_variables(model_spec.odds_expr))

odds_derivs.push_back(JIT::get_function("odds-" + to_string(*variable)));

for (auto &counts : c)

add_experiment(counts);

update_contributions();

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

add_covariates();

ensure_dimensions();

coeff_debug_dump("INITIAL");

if (construct_gp)

construct_GPs();

// coeff_debug_dump("BEFORE");

// remove_redundant_terms();

coeff_debug_dump("AFTER");

if (initialize)

for (auto &coeff : coeffs) {

coeff->sample();

if (coeff->type == Coefficient::Type::gp_coord

and not parameters.gp.free_mean)

dynamic_pointer_cast<Coefficient::Spatial::Coord>(coeff)

->subtract_mean();

}

size_t index = 0;

for (auto coeff : coeffs)

LOG(debug) << index++ << " " << *coeff << ": "

<< coeff->info.to_string(design.covariates);

auto fnc = [](const Experiment &a, const Experiment &b) -> bool {

return a.scale_ratio < b.scale_ratio;

};

double min_ratio = std::min_element(begin(experiments), end(experiments), fnc)

->scale_ratio;

for (auto &experiment : experiments)

experiment.scale_ratio /= min_ratio;

if (parameters.gp.center)

center();

verify_model(*this);

}

Model Model::clone() const {

// TODO reactivate coeffs make efficient

// avoid recompilation

// avoid recalculation of spectral decomposition

vector<Counts> counts;

for (auto experiment : experiments)

counts.push_back(experiment.counts);

Model model(counts, T, design, model_spec, parameters, false, false);

for (size_t idx = 0; idx < coeffs.size(); ++idx)

if (coeffs[idx]->type == Coefficient::Type::gp_coord) {

auto coord_coeff

= dynamic_pointer_cast<Coefficient::Spatial::Coord>(coeffs[idx]);

auto model_coord_coeff

= dynamic_pointer_cast<Coefficient::Spatial::Coord>(

model.coeffs[idx]);

model_coord_coeff->gp = coord_coeff->gp;

}

return model;

}

void Model::ensure_dimensions() const {

for (auto &experiment : experiments)

experiment.ensure_dimensions();

}

void Model::setZero() {

for (auto &coeff : coeffs)

coeff->values.setZero();

}

size_t Model::number_variable() const {

size_t s = 0;

for (auto &coeff : coeffs)

s += coeff->number_variable();

return s;

}

size_t Model::size() const {

size_t s = 0;

for (auto &coeff : coeffs)

s += coeff->size();

return s;

}

Vector Model::vectorize() const {

Vector v(size());

auto iter = begin(v);

for (auto &coeff : coeffs)

for (auto &x : coeff->vectorize())

*iter++ = x;

assert(iter == end(v));

return v;

}

void Model::from_vector(const Vector &v) {

auto iter = begin(v);

for (auto &coeff : coeffs)

coeff->from_vector(iter);

}

pair<Matrix, Matrix> Model::compute_mean_and_var(size_t e) const {

LOG(debug) << "Computing means and variances";

const auto &exp = experiments[e];

Matrix Mean = Matrix::Zero(G, exp.S);

Matrix Var = Matrix::Zero(G, exp.S);

const size_t num_rate_coeffs = rate_derivs.size();

const size_t num_odds_coeffs = odds_derivs.size();

#pragma omp parallel if (DO_PARALLEL)

{

Matrix mean = Matrix::Zero(G, exp.S);

Matrix var = Matrix::Zero(G, exp.S);

double rate, odds;

std::vector<std::vector<double>> rate_coeff_arrays, odds_coeff_arrays;

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

rate_coeff_arrays.push_back(std::vector<double>(num_rate_coeffs));

odds_coeff_arrays.push_back(std::vector<double>(num_odds_coeffs));

}

#pragma omp for schedule(guided)

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

for (size_t s = 0; s < exp.S; ++s) {

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

for (size_t i = 0; i < num_rate_coeffs; ++i)

rate_coeff_arrays[t][i] = exp.rate_coeffs[i]->get_actual(g, t, s);

for (size_t i = 0; i < num_odds_coeffs; ++i)

odds_coeff_arrays[t][i] = exp.odds_coeffs[i]->get_actual(g, t, s);

rate = std::exp(rate_fnc(rate_coeff_arrays[t].data()));

odds = std::exp(odds_fnc(odds_coeff_arrays[t].data()));

mean(g, s) += rate * odds;

var(g, s) += rate * odds * (1 + odds);

}

}

#pragma omp critical

{

Mean = Mean + mean;

Var = Var + var;

}

}

return {Mean, Var};

}

Model Model::compute_gradient(double &score) const {

LOG(debug) << "Computing gradient";

score = 0;

Model gradient = clone();

gradient.setZero();

gradient.contributions_gene_type.setZero();

for (auto &experiment : gradient.experiments) {

experiment.contributions_spot_type.setZero();

experiment.contributions_gene_type.setZero();

}

const size_t num_rate_coeffs = rate_derivs.size();

const size_t num_odds_coeffs = odds_derivs.size();

#pragma omp parallel if (DO_PARALLEL)

{

Model grad = gradient.clone();

auto rng = EntropySource::rngs[omp_get_thread_num()];

double score_ = 0;

Vector rate(T), odds(T);

std::vector<std::vector<double>> rate_coeff_arrays, odds_coeff_arrays;

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

rate_coeff_arrays.push_back(std::vector<double>(num_rate_coeffs));

odds_coeff_arrays.push_back(std::vector<double>(num_odds_coeffs));

}

#pragma omp for schedule(guided)

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

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

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

if (RandomDistribution::Uniform(rng)

>= parameters.dropout_gene_spot) {

const auto &exp = experiments[e];

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

for (size_t i = 0; i < num_rate_coeffs; ++i)

rate_coeff_arrays[t][i]

= exp.rate_coeffs[i]->get_actual(g, t, s);

for (size_t i = 0; i < num_odds_coeffs; ++i)

odds_coeff_arrays[t][i]

= exp.odds_coeffs[i]->get_actual(g, t, s);

rate(t) = std::exp(rate_fnc(rate_coeff_arrays[t].data()));

odds(t) = std::exp(odds_fnc(odds_coeff_arrays[t].data()));

}

double total_rate = std::accumulate(begin(rate), end(rate), 0.0);

assert(std::all_of(begin(odds), end(odds), [&odds](const auto &x) {

return x == odds[0];

}));

double total_odds = odds[0];

register_gradient(g, e, s, total_rate, total_odds, grad, rate, odds,

rate_coeff_arrays, odds_coeff_arrays, rng);

double p = odds_to_prob(total_odds);

score_ += log_negative_binomial(exp.counts(g, s), total_rate, p);

}

#pragma omp critical

{

gradient = gradient + grad;

score += score_;

}

}

gradient.update_contributions();

for (size_t i = 0; i < coeffs.size(); ++i)

if (coeffs[i]->type != Coefficient::Type::gp_coord

or iter_cnt >= parameters.gp.first_iteration)

score += coeffs[i]->compute_gradient(gradient.coeffs[i]);

return gradient;

}

void Model::register_gradient(

size_t g, size_t e, size_t s, double total_rate, double total_odds,

Model &gradient, const Vector &rate, const Vector &odds,

const std::vector<std::vector<double>> &rate_coeffs,

const std::vector<std::vector<double>> &odds_coeffs, RNG &rng) const {

const double K = experiments[e].counts(g, s);

double k = K;

if (parameters.adjust_seq_depth)

k = std::binomial_distribution<size_t>(k,

1 / experiments[e].scale_ratio)(rng);

if (parameters.downsample < 1)

k = std::binomial_distribution<size_t>(k, parameters.downsample)(rng);

const double r

= total_rate

* (parameters.adjust_seq_depth ? 1 / experiments[e].scale_ratio : 1);

const double p = odds_to_prob(total_odds);

const double log_one_minus_p = neg_odds_to_log_prob(total_odds);

// const double rate_term = r * (log_one_minus_p + digamma_diff(r, k));

const double rate_term = log_one_minus_p + digamma_diff(r, k);

const double odds_term = k - p * (r + k);

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

gradient.experiments[e].contributions_gene_type(g, t)

+= rate(t) / total_rate * k;

gradient.experiments[e].contributions_spot_type(s, t)

+= rate(t) / total_rate * k;

}

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

{ // loop over rate covariates

auto deriv_iter = begin(rate_derivs);

for (size_t idx = 0; deriv_iter != end(rate_derivs);

++idx, ++deriv_iter) {

gradient.experiments[e].rate_coeffs[idx]->get_raw(g, t, s)

+= rate_term * rate(t) * (*deriv_iter)(rate_coeffs[t].data());

}

}

{ // loop over odds covariates

auto deriv_iter = begin(odds_derivs);

for (size_t idx = 0; deriv_iter != end(odds_derivs);

++idx, ++deriv_iter) {

gradient.experiments[e].odds_coeffs[idx]->get_raw(g, t, s)

+= odds_term * (*deriv_iter)(odds_coeffs[t].data());

}

}

}

}

void Model::center() {

for (auto &coeff : coeffs)

if (coeff->type == Coefficient::Type::gp_points)

for (int t = 0; t < coeff->values.cols(); ++t)

coeff->values.col(t)

= coeff->values.col(t).array() - coeff->values.col(t).mean();

}

void Model::gradient_update(

size_t num_iterations,

std::function<bool(const CoefficientPtr)> is_included) {

LOG(verbose) << "Performing gradient update iterations";

for (auto coeff : coeffs)

if (is_included(coeff))

LOG(debug) << "Optimizing " << coeff->to_string();

size_t current_iteration = 0;

auto eval_and_compute_gradient = [&](const Vector &x, Vector &grad) {

if (((iter_cnt++) % parameters.report_interval) == 0) {

const size_t iteration_num_digits

= 1 + floor(log(num_iterations) / log(10));

store("iter" + to_string_embedded(iter_cnt - 1, iteration_num_digits)

+ "/");

}

// deactivate stochasticity in last iteration for correct contribution stats

auto temp_parameters = parameters;

if (++current_iteration == num_iterations) {

parameters.dropout_gene_spot = 0;

parameters.adjust_seq_depth = false;

parameters.downsample = 1;

}

from_vector(x.array());

if (parameters.gp.center)

center();

double score = 0;

Model model_grad = compute_gradient(score);

for (auto coeff : model_grad.coeffs)

LOG(debug) << coeff << " grad = " << Stats::summary(coeff->values);

// restore parameters from before deactivating stochasticity

parameters = temp_parameters;

// set gradient to zero for fixed coefficients

for (auto coeff : model_grad.coeffs)

if (coeff->type == Coefficient::Type::fixed)

coeff->values.setZero();

// set gradient to zero for coefficients that are not included

for (auto coeff : model_grad.coeffs) {

if (not is_included(coeff))

coeff->values.fill(0);

}

grad = model_grad.vectorize();

contributions_gene_type = model_grad.contributions_gene_type;

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

experiments[e].contributions_spot_type

= model_grad.experiments[e].contributions_spot_type;

experiments[e].contributions_gene_type

= model_grad.experiments[e].contributions_gene_type;

}

LOG(info) << "Iteration " << iter_cnt << ", score: " << score;

LOG(debug) << "x: " << endl << Stats::summary(x);

LOG(debug) << "grad: " << endl << Stats::summary(grad);

return score;

};

Vector x = vectorize().array();

double fx;

switch (parameters.optim_method) {

case Optimize::Method::RPROP: {

Vector grad;

Vector prev_sign(Vector::Zero(x.size()));

Vector rates(x.size());

rates.fill(parameters.grad_alpha);

for (size_t iter = 0; iter < num_iterations; ++iter) {

fx = eval_and_compute_gradient(x, grad);

rprop_update(grad, prev_sign, rates, x, parameters.rprop);

}

} break;

case Optimize::Method::Gradient: {

double alpha = parameters.grad_alpha;

for (size_t iter = 0; iter < num_iterations; ++iter) {

Vector grad;

fx = eval_and_compute_gradient(x, grad);

x = x + alpha * grad;

LOG(debug) << "iter " << iter << " alpha: " << alpha;

LOG(debug) << "iter " << iter << " fx: " << fx;

LOG(debug) << "iter " << iter << " x: " << endl << Stats::summary(x);

alpha *= parameters.grad_anneal;

}

} break;

case Optimize::Method::AdaGrad: {

Vector grad;

Array agrad;

Array ax;

Array scale(Array::Zero(x.size()));

for (size_t iter = 0; iter < num_iterations; ++iter) {

fx = eval_and_compute_gradient(x, grad);

agrad = grad.array();

ax = x.array();

adagrad_update(agrad, scale, ax, parameters.adagrad);

x = ax.matrix();

}

} break;

case Optimize::Method::Adam: {

Vector grad;

Array agrad;

Array ax;

Array mom1(Array::Zero(x.size()));

Array mom2(Array::Zero(x.size()));

auto updater = parameters.adam_nesterov_momentum ? nadam_update<Array>

: adam_update<Array>;

for (size_t iter = 1; iter <= num_iterations; ++iter) {

fx = eval_and_compute_gradient(x, grad);

agrad = grad.array();

ax = x.array();

updater(agrad, mom1, mom2, ax, iter, parameters.adam);

x = ax.matrix();

}

} break;

}

LOG(info) << "Final f(x) = " << fx;

from_vector(x.array());

}

void Model::update_contributions() {

contributions_gene_type.setZero();

contributions_gene.setZero();

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);

}

}

void Model::add_experiment(const Counts &counts) {

experiments.push_back({this, counts, T, parameters});

E++;

S += experiments.back().S;

}

ostream &operator<<(ostream &os, const Model &model) {

size_t n_variable = model.number_variable();

os << "Spatial Transcriptome Deconvolution "

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

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

<< "E = " << model.E << " "

<< "S = " << model.S << endl

<< model.size() << " parameters, " << n_variable << " variable" << endl

<< "G * S = " << (model.G * model.S) << " -> "

<< 100.0 * n_variable / (model.G * model.S) << "%." << endl;

for (auto &experiment : model.experiments)

os << experiment;

return os;

}

Model operator+(const Model &a, const Model &b) {

Model model = a;

model.contributions_gene_type += b.contributions_gene_type;

model.contributions_gene += b.contributions_gene;

for (size_t i = 0; i < a.coeffs.size(); ++i)

model.coeffs[i]->values.array() += b.coeffs[i]->values.array();

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

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

return model;

}

} // namespace