CovariateInformation drop_covariate(const CovariateInformation &info,

const Design::Design &design,

const string &cov_label) {

auto mod_info = info;

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

if (design.covariates[info.idxs[idx]].label == cov_label) {

mod_info.idxs.erase(begin(mod_info.idxs) + idx);

mod_info.vals.erase(begin(mod_info.vals) + idx);

}

return mod_info;

}

CovariateInformation drop_covariates(CovariateInformation info,

const Design::Design &design,

const vector<string> &cov_labels) {

for (auto &cov_label : cov_labels)

info = drop_covariate(info, design, cov_label);

return info;

}

CovariateInformation get_covariate_info(const Design::Design &design,

const set<string> &covariates,

size_t experiment) {

auto covariates_ = covariates;

using namespace Design;

// spot dependency implies section dependency

if (covariates_.find(spot_label) != covariates_.end()) {

covariates_.insert(section_label);

}

vector<size_t> cov_idxs = design.determine_covariate_idxs(covariates_);

vector<size_t> cov_vals;

for (auto &covariate_idx : cov_idxs)

cov_vals.push_back(design.dataset_specifications[experiment]

.covariate_values[covariate_idx]);

return CovariateInformation{cov_idxs, cov_vals};

}

// Removes trailing zeros in a numeric string's decimals.

string remove_trailing_zeros(const string &str) {

if (str.find('.') == string::npos) {

return str;

}

auto pos = str.end();

while (*(--pos) != '.') {

if (*pos != '0') {

pos += 1;

break;

}

}

return string(str.begin(), pos);

}

void verify_model(const Model &m) {

{ // check for overspecification

static const auto input_dim

= Coefficient::Kind::gene | Coefficient::Kind::spot;

for (auto &x : m.coeffs) {

if ((x->kind & input_dim) == input_dim) {

throw runtime_error(

"Error: coefficient '" + x->name

+ "' has dimensionality greater than or equal to the input data.");

}

}

}

/* TODO FIXUP coeffs

{ // check for cycles in model spec

auto check_cycles = [&m](size_t root) {

vector<bool> visited(m.coeffs.size());

function<void(size_t)> go = [&m, &visited, &go](size_t x) {

if (visited[x]) {

throw runtime_error(

"Error: cyclic model specifications are currently not "

"supported.");

}

visited[x] = true;

for (auto &p : m.coeffs[x]->prior_idxs) {

go(p);

}

visited[x] = false;

};

return go(root);

};

unordered_set<size_t> coeffs;

for (auto &e : m.experiments) {

coeffs.insert(e.rate_coeff_idxs.begin(), e.rate_coeff_idxs.end());

coeffs.insert(e.odds_coeff_idxs.begin(), e.odds_coeff_idxs.end());

}

for (auto &x : coeffs) {

check_cycles(x);

}

} */

}

} // namespace

vector<CoefficientPtr>::iterator Model::find_coefficient(

const Coefficient::Id &cid) {

return find_if(begin(coeffs), end(coeffs), [&](const CoefficientPtr &coeff) {

return coeff->name == cid.name and coeff->kind == cid.kind

and coeff->type == cid.type and coeff->info == cid.info;

});

}

CoefficientPtr Model::register_coefficient(

const unordered_map<string, ModelSpec::Variable> &variable_map, string id,

size_t experiment) {

// Register coefficient if it doesn't already exist and return a pointer to it

auto do_registration

= [this, &experiment](

const Coefficient::Id &cid, size_t _G, size_t _T, size_t _S,

const vector<CoefficientPtr> &priors) -> CoefficientPtr {

auto it = find_coefficient(cid);

if (it != end(coeffs)) {

(*it)->experiments.push_back(&experiments[experiment]);

return *it;

} else {

LOG(debug) << "Adding new coefficient for " << cid.name << ".";

coeffs.emplace_back(Coefficient::make_shared(_G, _T, _S, cid, priors));

auto coeff = coeffs.back();

coeff->experiments.push_back(&experiments[experiment]);

LOG(debug) << "Added new coefficient: " << *coeff << ".";

return coeff;

}

};

auto register_fixed = [&](double value) {

Coefficient::Id cid{

.name = id,

.kind = Coefficient::Kind::scalar,

.type = Coefficient::Type::fixed,

.info = CovariateInformation{},

};

auto coeff = do_registration(cid, 0, 0, 0, {});

coeff->get_raw(0, 0, 0) = value;

return coeff;

};

auto register_random = [&]() {

auto it = variable_map.find(id);

if (it == variable_map.end()) {

throw runtime_error("Unknown variable id '" + id + "'.");

}

auto variable = it->second;

auto info = get_covariate_info(design, variable->covariates, experiment);

auto kind = Coefficient::determine_kind(variable->covariates);

if (variable->distribution == nullptr) {

auto dist = parameters.default_distribution;

LOG(debug) << id << " does not have a distribution specification. Using "

<< to_string(dist) << " as per defaults.";

variable->distribution = make_shared<Distribution>(

dist, vector<string>{

remove_trailing_zeros(to_string(

parameters.hyperparameters.get_param(dist, 0))),

remove_trailing_zeros(to_string(

parameters.hyperparameters.get_param(dist, 1))),

});

}

Coefficient::Id cid{

.name = id,

.kind = kind,

.type = variable->distribution->type,

.info = info,

};

vector<CoefficientPtr> priors;

{ // register prior coefficients

size_t i = 1;

for (auto &argument : variable->distribution->arguments) {

auto prior = register_coefficient(variable_map, argument, experiment);

priors.emplace_back(prior);

LOG(debug) << "Prior " << i++ << " of " << id << " is " << argument

<< " (" << *prior << ").";

}

}

CoefficientPtr coeff

= do_registration(cid, G, T, experiments[experiment].S, priors);

if (variable->distribution->type == Coefficient::Type::gp_points) {

auto gp_kind = kind & ~Coefficient::Kind::spot;

// Create or update coordinate system

auto gp_coord_info = drop_covariates(

info, design, {Design::spot_label, Design::section_label});

auto gp_coord_id = Coefficient::Id{

.name

= id + "-gp-coord-"

+ design.get_covariate_value(experiment, Design::coordsys_label),

.kind = gp_kind,

.type = Coefficient::Type::gp_coord,

.info = gp_coord_info,

};

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

do_registration(gp_coord_id, G, T, 0, priors));

LOG(debug) << "Updating GP coordinate system (" << *gp_coord_coeff

<< ").";

gp_coord_coeff->points.emplace_back(

dynamic_pointer_cast<Coefficient::Spatial::Points>(coeff));

}

return coeff;

};

double value;

try {

value = stod(id);

} catch (const invalid_argument &) {

return register_random();

}

return register_fixed(value);

}

void Model::add_covariates() {

auto rate_variables = collect_variables(model_spec.rate_expr);

auto odds_variables = collect_variables(model_spec.odds_expr);

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

LOG(debug) << "Registering coefficients for experiment " << e;

for (auto &variable : rate_variables) {

auto coeff

= register_coefficient(model_spec.variables, variable->full_id(), e);

experiments[e].rate_coeffs.emplace_back(coeff);

}

for (auto &variable : odds_variables) {

auto coeff

= register_coefficient(model_spec.variables, variable->full_id(), e);

experiments[e].odds_coeffs.emplace_back(coeff);

}

}

}

void Model::construct_GPs() {

LOG(debug) << "Constructing GPs";

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

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

LOG(debug) << "Constructing GP " << idx << ": " << *coeffs[idx];

auto coord_coeff

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

coord_coeff->construct_gp();

}

}

void Model::coeff_debug_dump(const string &tag) const {

size_t index = 0;

for (auto coeff : coeffs)

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

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

auto fnc = [&](const string &s, CoefficientPtr coeff, size_t e) {

LOG(debug) << tag << " " << s << " experiment " << e << " " << *coeff

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

};

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

for (auto coeff : experiments[e].rate_coeffs)

fnc("rate", coeff, e);

for (auto coeff : experiments[e].odds_coeffs)

fnc("odds", coeff, e);

}

}

template <typename V>

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

size_t N = v.size();

vector<size_t> order(N);

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

sort(begin(order), end(order),

[&v](size_t a, size_t b) { return v[a] > v[b]; });

return order;

}

void Model::store(const string &prefix_, bool mean_and_var,

bool reorder) const {

string prefix = parameters.output_directory + prefix_;

{

using namespace boost::filesystem;

if (not((exists(prefix) and is_directory(prefix))

or create_directory(prefix)))

throw(std::runtime_error("Couldn't create directory " + prefix));

}

auto type_names = form_factor_names(T);

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

vector<size_t> order;

if (reorder) {

auto cs = colSums<Vector>(contributions_gene_type);

order = get_order(cs);

}

#pragma omp parallel sections if (DO_PARALLEL)

{

#pragma omp section

{

ofstream ofs(prefix + "design.txt");

ofs << design;

}

#pragma omp section

{

ofstream ofs(prefix + "model.txt");

log([&](const string &s) { ofs << s << endl; }, model_spec);

}

#pragma omp section

{

// TODO use parse-able format

ofstream ofs(prefix + "hyperparameters.txt");

ofs << parameters.hyperparameters;

}

// TODO perhaps write out a single file for the scalar covariates

#pragma omp section

{

for (auto &coeff : coeffs) {

vector<string> spot_names;

if (coeff->spot_dependent())

for (Experiment *experiment : coeff->experiments)

spot_names.insert(begin(spot_names),

begin(experiment->counts.col_names),

end(experiment->counts.col_names));

coeff->store(prefix + "covariate-" + storage_type(coeff->kind) + "-"

+ coeff->name + "-"

+ coeff->info.to_string(design.covariates)

+ FILENAME_ENDING,

parameters.compression_mode, gene_names, spot_names,

type_names, order);

}

}

#pragma omp section

write_matrix(contributions_gene_type,

prefix + "contributions_gene_type" + FILENAME_ENDING,

parameters.compression_mode, gene_names, type_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) {

string exp_prefix = prefix + "experiment"

+ to_string_embedded(e, EXPERIMENT_NUM_DIGITS) + "-";

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

}

if (mean_and_var)

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

string exp_prefix = prefix + "experiment"

+ to_string_embedded(e, EXPERIMENT_NUM_DIGITS) + "-";

auto mean_var = compute_mean_and_var(e);

write_matrix(mean_var.first,

exp_prefix + "counts_expected" + FILENAME_ENDING,

parameters.compression_mode, gene_names,

experiments[e].counts.col_names);

write_matrix(mean_var.second,

exp_prefix + "counts_variance" + FILENAME_ENDING,

parameters.compression_mode, gene_names,

experiments[e].counts.col_names);

}

}

/* TODO enable loading of subsets of covariates */

void Model::restore(const string &prefix) {

{

for (auto &coeff : coeffs) {

coeff->restore(

prefix + "covariate-" + storage_type(coeff->kind) + "-" + coeff->name

+ "-" + coeff->info.to_string(design.covariates) + FILENAME_ENDING);

}

}

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

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

string exp_prefix = prefix + "experiment"

+ to_string_embedded(e, EXPERIMENT_NUM_DIGITS) + "-";

experiments[e].restore(exp_prefix);

}

}