/*
Copyright (C) 2017 Sascha Meiers
Distributed under the MIT software license, see the accompanying
file LICENSE.md or http://www.opensource.org/licenses/mit-license.php.
*/
#include <iostream>
#include <fstream>
#include <vector>
#include <unordered_map>
#include <tuple>
#include <boost/program_options/cmdline.hpp>
#include <boost/program_options/options_description.hpp>
#include <boost/program_options/parsers.hpp>
#include <boost/program_options/variables_map.hpp>
#include <boost/tokenizer.hpp>
#include <boost/filesystem.hpp>
#include <boost/progress.hpp>
#include <boost/iostreams/stream.hpp>
#include <boost/iostreams/stream_buffer.hpp>
#include <boost/iostreams/device/file.hpp>
#include <boost/iostreams/filtering_stream.hpp>
#include <boost/iostreams/filter/zlib.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <htslib/sam.h>
#include "version.hpp"
#include "intervals.hpp"
#include "counter.hpp"
#include "distribution.hpp"
#include "hmm.hpp"
#include "iocounts.hpp"
using interval::Interval;
using count::TGenomeCounts;
using count::Counter;
std::vector<std::string> get_states(unsigned ploidy)
{
assert(ploidy>0 && ploidy < 6);
std::vector<std::string> states;
for (unsigned len=1; len < ploidy+1; ++len) {
for (unsigned w = 0; w <= len; ++w) {
unsigned c = len - w;
std::string s;
for (unsigned i = 0; i < w; ++i)
s.push_back('W');
for (unsigned i = 0; i < c; ++i)
s.push_back('C');
states.push_back(std::move(s));
}
}
return states;
}
void setup_HMM_emissions(hmm::HMM<unsigned, hmm::MultiVariate<hmm::NegativeBinomial>> & hmm,
unsigned ploidy,
double nb_p,
double nb_r,
double nb_a,
double prior,
unsigned max_ploidy = 4)
{
assert(max_ploidy>0 && max_ploidy<6);
std::vector<hmm::MultiVariate<hmm::NegativeBinomial>> emissions;
double prior_others = (1-prior)/static_cast<double>(max_ploidy);
//std::cout << "MultiVariate<NegativeBinomial>:" << std::endl;
for (unsigned len=1; len < max_ploidy+1; ++len) {
for (unsigned w = 0; w <= len; ++w) {
// e.g. WWWW (with N=4) should be NB[p,a] x NB[p,(1-a)r]
// or WWWC NB[p,3/4r] x NB[p,1/4r]
// but WW ?? I decide for NB[p,2/4r] x NB[p,a]
unsigned c = len - w;
double fac_w = (w==0 ? nb_a : w/static_cast<double>(max_ploidy) );
double fac_c = (c==0 ? nb_a : c/static_cast<double>(max_ploidy) );
hmm::MultiVariate<hmm::NegativeBinomial> distribution
({ hmm::NegativeBinomial(nb_p, fac_w * nb_r),
hmm::NegativeBinomial(nb_p, fac_c * nb_r)
}, (len == ploidy ? prior : prior_others));
emissions.push_back(std::move(distribution));
// DEBUG
//std::cout << std::setprecision(3) << "\t" << "p = " << nb_p << "\tcell_mean = " << nb_r/nb_p*(1-nb_p) << "\tState = ";
//for (unsigned x = 0; x < w; ++x) std::cout << "W";
//for (unsigned x = 0; x < c; ++x) std::cout << "C";
//std::cout << "\tFactors = [" << fac_w << "\t" << fac_c << "]\t means = [" << fac_w * nb_r/nb_p*(1-nb_p) << "\t" << fac_c * nb_r/nb_p*(1-nb_p) << "]" << std::endl;
}
}
//std::cout << "-" << std::endl;
hmm.set_emissions(emissions);
}
void setup_HMM_emissions(hmm::HMM<unsigned, hmm::CombinedNegBinAndBinomial> & hmm,
unsigned ploidy,
double nb_p,
double nb_r,
double nb_a,
double prior,
unsigned max_ploidy)
{
assert(max_ploidy>0 && max_ploidy<6);
std::vector<hmm::CombinedNegBinAndBinomial> emissions;
double prior_others = (1-prior)/static_cast<double>(max_ploidy);
//std::cout << "CombinedNegBinAndBinomial:" << std::endl;
for (unsigned len=1; len < max_ploidy+1; ++len) {
for (unsigned w = 0; w <= len; ++w) {
unsigned c = len - w;
double ratio = (c == 0 || c == len ? (c == 0 ? 0.02 : 0.98) : static_cast<double>(c)/static_cast<double>(len) );
double r = (len == 0 ? nb_a : nb_r * static_cast<double>(len)/static_cast<double>(max_ploidy));
emissions.push_back(hmm::CombinedNegBinAndBinomial(nb_p, r, ratio, (len == ploidy ? prior : prior_others)));
// DEBUG
//std::cout << std::setprecision(3) << "\t" << "p = " << nb_p << "\tcell_mean = " << nb_r/nb_p*(1-nb_p) << "\tState = ";
//for (unsigned x = 0; x < w; ++x) std::cout << "W";
//for (unsigned x = 0; x < c; ++x) std::cout << "C";
//std::cout << "\t NB mean = " << r/nb_p*(1-nb_p) << "\tBinomial ratio = " << ratio << std::endl;
}
}
//std::cout << "-" << std::endl;
hmm.set_emissions(emissions);
std::cout << "NB params for ploidy " << ploidy << ": p = " << nb_p << "\tr = " << nb_r << "\tmean = " << nb_r/nb_p*(1-nb_p) << std::endl;
}
template <typename TDistribution>
typename hmm::HMM<unsigned, TDistribution> setup_HMM(unsigned max_ploidy, double p_trans)
{
assert(max_ploidy>0 && max_ploidy<6);
std::vector<std::string> states = get_states(max_ploidy);
std::vector<double> initials(states.size(), 1/static_cast<double>(states.size()));
std::vector<double> transitions(states.size() * states.size(), p_trans);
for (unsigned i = 0; i < states.size(); ++i) {
transitions[i * states.size() + i] = 1 - (states.size() - 1) * p_trans;
}
hmm::HMM<unsigned, TDistribution> hmm(states);
hmm.set_initials(initials);
hmm.set_transitions(transitions);
return hmm;
}
template <typename TDistribution>
void run_generic_HMM(std::vector<TGenomeCounts> & counts,
std::vector<unsigned> const & good_cells,
std::vector<CellInfo> & cells,
std::vector<unsigned> const & good_bins,
std::vector<int32_t> const & good_map,
std::unordered_map<std::string, SampleInfo> const & samples,
float p_trans,
double prior,
unsigned ploidy = 2,
unsigned max_ploidy = 4)
{
// Set up and run HMM:
hmm::HMM<unsigned, TDistribution> hmm = setup_HMM<TDistribution>(max_ploidy, p_trans);
//std::cout << "HMM states: ";
//for (auto x: hmm.state_labels) {std::cout << x << "\t";}
//std::cout << std::endl;
for (auto i = good_cells.begin(); i != good_cells.end(); ++i)
{
// set NB(n,p) parameters according to `p` of sample and mean of cell.
float p = samples.at(cells[*i].sample_name).p;
float r = (float)cells[*i].mean_bin_count * p / (1-p);
float a = 0.05;
setup_HMM_emissions(hmm, ploidy, p, r, a, prior, max_ploidy);
run_HMM(hmm, counts[*i], good_bins, good_map);
}
}
struct Conf_ploidy {
std::vector<boost::filesystem::path> f_in;
boost::filesystem::path f_out;
boost::filesystem::path f_bins;
boost::filesystem::path f_excl;
boost::filesystem::path f_info;
boost::filesystem::path f_segments;
unsigned ploidy;
unsigned max_ploidy;
int minMapQual;
unsigned int window;
std::string model;
double prior;
Conf_ploidy() : max_ploidy(4)
{}
};
int main_hmm(int argc, char **argv)
{
// Command line options
Conf_ploidy conf;
boost::program_options::options_description generic("Generic options");
generic.add_options()
("help,?", "show help message")
("verbose,v", "Be more verbose in the output")
("mapq,q", boost::program_options::value<int>(&conf.minMapQual)->default_value(10), "min mapping quality")
("window,w", boost::program_options::value<unsigned int>(&conf.window)->default_value(500000), "window size of fixed windows")
("out,o", boost::program_options::value<boost::filesystem::path>(&conf.f_out)->default_value("out.txt.gz"), "output file for counts + strand state (gz)")
("bins,b", boost::program_options::value<boost::filesystem::path>(&conf.f_bins), "BED file with manual bins (disables -w). See also 'makebins'")
("exclude,x", boost::program_options::value<boost::filesystem::path>(&conf.f_excl), "Exclude chromosomes and regions")
("info,i", boost::program_options::value<boost::filesystem::path>(&conf.f_info), "Write info about samples")
("ploidy,p", boost::program_options::value<unsigned>(&conf.ploidy)->default_value(2), "Assume cells have this ploidy level (max 4)")
("prior,P", boost::program_options::value<double>(&conf.prior)->notifier(in_range(0,1,"prior")), "Prior probability for copy number <p> (penalize other copy number states. Default: All copy numbers are equally probable)")
("model,m", boost::program_options::value<std::string>(&conf.model)->default_value("multiNB"), "Models for HMM: multiNB, Binomial+NB")
;
boost::program_options::options_description hidden("Hidden options");
hidden.add_options()
("input-file", boost::program_options::value<std::vector<boost::filesystem::path> >(&conf.f_in), "input bam file(s)")
;
boost::program_options::positional_options_description pos_args;
pos_args.add("input-file", -1);
boost::program_options::options_description cmdline_options;
cmdline_options.add(generic).add(hidden);
boost::program_options::options_description visible_options;
visible_options.add(generic);
boost::program_options::variables_map vm;
boost::program_options::store(boost::program_options::command_line_parser(argc, argv).options(cmdline_options).positional(pos_args).run(), vm);
boost::program_options::notify(vm);
// Check command line arguments
if (!vm["window"].defaulted() && vm.count("bins")) {
std::cerr << "[Error] -w and -b cannot be specified together" << std::endl << std::endl;
goto print_usage_and_exit;
}
if (vm.count("bins") && vm.count("exclude")) {
std::cerr << "[Error] Exclude chromosomes (-x) have no effect when -b is specified. Stop" << std::endl << std::endl;
goto print_usage_and_exit;
}
if (conf.model != "multiNB" && conf.model != "Binomial+NB") {
std::cerr << "[Error] Unknown --model for the HMM." << std::endl << std::endl;
goto print_usage_and_exit;
}
if (vm.count("prior") && (conf.prior < 0 || conf.prior > 1)) {
std::cerr << "[Error] Prior probability (for CN " << conf.ploidy << ") has to be between 0 and 1. By default all CNs will be equally likely" << std::endl << std::endl;
goto print_usage_and_exit;
}
if (vm.count("help") || !vm.count("input-file"))
{
print_usage_and_exit:
std::cout << std::endl;
std::cout << "Mosaicatcher " << STRINGIFYMACRO(MOSAIC_VERSION_MAJOR);
std::cout << "." << STRINGIFYMACRO(MOSAIC_VERSION_MINOR) << std::endl;
std::cout << "> Count reads from Strand-seq BAM files..." << std::endl;
std::cout << " Now for different ploidy levels, too!" << std::endl;
std::cout << std::endl;
std::cout << "Usage: " << argv[0] << " --ploidy N [OPTIONS] <cell1.bam> <cell2.bam> ..." << std::endl << std::endl;
std::cout << visible_options << std::endl;
std::cout << "Notes:" << std::endl;
std::cout << " * writes a table of bin counts and state classifcation as a gzip file (default: out.txt.gz)" << std::endl;
std::cout << " * Reads are counted by start position" << std::endl;
std::cout << " * One cell per BAM file, including SM tag in header" << std::endl;
std::cout << " * For paired-end data, only read 1 is counted" << std::endl;
return vm.count("help") ? 0 : 1;
}
/////////////////////////////////////////////////////////// global variables
/* leave one BAM header open to get chrom names & lengths */
bam_hdr_t* hdr = NULL;
/* regarding each cell */
std::vector<CellInfo> cells(conf.f_in.size());
std::vector<TGenomeCounts> counts(conf.f_in.size());
std::vector<unsigned> good_cells;
/* regarding each sample */
std::unordered_map<std::string, SampleInfo> samples;
/* regarding bins */
std::vector<Interval> bins;
std::vector<int32_t> chrom_map;
std::vector<unsigned> good_bins;
std::vector<int32_t> good_map;
////////////////////////////////////////////////////////////////////////////
//
// Chapter: Binning & counting
// ===========================
//
// Read sample names from headers.
// Keep one header throughout the program.
if (vm.count("verbose")) std::cout << "[Info] Exploring SAM headers..." << std::endl;
for(unsigned i = 0; i < conf.f_in.size(); ++i)
{
cells[i].id = (int32_t)i;
cells[i].bam_file = conf.f_in[i].string();
samFile* samfile = sam_open(conf.f_in[i].string().c_str(), "r");
if (samfile == NULL) {
std::cerr << "[Error] Fail to open file " << conf.f_in[i].string() << std::endl;
return 1;
}
hdr = sam_hdr_read(samfile);
if (!get_RG_tag("SM", hdr->text, cells[i].sample_name)) {
std::cerr << "[Error] Each BAM file has to have exactly one RG tag. Group cells " << std::endl;
std::cerr << " belonging to the same sample by the SM tag." << std::endl;
std::cerr << " Problematic file: " << conf.f_in[i].string() << std::endl << std::endl;
goto print_usage_and_exit;
}
if (!get_RG_tag("ID", hdr->text, cells[i].cell_name, /*allow_multiple_matches = */ true)) {
std::cerr << "[Error] Each BAM file has to have exactly one RG tag." << std::endl;
std::cerr << " Problematic file: " << conf.f_in[i].string() << std::endl;
goto print_usage_and_exit;
}
sam_close(samfile);
}
// Bin the genome
unsigned median_binsize;
chrom_map = std::vector<int32_t>(hdr->n_targets, -1);
if (vm.count("bins"))
{
if (!read_dynamic_bins(bins,
chrom_map,
conf.f_bins.string().c_str(),
hdr))
return 1;
TMedianAccumulator<unsigned> med_acc;
for (Interval const & b : bins)
med_acc(b.end - b.start);
median_binsize = boost::accumulators::median(med_acc);
if (vm.count("verbose")) std::cout << "[Info] Reading " << bins.size() << " variable-width bins with median bin size of " << round(median_binsize/1000) << "kb" << std::endl;
}
else
{
std::vector<Interval> exclude;
if (vm.count("exclude")) {
read_exclude_file(conf.f_excl.string(), hdr, exclude, vm.count("verbose"));
sort(exclude.begin(), exclude.end(), interval::invt_less);
}
if (vm.count("verbose")) std::cout << "[Info] Creating " << round(conf.window/1000) << "kb bins with " << exclude.size() << " excluded regions" << std::endl;
create_fixed_bins(bins,
chrom_map,
conf.window,
exclude,
hdr->n_targets,
hdr->target_len);
median_binsize = conf.window;
}
// add last element for easy calculation of number of bins
chrom_map.push_back((int32_t)bins.size());
// Count in bins. If A bam file cannot be read, the cell is ignored and
// the respective entry in `counts` and `cells` will be erased.
if (vm.count("verbose")) std::cout << "[Info] Reading " << conf.f_in.size() << " BAM files...";
boost::progress_display show_progress1(conf.f_in.size());
for (unsigned i = 0, i_f = 0; i_f < conf.f_in.size(); ++i, ++i_f)
{
if (!count_sorted_reads(conf.f_in[i_f].string(),
bins,
chrom_map,
hdr,
conf.minMapQual,
counts[i],
cells[i]))
{
std::cerr << "[Warning] Ignoring cell " << conf.f_in[i_f].string() << std::endl;
counts.erase(counts.begin()+i);
cells.erase(cells.begin()+i);
--i;
}
++show_progress1;
}
//
// Chapter: Filter cells and bins and estimate NB parameter p
// ==========================================================
//
// median per cell
count::set_median_per_cell(counts, cells);
// filter cells with low counts
good_cells = count::get_good_cells(counts, cells);
// filter bins with abnormal counts
if (good_cells.size() < 5) {
std::cerr << "[Warning] Only few cells with sufficient coverage. I will not filter bad bins" << std::endl;
good_bins.resize(bins.size());
std::iota(good_bins.begin(), good_bins.end(), 0); // fill with 0,1,2,...
} else {
good_bins = count::get_good_bins(counts, cells, good_cells);
if (vm.count("verbose")) std::cout << "[Info] Filtered out " << bins.size() - good_bins.size() << " bad bins" << std::endl;
}
// build chrom_map for good bins
good_map = std::vector<int32_t>(chrom_map.size() - 1, -1);
int32_t pos = 0;
for (int32_t chr = 0; chr < static_cast<int32_t>(good_map.size()); ++chr) {
while (pos < good_bins.size() && bins[good_bins[pos]].chr < chr)
++pos;
// now goodit is either at first occurence of chr, or at the end.
if (pos >= good_bins.size()) good_map[chr] = (int32_t)good_bins.size();
else good_map[chr] = pos;
}
// add last element for easy calculation of number of bins
good_map.push_back((int32_t)good_bins.size());
// calculate cell means and cell variances, grouped by sample (not cell)
calculate_new_cell_mean(samples, cells, counts, good_cells, good_bins);
// Estimation of parameter p per sample (should work even with one cell only)
for (auto it = samples.begin(); it != samples.end(); ++it) {
SampleInfo & s = it->second;
s.p = std::inner_product(s.means.begin(), s.means.end(), s.means.begin(), 0.0f) \
/ std::inner_product(s.means.begin(), s.means.end(), s.vars.begin(), 0.0f);
}
//
// Chapter: Run HMM
// ================
//
double prior = (vm.count("prior") ? conf.prior : 1.0/static_cast<double>(conf.max_ploidy));
std::cout << "[Info] Running HMM with model " << conf.model << " and expected ploidy ";
std::cout << std::setprecision(3) << conf.ploidy << " (prior = " << prior << ")" << std::endl;
if (conf.model == "multiNB")
run_generic_HMM<hmm::MultiVariate<hmm::NegativeBinomial>>(
counts,
good_cells,
cells,
good_bins,
good_map,
samples,
10.0f / bins.size(),
prior,
conf.ploidy,
conf.max_ploidy);
if (conf.model == "Binomial+NB")
run_generic_HMM<hmm::CombinedNegBinAndBinomial>(
counts,
good_cells,
cells,
good_bins,
good_map,
samples,
10.0f / bins.size(),
prior,
conf.ploidy,
conf.max_ploidy);
// Print cell information:
if (vm.count("info")) {
if (vm.count("verbose")) std::cout << "[Write] Cell summary: " << conf.f_info.string() << std::endl;
write_cell_info(conf.f_info.string(), cells);
}
// Write final counts + classification
std::cout << "[Write] count table: " << conf.f_out.string() << std::endl;
{
// TODO: why do I pass vector<pair>? I could make it two separate vectors. Just check where else the io function is called.
struct sample_cell_name_wrapper {
std::vector<CellInfo> const & cells;
sample_cell_name_wrapper(std::vector<CellInfo> const & cells) : cells(cells)
{}
std::pair<std::string,std::string> operator[](size_t i) const {
return std::make_pair(cells[i].sample_name, cells[i].cell_name);
}
};
if (!io::write_counts_gzip(conf.f_out.string(),
counts,
bins,
hdr->target_name,
sample_cell_name_wrapper(cells)) )
return 1;
}
return 0;
}