pedphase.cpp

#include "pedphase.hh"

using namespace std;

static void usage()
{
    fprintf(stderr, "\n");
    fprintf(stderr, "About:   akt pedphase - simple Mendel inheritance phasing of duos/trios\n");
    fprintf(stderr, "Usage:   ./akt pedphase input.vcf.gz -p pedigree.fam\n");
    fprintf(stderr, "\n");
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "    -p, --pedigree                 pedigree information in plink .fam format\n");
    fprintf(stderr, "    -o, --output-file <file>       output file name [stdout]\n");
    fprintf(stderr, "    -O, --output-type <b|u|z|v>    b: compressed BCF, u: uncompressed BCF, z: compressed VCF, v: uncompressed VCF [v]\n");
    fprintf(stderr, "    -@, --threads                  number of compression/decompression threads to use\n");
    fprintf(stderr, "    -x, --exclude-chromosome       leave these chromosomes unphased (unphased lines will still be in in output)  eg. -x chrM,chrY\n");
    exit(1);
}

static int bcf_int32_count_missing(int32_t *array, int n)
{
    int num_missing = 0;
    for (int i = 0; i < n; i++)
    {
        if (array[i] == bcf_int32_missing)
        {
            num_missing++;
        }
    }
    return (num_missing);
}

static void bcf_int32_set_missing(int32_t *array, int n)
{
    for (int i = 0; i < n; i++)
    {
        array[i] = bcf_int32_missing;
    }
}

PedPhaser::PedPhaser(args &a)
{
    setup_io(a);
    cerr << "Reading input from " << a.inputfile << endl;    
    _num_sample = bcf_hdr_nsamples(_out_header);
    _parental_genotypes.assign(2*_num_sample,pair<int,int>(bcf_gt_missing,bcf_gt_missing));
    _num_gt=_num_sample*2;    
    _gt_array = (int *)malloc(sizeof(int)*_num_gt);
    _ps_array = (int32_t *)malloc(sizeof(int32_t)*_num_sample);
    _rps_array = (int32_t *)malloc(sizeof(int32_t)*_num_sample);
    _mendel_conflict = (const char **)malloc(sizeof(char *)*_num_sample);
    
    _num_rps=_num_ps=_num_sample;
    if (!a.exclude_chromosomes.empty())
    {
        cerr << "Will not phase chromosomes: " << a.exclude_chromosomes << " (--exclude-chromosomes)" << endl;
        vector<string> chromosomes;
        stringSplit(a.exclude_chromosomes, ',', chromosomes);
        if (chromosomes.empty())
        {
            die("Problem parsing --exclude-chromosomes: " + a.exclude_chromosomes);
        }
        for (vector<string>::iterator chromosome = chromosomes.begin(); chromosome != chromosomes.end(); chromosome++)
        {
            int rid = bcf_hdr_id2int(_in_header, BCF_DT_CTG, chromosome->c_str());
            if (rid == -1)
            {
                cerr << "WARNING: chromosome " << *chromosome << " was not in header" << endl;
            }
            else
            {
                _chromosomes_to_ignore.push_back(rid);
            }
        }
    }
    main();
}

void PedPhaser::main()
{
    int min_distance_to_flush=10000;
    int prev_rid = -1;
    bcf1_t *line;
    vector<int32_t> last_phase_set(_num_sample, bcf_int32_missing); //stores the current phase set for each sample
    while (bcf_sr_next_line(_bcf_reader))
    {
        line = bcf_sr_get_line(_bcf_reader, 0);
	if(line->rid!=prev_rid) flush_buffer();
	prev_rid = line->rid;
	
        if (chromosome_is_in_ignore_list(line))
        {
            flush_buffer(); //just in case something was sitting in buffer from previous chromosome
            bcf_write(_out_file, _out_header, line);
        }
        else
        {
            bcf_unpack(line, BCF_UN_ALL);
	    int ps_status=bcf_get_format_int32(_in_header, line, "PS", &_ps_array, &_num_ps);
	    if(ps_status == -1)
	    {
                bcf1_t *new_line = bcf_dup(line);
                _line_buffer.push_back(new_line);
		flush_buffer();
	    }
	    else if(ps_status == -2)
	    {
		die("FORMAT/PS had incorrect type");
	    }
	    else
	    {
		if(ps_status == -3)
		{
		    int distance_from_last_phase_set = INT_MAX;
		    for(size_t i=0;i<_num_sample;i++)
		    {
			int d = last_phase_set[i]==bcf_int32_missing ? INT_MAX : (line->pos-last_phase_set[i]);
			distance_from_last_phase_set = min(d,distance_from_last_phase_set);
		    }
		    if(distance_from_last_phase_set>min_distance_to_flush)
		    {
			flush_buffer();
			std::fill(last_phase_set.begin(),last_phase_set.end(),bcf_int32_missing);			
		    }
		}
		else
		{
		    for(size_t i=0;i<_num_sample;i++)
			if(_ps_array[i]!=bcf_int32_missing)
			    last_phase_set[i] = _ps_array[i];
		}
		
                bcf1_t *new_line = bcf_dup(line);
                _line_buffer.push_back(new_line);
	    }
	}
    }
    flush_buffer();
}

//performs simple duo/trio phasing using mendelian inheritance.
//returns
//-2: not a child in a duo/trio
//-1: mendelian inconsistent
//0:  unphaseable
//1:  phased
int phase_by_transmission(Genotype & kid_gt,Genotype & dad_gt,Genotype & mum_gt)
{
    const int NUM_LEAVES= 8;//maximum number of leaves on the binary tree(we are only doing duos/trios so it never gets this big)    
    int pedigree_size = 3; //we might make this dynamic later
    if ((dad_gt.isMissing() && mum_gt.isMissing()) || kid_gt.isMissing())  return (-2); //unphaseable due to missingness

    //phasetree is perfect binary tree (stored as array) that enumerates every genotype configuration in the pedigree
    //the leaf of each tree is 0 if the genotype configuration is inconsistent with inheritance and 1 otherwise.
    //redundant leaves (eg. where a sample is homozygous) are also 0
    vector<bool> phasetree(pow(2, pedigree_size), 0);
    assert(phasetree.size() <= NUM_LEAVES);
    //this loop enumerates the 2**n possible phase configurations and checks which are compatible with inheritance
    for (size_t i = 0; i < phasetree.size(); i++)
    {
        bitset<NUM_LEAVES> leaf((int)i);
        bool kid_branch = leaf[0];
        bool dad_branch = leaf[1];
        bool mum_branch = leaf[2];
        if ((!kid_gt.is_phased() || !kid_branch) && (!dad_gt.is_phased() || !dad_branch) && (!mum_gt.is_phased() || !mum_branch))
        {
            bool is_inheritance_consistent = dad_gt.isMissing() || dad_gt.getGenotype(dad_branch) == kid_gt.getGenotype((kid_branch + 1) % 2);
            is_inheritance_consistent &= mum_gt.isMissing() || mum_gt.getGenotype(mum_branch) == kid_gt.getGenotype(kid_branch);
            if (is_inheritance_consistent)
            {
                phasetree[i] = 1;
            }
        }
    }
    int sum = accumulate(phasetree.begin(), phasetree.end(), 0);

    if (sum > 1)
    {
        return (0);  //multiple solutions - cannot phase
    }
    else if (sum == 1) //found a unique solution for phasing. update the genotype array.
    {
        int leaf = find(phasetree.begin(), phasetree.end(), 1) - phasetree.begin();
        if ((leaf >> 0) % 2)   kid_gt.swap();
        if ((leaf >> 1) % 2)   dad_gt.swap();
        if (leaf >> 2)         mum_gt.swap();

	kid_gt.setPhase(true);
	dad_gt.setPhase(true);
	mum_gt.setPhase(true);
        return (1);
    }
    else //inconsistent with mendelian inheritance.
    {
        return (-1);
    }
}

bool is_mendel_inconsistent(Genotype  kid,Genotype  dad,Genotype  mum)
{
    if(kid.isMissing() || !kid.is_phased()) return false;
    int k0=kid.first();
    int k1=kid.second();
    int m_transmitted = k0;
    int d_transmitted = k1;
    if(!dad.isMissing())
	d_transmitted=dad.first();
    if(!mum.isMissing())
	m_transmitted=mum.first();
    return(k0!=m_transmitted || k1!=d_transmitted);
}

bool is_mendel_inconsistent(pair<int,int> kid,pair<int,int> dad,pair<int,int> mum)
{
    if(bcf_gt_is_missing(kid.first) || bcf_gt_is_missing(kid.second))
	return false;
    int k0=kid.first;
    int k1=kid.second;
    int m_transmitted = k0;
    int d_transmitted = k1;
    if(!bcf_gt_is_missing(dad.first)&&!bcf_gt_is_missing(dad.second))
	d_transmitted=dad.first;
    if(!bcf_gt_is_missing(mum.first)&&!bcf_gt_is_missing(mum.second))
	m_transmitted=mum.first;    
    return(k0!=m_transmitted || k1!=d_transmitted);
}

int PedPhaser::mendel_phase(int kid_index, int *gt_array, int *ps_array)
{
    int dad_index = _pedigree->getDadIndex(kid_index);
    int mum_index = _pedigree->getMumIndex(kid_index);
    Genotype kid_gt(kid_index, gt_array, ps_array);
    Genotype dad_gt(dad_index, gt_array, ps_array);
    Genotype mum_gt(mum_index, gt_array, ps_array);        
    if(mum_index>=0 || dad_index>=0) _sample_has_been_phased[kid_index]=true;
    bool update_dad=false,update_mum=false;
    if(mum_index>=0)
    {
	update_mum = !_sample_has_been_phased[mum_index];        
	_sample_has_been_phased[mum_index]=true;
    }
    if(dad_index>=0)
    {
	update_dad = !_sample_has_been_phased[dad_index];	
	_sample_has_been_phased[dad_index]=true;
    }
    int ret = phase_by_transmission(kid_gt,dad_gt,mum_gt);
    if(ret)
    {
	if(dad_index>=0) _parental_genotypes[kid_index*2]   = pair<int,int>(dad_gt.first(),dad_gt.second());
	if(mum_index>=0) _parental_genotypes[kid_index*2+1] = pair<int,int>(mum_gt.first(),mum_gt.second());	
	if(update_dad) dad_gt.update_bcf_gt_array(gt_array, dad_index);
	if(update_mum) mum_gt.update_bcf_gt_array(gt_array, mum_index);
	kid_gt.update_bcf_gt_array(gt_array, kid_index);
    }    
    return(ret);
}

bool PedPhaser::chromosome_is_in_ignore_list(bcf1_t *record)
{
    for (vector<int>::iterator rid = _chromosomes_to_ignore.begin(); rid != _chromosomes_to_ignore.end(); rid++)
        if (*rid == record->rid)
            return true;

    return false;
}

//Turns haploid GTs (1) into padded "diploid" ones like (1,bcf_vector32_end).
int *diplofy(int *gt,int num_sample)
{
    std::vector<int> src(gt,gt+num_sample);
    gt=(int *)realloc(gt,2*num_sample*sizeof(int));//this should be redundant
    for(int i=0;i<num_sample;i++)
    {
	gt[i*2] = src[i];
	gt[i*2+1] = bcf_int32_vector_end;
    }
    return(gt);
}

int PedPhaser::flush_buffer()
{
//    std::cerr<<"flushing buffer "<<_line_buffer.size()<<std::endl;//debug
    if (_line_buffer.empty()) return(0); 
    int _num_gt=0,_num_ps=0;
    HaplotypeBuffer hap_transmission(_num_sample,_pedigree);//stores the transmission phased haplotypes
    HaplotypeBuffer hap_phaseset(_num_sample,_pedigree);//stores the phase-set phased haplotypes
    for (deque<bcf1_t *>::iterator it1 = _line_buffer.begin(); it1 != _line_buffer.end(); it1++)
    {
	bcf1_t *line = *it1;	
	int status = bcf_get_genotypes(_out_header, line, &_gt_array, &_num_gt);
	assert(status== 2 * _num_sample || status==_num_sample);
	if(status==_num_sample)//This is a  hack to handle all-haploid VCF rows.
	    _gt_array=diplofy(_gt_array,_num_sample);
	
	hap_transmission.push_back(_gt_array);
	status = bcf_get_format_int32(_in_header, line,"PS", &_ps_array, &_num_ps);
	if(status==_num_sample)
	    hap_phaseset.push_back(_gt_array,_ps_array);
	else if(status<0)
	    hap_phaseset.push_back(_gt_array,nullptr);
	else
	    die(("Invalid PS length: "+std::to_string(status)).c_str());
    }
    hap_transmission.phase();
    hap_transmission.align(hap_phaseset);
    
    int count=0;
    //Finally, flush the buffer to the output file.
    while (!_line_buffer.empty()) 
    {
        bcf1_t *line = _line_buffer.front();
        _line_buffer.pop_front();
	hap_transmission.update_bcf1_genotypes(count,_gt_array,_ps_array,_rps_array);
	bcf_update_genotypes(_out_header, line, _gt_array, 2*_num_sample);
	if(bcf_int32_count_missing(_ps_array,_num_sample)==_num_sample)
	    bcf_update_format_int32(_out_header, line, "PS", nullptr,0);
	else
	    bcf_update_format_int32(_out_header, line, "PS", _ps_array, _num_sample);
	if(bcf_int32_count_missing(_rps_array,_num_sample)==_num_sample)	
	    bcf_update_format_int32(_out_header, line, "RPS", nullptr,0);
	else
	    bcf_update_format_int32(_out_header, line, "RPS", _rps_array, _num_sample);

	if(bcf_update_format_int32(_out_header, line, "ME", hap_transmission.get_mendel_conflict(count), _num_sample)!=0)
	    die("problem writing FORMAT/ME");
	
	bcf_write(_out_file, _out_header, line);	
        bcf_destroy(line);
	count++;
    }
    return (0);
}


void PedPhaser::setup_io(args &a)
{
    //open a file.
    _bcf_reader = bcf_sr_init();

    if (a.targets != nullptr)
    {
        if (bcf_sr_set_targets(_bcf_reader, a.targets, a.targets_is_file, 0) < 0)
        {
            cerr << "ERROR: Failed to set targets " << a.targets << endl;
            exit(1);
        }
    }

    if (a.regions != nullptr)
    {
        if (bcf_sr_set_regions(_bcf_reader, a.regions, a.regions_is_file) < 0)
        {
            cerr << "ERROR: Failed to read the regions: " << a.regions << endl;
            exit(1);
        }
    }

    if (bcf_sr_add_reader(_bcf_reader, a.inputfile) != 1)
    {
        cerr << "ERROR: problem opening " << a.inputfile << endl;
        exit(1);
    }

    _in_header = _bcf_reader->readers[0].header;

    if (a.pedigree == nullptr)
    {
        _pedigree = new sampleInfo(_in_header);
    }
    else
    {
        _pedigree = new sampleInfo(a.pedigree, _in_header);
    }
    if (_pedigree->N <= 0)
    {
        die("no pedigree detected");
    }
    setup_output(a);
}

void PedPhaser::setup_output(args &a)
{
    _out_header = bcf_hdr_dup(_in_header);
    char output_type[] = "wv";
    output_type[1] = a.output_type;
    _out_file = hts_open(a.outfile, output_type);
    if (a.nthreads > 0)
    {
        bcf_sr_set_threads(_bcf_reader, a.nthreads);
        hts_set_threads(_out_file, a.nthreads);
    }

    bcf_hdr_remove(_out_header, BCF_HL_FMT, "PS"); //remove the old PS descripion
    bcf_hdr_append(_out_header, "##FORMAT=<ID=PS,Number=1,Type=Integer,Description=\"Read-backed phase set. If missing from a phased genotype then it indicates the genotype was pedigree-phased such that children are phased as 'maternal allele | paternal allele' and parents are phased as 'allele transmitted to first child | untransmitted allele'\">");
    bcf_hdr_append(_out_header, "##FORMAT=<ID=RPS,Number=1,Type=Integer,Description=\"Read-backed phase set. The phase set (PS) value before this phased genotype was incorporated into the pedigree phase set\">");
    bcf_hdr_append(_out_header, "##FORMAT=<ID=ME,Number=1,Type=Integer,Description=\"Mendel error. A value of ME=1 indicates that this sample is a child in a duo/trio with genotypes that are inconsistent with Mendelian inheritance. The value is 0 if the child is in a trio that is Mendel consistent and has no missing genotypes. The value is missing otherwise.\">");
    bcf_hdr_append(_out_header, ("##akt_pedphase_version=" + (string)AKT_VERSION).c_str());
    bcf_hdr_write(_out_file, _out_header);
}

PedPhaser::~PedPhaser()
{
    delete _pedigree;
    hts_close(_out_file);
    free(_ps_array);
    free(_gt_array);
    bcf_sr_destroy(_bcf_reader);
    bcf_hdr_destroy(_out_header);
    if(_rps_array) free(_rps_array);
    free(_mendel_conflict);
}

int pedphase_main(int argc, char **argv)
{
    int c;
    args arguments;
    arguments.output_type = 'v';
    if (argc < 3)
        usage();
    static struct option loptions[] = {
        {"out", 1, 0, 'o'},
        {"output-type", 1, 0, 'O'},
        {"pedigree", 1, 0, 'p'},
        {"threads", required_argument, nullptr, '@'},
        {"targets", required_argument, nullptr, 't'},
        {"targets-file", required_argument, nullptr, 'T'},
        {"regions-file", required_argument, nullptr, 'R'},
        {"regions", required_argument, nullptr, 'r'},
        {"exclude-chromosome", required_argument, nullptr, 'x'},
        {0, 0, 0, 0}};
    arguments.regions_is_file = false;
    arguments.targets_is_file = false;
    arguments.targets = arguments.pedigree = arguments.inputfile = arguments.include = arguments.regions = nullptr;
    arguments.outfile = "-";
    arguments.nthreads = 0;
    arguments.exclude_chromosomes = "";

    while ((c = getopt_long(argc, argv, "o:p:t:T:r:R:O:@:x:", loptions, nullptr)) >= 0)
    {
        switch (c)
        {
        case 'o':
            arguments.outfile = optarg;
            break;
        case 'O':
            arguments.output_type = optarg[0];
            break;
        case 'p':
            arguments.pedigree = optarg;
            break;
        case 'i':
            arguments.include = optarg;
            break;
        case 't':
            arguments.targets = optarg;
            break;
        case 'T':
            arguments.targets = optarg;
            break;
        case 'r':
            arguments.regions = optarg;
            break;
        case 'x':
            arguments.exclude_chromosomes = optarg;
            break;
        case '@':
            arguments.nthreads = atoi(optarg);
            break;
        case 'R':
            arguments.regions = optarg;
            arguments.regions_is_file = true;
            break;
        default:
            die("unknown argument");
        }
    }
    optind++;
    arguments.inputfile = argv[optind];
    if (arguments.inputfile == nullptr)
    {
        die("no input provided");
    }

    cerr << "Output file: " << arguments.outfile << endl;
    PedPhaser p(arguments);
    return (0);
}