jobs.cc

/* Spell-QTL  Software suite for the QTL analysis of modern datasets.
 * Copyright (C) 2016,2017  Damien Leroux <damien.leroux@inra.fr>, Sylvain Jasson <sylvain.jasson@inra.fr>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "dispatch.h"
#include "output.h"
#include "pedigree.h"
/*#include "factor_var4.h"*/
#include "graphnode2.h"
#include "tensor.h"

bool job_check_fourcc(ifile& ifs, const char* fourcc)
{
    if (check_fourcc(ifs, fourcc)) {
        MSG_ERROR("NOT A JOB DATA FILE OR HAS BEEN CORRUPTED", "Re-run this tool.");
        return true;
    }
    return false;
}


inline
std::ostream&
operator << (std::ostream& os, const std::pair<char, char>& p) { return os << ((int) p.first) << ((int) p.second); }


inline
std::ostream&
operator << (std::ostream& os, const std::map<bn_label_type, double>& obs)
{
    os << '{';
    auto i = obs.begin(), j = obs.end();
    if (i != j) {
        os << i->first << '=' << i->second;
        for (++i; i != j; ++i) {
            os << ' ' << i->first << '=' << i->second;
        }
    }
    return os << '}';
}


void
dispatch_geno_probs(
        const std::vector<gencomb_type>& lc,  /* table parent states -> spawnling states */
        const std::vector<label_type>& labels, /* spawnling labels */
        const std::map<label_type, double>& geno_probs, /* spawnling genotype probabilities */
        size_t ind, /* spawnling number */
        std::map<size_t, std::vector<double>>& state_prob) /* all computed locus vectors to fetch the parents' states */
{
/*
    scoped_indent _(MESSAGE("[dispatchGP] "));
    MSG_DEBUG("ind " << ind);
    MSG_DEBUG("lc " << lc);
    MSG_DEBUG("labels " << labels);
    MSG_DEBUG("geno_probs " << geno_probs);
    MSG_DEBUG("state_prob " << state_prob);
*/
//    MSG_QUEUE_FLUSH();
    std::vector<double> ret(lc.size());
    std::map<label_type, double> norms;
    for (size_t i = 0; i < ret.size(); ++i) {
        ret[i] = 0;
        for (const auto& e: lc[i]) {
            double coef = e.coef;
            for (const auto& k: e.keys) {
                coef *= (state_prob[k.parent][k.state] > 0 ? 1. : 0.);
            }
            ret[i] += coef;
        }
        norms[labels[i]] += ret[i];
    }
    for (auto& kv: norms) if (kv.second) {
        auto it = geno_probs.find(kv.first);
        if (it == geno_probs.end()) {
            kv.second = 0;
        } else {
            kv.second = it->second / kv.second;
        }
    }
    for (size_t i = 0; i < ret.size(); ++i) { ret[i] *= norms[labels[i]]; }
    state_prob[ind].swap(ret);
}



std::vector<bn_label_type>
get_domain(std::map<var_vec, genotype_comb_type>& domains, variable_index_type var)
{
    std::vector<bn_label_type> ret;
    const auto& dom = domains[{var}];
    ret.reserve(dom.size());
    for (const auto& e: dom) {
        ret.push_back(e.keys.keys.front().state);
    }
    return ret;
}



genotype_comb_type
build_evidence(variable_index_type id, const std::map<bn_label_type, double>& obs)
{
    genotype_comb_type tmp;
    for (const auto& kv: obs) {
        tmp.m_combination.emplace_back(genotype_comb_type::element_type{{{id, kv.first}}, kv.second});
    }
    return tmp;
}

std::map<std::string, std::pair<count_jobs_type, job_type>>
job_registry = {
    {"dummy-one", {
        [](bn_settings_t*) { return 1; },
        [](bn_settings_t*, size_t i) { MSG_ERROR("running #" << i, ""); return true; }
    }},
    {"dummy", {
        [](bn_settings_t* settings) { return settings->count_markers(); },
        [](bn_settings_t*, size_t i) { MSG_ERROR("running #" << i, ""); return true; }
    }},
    {"compute-alleles-per-marker", {
        [](bn_settings_t* settings) { return settings->count_markers(); },
        [](bn_settings_t* settings, size_t mark)
        {
            std::string filename = settings->job_filename("compute-alleles-per-marker", mark);
            
            if (check_file(filename, false, true, false)) {
                return true;
            }

            /*const std::string& marker_name = settings->marker_names[mark];*/
            auto get_obs
                = [&] (const std::string& gen, size_t i)
                {
                    const auto& pop_obs = settings->observed_mark.find(gen)->second;
                    const marker_data& obs = pop_obs.observations;
                    auto ret = settings->marker_observation_specs
                        .find(pop_obs.format_name)->second
                        .score(obs.data.find(settings->marker_names[mark])->second[i]);
                    /*MSG_DEBUG("FETCHING OBS[" << marker_name << ':' << gen << ':' << i << "] => " << ret);*/
                    return ret;
                };

            ofile dump(filename);

            std::set<char> used_alleles;
            for (const auto& kv: settings->observed_mark) {
                if (settings->marker_observation_specs.find(kv.second.format_name)->second.domain != ODAllele) {
                    continue;
                }
                const std::vector<int>&
                    gen_ind = settings->pedigree.individuals_by_generation_name.find(kv.first)->second;
                for (size_t i = 0; i < gen_ind.size(); ++i) {
                    for (const auto& score: get_obs(kv.first, i)) {
                        used_alleles.insert(score.first);
                        used_alleles.insert(score.second);
                    }
                }
                /*MSG_DEBUG("Used alleles in generation " << kv.first << ": " << used_alleles.size());*/
            }
            if (used_alleles.size() == 0) {
                used_alleles.insert(1);
            }
            rw_base()(dump, used_alleles);
            return true;
        }
    }},
    {"compute-factor-graphs", {
        [](bn_settings_t* settings) { return settings->unique_n_alleles.size(); },
        [](bn_settings_t* settings, size_t n)
        {
            std::vector<size_t> unique_n_alleles;
            {
                ifile unafs(settings->job_filename("unique_n_alleles", 0));
                rw_base()(unafs, unique_n_alleles);
            }
            std::string filename = settings->job_filename("factor-graph", unique_n_alleles[n]);
            if (check_file(filename, false, true, false)) {
                return true;
            }
            /*MSG_DEBUG("Pedigree" << std::endl << settings->pedigree);*/
            /*MSG_DEBUG("COMPUTING FACTOR GRAPH FOR " << unique_n_alleles[n] << " ALLELES.");*/
            /*factor_graph fg(settings->pedigree, unique_n_alleles[n], settings->noise);*/
            auto fg = factor_graph_type::from_pedigree(settings->pedigree.items, unique_n_alleles[n], settings->output_mode & bn_settings_t::OutputGametes, settings->input_generations, /*settings->output_generations*/ {});
            /*MSG_DEBUG("COMPUTED FACTOR GRAPH" << std::endl << fg);*/
            /*fg.finalize();*/
            /*fg.save(settings->job_filename("factor-graph", unique_n_alleles[n]));*/
            /*fg->dump();*/
            //fg->to_image(MESSAGE("factor-graph-" << unique_n_alleles[n]), "png");
            auto I = instance(fg);
            ofile of(filename);
            I->file_io(of);
            rw_comb<int, bn_label_type>()(of, fg->domains);
            return true;
        }
    }},
    {"compute-LV", {
        [](bn_settings_t* settings) { return settings->count_markers(); },
        [](bn_settings_t* settings, size_t mark)
        {
            std::string filename = settings->job_filename("compute-LV", mark);
            if (check_file(filename, false, true, false)) {
                return true;
            }

            const std::string& marker_name = settings->marker_names[mark];
            std::set<char> used_alleles;
            std::map<char, char> allele_obs_to_idx;
            {
                ifile apm(settings->job_filename("compute-alleles-per-marker", mark));
                rw_base()(apm, used_alleles);
                char idx = 0;
                for (char c: used_alleles) {
                    allele_obs_to_idx[c] = idx++;
                    /*MSG_DEBUG("ALLELE CONVERT " << c << " => " << ((int) allele_obs_to_idx[c]));*/
                }
            }
            /*graph_type fg;*/
            /*fg.load(settings->job_filename("factor-graph", used_alleles.size()));*/
            ifile ifs(settings->job_filename("factor-graph", used_alleles.size()));
            std::unique_ptr<instance_type> I(new instance_type());
            std::map<var_vec, genotype_comb_type> domains;
            I->file_io(ifs);
            rw_comb<int, bn_label_type>()(ifs, domains);

            std::map<char, int> allele_obs2bn;
            for (char c: used_alleles) { allele_obs2bn[c] = allele_obs2bn.size(); }

            auto make_obs
                = [&] (const population_marker_observation& pop_obs, size_t i, size_t n)
                {
                    std::map<bn_label_type, double> ind_obs;
                    const auto& obs_spec = settings->marker_observation_specs.find(pop_obs.format_name)->second;
                    const auto& obsdat = pop_obs.observations.data.find(settings->marker_names[mark])->second[i];
                    std::vector<std::pair<char, char>>
                        obs_vec = obs_spec.score(obsdat);
                    genotype_comb_type::key_type key;
                    key.parent = n;
                    for (const auto& label: get_domain(domains, n)) {
                        ind_obs[label] = 0;
                    }
//                     MSG_DEBUG("i " << i);
//                     MSG_DEBUG("obs " << obsdat);
//                     MSG_DEBUG("obs_vec " << obs_vec);
//                     MSG_DEBUG("ind_dom = " << get_domain(domains, n));
//                     MSG_DEBUG("ind_obs = " << ind_obs);
                    /*MSG_DEBUG("obs spec");*/
                    /*for (const auto& s: obs_spec.scores) {*/
                        /*MSG_DEBUG("* " << s.first);*/
                        /*for (const auto& p: s.second) {*/
                            /*MSG_DEBUG("  " << p.first << p.second);*/
                        /*}*/
                    /*}*/
//                     MSG_DEBUG("obs for #" << n << " '" << pop_obs.observations.data.find(settings->marker_names[mark])->second[i] << "': " << obs_vec);
                    if (obs_spec.domain == ODAllele) {
                        for (const auto& obs: obs_vec) {
                            for (const auto& label: get_domain(domains, n)) {
                                if (label.first_allele == allele_obs_to_idx[obs.first] && label.second_allele == allele_obs_to_idx[obs.second]) {
                                    ind_obs[label] = 1;
                                }
                            }
                        }
                    } else {
                        for (const auto& obs: obs_vec) {
                            for (const auto& label: get_domain(domains, n)) {
                                if (label.first == obs.first && label.second == obs.second) {
                                    ind_obs[label] = 1;
                                }
                            }
                        }
                    }
//                     MSG_DEBUG("ind_obs = " << ind_obs);
                    return ind_obs;
                };

            /*auto instance = fg.instance();*/
            /*instance.clear_evidence();*/
            /*std::vector<std::shared_ptr<message_type>> evidence;*/

            for (const auto& pop_obs: settings->observed_mark) {
                /*MSG_DEBUG("Obs in generation " << pop_obs.first);*/
                size_t n = 0;
                for (size_t i: settings->pedigree.individuals_by_generation_name.find(pop_obs.first)->second) {
                    int id = settings->pedigree.ind2id(i);
                    auto obs = make_obs(pop_obs.second, n, id);
                    /*evidence.emplace_back(std::make_shared<message_type>(message_type{build_evidence(id, obs)}));*/
                    I->add_evidence(id, build_evidence(id, obs));
                    MSG_DEBUG("observations for individual #" << n << " (id #" << id << ") = " << obs);
                    /*MSG_DEBUG("id=" << id << " id=" << id << " evidence " << evidence);*/
                    double accum = 0;
                    /*for (const auto& kv: obs) {*/
                        /*evidence.force_set({{node, kv.first}}, kv.second);*/
                        /*instance.evidence(node, kv.first, kv.second);*/
                        /*accum += kv.second;*/
                    /*}*/
                    if (0 && /* FIXME */ accum == 0) {
                        msg_handler_t::cout() << std::endl;
                        MSG_ERROR("Inconsistent observations for individual #" << n << " of generation " << pop_obs.first << " on marker " << marker_name, MESSAGE("Remove inconsistencies in your genotype data file for generation " << pop_obs.first));
                        return false;
                    }
                    ++n;
                }
            }

            /*instance.evidence(evidence);*/
            /*MSG_DEBUG("EVIDENCE");*/
            /*MSG_DEBUG("" << instance.evidence().dump());*/

            /*std::vector<graph_type::query_operation_type> marginals;*/
            std::map<edge_type, std::shared_ptr<message_type>> messages;
            /*for (const auto& i: settings->pedigree.items) {*/
                /*auto q = fg.build_query_operation({i.id});*/
                /*marginals.insert(marginals.end(), q.begin(), q.end());*/
            /*}*/
            /*std::map<int, genotype_comb_type> marginals;*/
            /*fg.compute_messages(evidence.begin(), evidence.end(), messages);*/
            /*fg.compute_full_factor_state(messages, marginals);*/
            message_type output = I->compute(0, domains);
            MSG_DEBUG("OUTPUT:");
            std::map<variable_index_type, std::map<label_type, double>> marginals;
            for (const auto& t: output) {
                MSG_DEBUG(" * " << t << std::endl);
                for (const auto& e: t) {
                    for (const auto& k: e.keys) {
                        label_type l = {k.state.first, k.state.second};
                        marginals[k.parent][l] += e.coef;
                    }
                }
            }

            if (settings->output_mode & bn_settings_t::OutputOnePoint) {
                std::string op_dirname = MESSAGE(settings->work_directory << '/' << settings->pop_name << ".1-point");
                ensure_directories_exist(op_dirname);
                std::string op_filename = MESSAGE(op_dirname << '/' << settings->pop_name << ".pedigree-and-probabilities." << marker_name << ".csv");
                std::ofstream ofs(op_filename);
                ofs << "Gen;Id;P1;P2;Prob" << std::endl;
                for (const auto& pi: settings->pedigree.items) {
                    ofs << pi.gen_name << ';' << pi.id << ';' << pi.p1 << ';' << pi.p2 << ";{";
                    bool first = true;
                    for (const auto& kv: marginals[pi.id]) {
                        ofs << (first ? "\"" : ", \"") << kv.first << "\": " << kv.second;
                        first = false;
                    }
                    ofs << '}' << std::endl;
                }
                settings->closing_messages.push_back(MESSAGE(GREEN << "1-point Parental Origin Probabilities for marker `" << marker_name << "' written in file " << WHITE << op_filename << NORMAL << std::endl));
            }

            if (settings->output_mode & bn_settings_t::OutputPopData) {
                MSG_DEBUG("PROJECTED MARGINALS");
                for (const auto &km: marginals) {
                    std::stringstream ss;
                    ss << " * " << km.first;
                    for (const auto &kv: km.second) {
                        ss << ' ' << kv.first << '=' << kv.second;
                    }
                    MSG_DEBUG(ss.str());
                }
                //            MSG_QUEUE_FLUSH();

                std::map<size_t, std::vector<double>> state_prob, output_prob;
                std::map<int, VectorXd> gamete_prob;

                for (size_t ind = 1; ind < settings->pedigree.tree.m_ind_number_to_node_number.size(); ++ind) {
                    size_t node = settings->pedigree.tree.m_ind_number_to_node_number[ind];
                    auto gen = settings->pedigree.get_gen(ind);
                    const auto &labels = gen->labels;
                    const auto &LC = settings->pedigree.LC[node];
                    int variable = settings->pedigree.ind2id(ind);
                    dispatch_geno_probs(LC, labels, marginals[variable], node, state_prob);
                    if (std::find(settings->output_generations.begin(), settings->output_generations.end(),
                                  gen->name) != settings->output_generations.end()) {
                        output_prob[node] = state_prob[node];
                        if ((settings->output_mode & bn_settings_t::OutputGametes) && !settings->pedigree.items[ind - 1].is_ancestor()) {
                            auto& mat = marginals[-variable];
                            MSG_DEBUG("marginals for gamete " << mat);
                            std::map<label_type, double> prob;
                            for (const auto& kv: mat) {
                                prob[{kv.first.first(), kv.first.second()}] = kv.second;
                            }
                            if (settings->pedigree.items[ind - 1].is_dh()) {
                                gamete_prob[ind].resize(2);
                                gamete_prob[ind] << prob[{GAMETE_L, 0}], prob[{GAMETE_R, 0}];
                            } else {
                                gamete_prob[ind].resize(4);
                                gamete_prob[ind] << prob[{GAMETE_L, GAMETE_L}], prob[{GAMETE_L, GAMETE_R}], prob[{GAMETE_R, GAMETE_L}], prob[{GAMETE_R, GAMETE_R}];
                            }
                            MSG_DEBUG("GAMETE #" << ind << " prob " << gamete_prob[ind - 1].transpose());
                        }
                    }
                }

                MSG_DEBUG("LOCUS VECTORS");

                for (const auto &kv: state_prob) {
                    MSG_DEBUG(" * " << kv.first << ' ' << kv.second);
                }

                ofile ofs(filename);
                rw_base()(ofs, output_prob);
                
                if (settings->output_mode & bn_settings_t::OutputGametes) {
                    ofile og(settings->job_filename("gametes-LV", mark));
                    rw_base()(og, gamete_prob);
                }
            }
            
            return true;

        }
    }},
    {"collect-LV", {
        [](bn_settings_t*) { return 1; },
        [](bn_settings_t* settings, size_t)
        {
            settings->load_full_pedigree_data();
            const pedigree_type& ped = settings->pedigree;
            std::map<std::string, std::map<size_t, std::vector<double>>> mark_state_prob;
            size_t n_mark = settings->count_markers();
            bool have_gametes = settings->output_mode & bn_settings_t::OutputGametes;
            gamete_LV_database gamete_LV;
            for (size_t m = 0; m < n_mark; ++m) {
                ifile ifs(settings->job_filename("compute-LV", m));
                rw_base() (ifs, mark_state_prob[settings->marker_names[m]]);
                if (have_gametes) {
                    ifile ig(settings->job_filename("gametes-LV", m));
                    std::map<int, MatrixXd> gp;
                    rw_base() (ig, gp);
                    for (const auto& kv: gp) {
                        if (kv.second.size() == 0) {
                            MSG_DEBUG("HAVE EMPTY GAMETE LV " << kv.first);
                        } else {
                            if (kv.second.sum() == 0) {
                                MSG_DEBUG("HAVE ZERO GAMETE LV " << kv.first);
                            } else {
                                gamete_LV.add_gametes(settings->marker_names[m], kv.first, kv.second, ped.items[kv.first - 1].is_dh());
                            }
                        }
                    }
                }
            }

            if (settings->output_mode & bn_settings_t::OutputGametes) {
                std::string filename = MESSAGE(settings->work_directory << "/" << settings->pop_name << ".cache/gamete.data");
                ofile og(filename);
                rw_base() (og, gamete_LV);
                MSG_DEBUG("Gametes: " << gamete_LV.data);
//                 abort();
            }

            pop_data_type pop_data;
            pop_data.ancestor_names = settings->pedigree.ancestor_names;

            for (size_t m = 0; m < n_mark; ++m) {
                std::map<size_t, std::vector<double>> mark_states;
                ifile ifs(settings->job_filename("compute-LV", m));
                rw_base() (ifs, mark_states);
                const std::string& mark = settings->marker_names[m];
                for (auto& ind_sp: mark_states) {
                    const std::string& gen_name = ped.generations[ped.node_generations[ind_sp.first]]->name;
                    Eigen::Map<Eigen::VectorXd> sp(ind_sp.second.data(), ind_sp.second.size());
                    pop_data.LV.data_by_marker[gen_name][mark].emplace_back(sp);
                }
            }


            pop_data.name = settings->pop_name;
            for (const auto& om: settings->observed_mark) {
                pop_data.marker_observation_filenames[om.first] = om.second.filename;
            }

//             settings->load_full_pedigree_data();

            /*pop_data.genetic_map_filename = settings->map_filename;*/
            pop_data.pedigree_filename = settings->pedigree_filename;
            /*pop_data.qtl_generation_name = settings->qtl_generation_name;*/
            pop_data.generations = ped.generations;
            pop_data.families.clear();
            for (size_t ind_node: ped.tree.m_ind_number_to_node_number) {
                if (ind_node == (size_t) -1) {
                    continue;
                }
                auto gen = pop_data.generations[ped.node_generations[ind_node]];
                pop_data.families[gen->name].push_back(ind_node);
                pop_data.gen_by_id[ind_node] = gen;
            }
            MSG_DEBUG("POP-DATA:");
            MSG_DEBUG("" << pop_data);
            std::string filename = pop_data.save(settings->work_directory + "/" + settings->pop_name + ".cache");
            MSG_DEBUG("Saved output in " << filename);
            {
                ofile output(settings->job_filename("output", 0));
                rw_base() (output, filename);
            }
            return true;
        }
    }},
};

#include "output_impl.h"