BamWriter.h

   1 #ifndef BAMWRITER_H_
   2 #define BAMWRITER_H_
   3
   4 #include<cmath>
   5 #include<cstdio>
   6 #include<cstring>
   7 #include<cassert>
   8 #include<string>
   9 #include<map>
  10 #include<sstream>
  11
  12 #include "sam/bam.h"
  13 #include "sam/sam.h"
  14
  15 #include "utils.h"
  16 #include "SingleHit.h"
  17 #include "PairedEndHit.h"
  18
  19 #include "HitWrapper.h"
  20 #include "Transcript.h"
  21 #include "Transcripts.h"
  22
  23 class BamWriter {
  24 public:
  25         BamWriter(char, const char*, const char*, const char*, const char*);
  26         ~BamWriter();
  27
  28         void work(HitWrapper<SingleHit>, Transcripts&);
  29         void work(HitWrapper<PairedEndHit>, Transcripts&);
  30 private:
  31         samfile_t *in, *out;
  32
  33         std::map<std::string, int> refmap;
  34         std::map<std::string, int>::iterator iter;
  35
  36         struct SingleEndT {
  37                 bam1_t *b;
  38
  39                 SingleEndT(bam1_t *b = NULL) {
  40                         this->b = b;
  41                 }
  42
  43                 bool operator< (const SingleEndT& o) const {
  44                         int strand1, strand2;
  45                         uint32_t *p1, *p2;
  46
  47                         if (b->core.tid != o.b->core.tid) return b->core.tid < o.b->core.tid;
  48                         if (b->core.pos != o.b->core.pos) return b->core.pos < o.b->core.pos;
  49                         strand1 = b->core.flag & 0x0010; strand2 = o.b->core.flag & 0x0010;
  50                         if (strand1 != strand2) return strand1 < strand2;
  51                         if (b->core.n_cigar != o.b->core.n_cigar) return b->core.n_cigar < o.b->core.n_cigar;
  52                         p1 = bam1_cigar(b); p2 = bam1_cigar(o.b);
  53                         for (int i = 0; i < (int)b->core.n_cigar; i++) {
  54                                 if (*p1 != *p2) return *p1 < *p2;
  55                                 ++p1; ++p2;
  56                         }
  57                         return false;
  58                 }
  59         };
  60
  61         //b is mate 1, b2 is mate 2
  62         struct PairedEndT {
  63                 bam1_t *b, *b2;
  64
  65                 PairedEndT() { b = NULL; b2 = NULL;}
  66
  67                 PairedEndT(bam1_t *b, bam1_t *b2) {
  68                         this->b = b;
  69                         this->b2 = b2;
  70                 }
  71
  72                 bool operator< (const PairedEndT& o) const {
  73                         int strand1, strand2;
  74                         uint32_t *p1, *p2;
  75
  76                         //compare b
  77                         if (b->core.tid != o.b->core.tid) return b->core.tid < o.b->core.tid;
  78                         if (b->core.pos != o.b->core.pos) return b->core.pos < o.b->core.pos;
  79                         strand1 = b->core.flag & 0x0010; strand2 = o.b->core.flag & 0x0010;
  80                         if (strand1 != strand2) return strand1 < strand2;
  81                         if (b->core.n_cigar != o.b->core.n_cigar) return b->core.n_cigar < o.b->core.n_cigar;
  82                         p1 = bam1_cigar(b); p2 = bam1_cigar(o.b);
  83                         for (int i = 0; i < (int)b->core.n_cigar; i++) {
  84                                 if (*p1 != *p2) return *p1 < *p2;
  85                                 ++p1; ++p2;
  86                         }
  87
  88                         //compare b2
  89                         if (b2->core.tid != o.b2->core.tid) return b2->core.tid < o.b2->core.tid;
  90                         if (b2->core.pos != o.b2->core.pos) return b2->core.pos < o.b2->core.pos;
  91                         strand1 = b2->core.flag & 0x0010; strand2 = o.b2->core.flag & 0x0010;
  92                         if (strand1 != strand2) return strand1 < strand2;
  93                         if (b2->core.n_cigar != o.b2->core.n_cigar) return b2->core.n_cigar < o.b2->core.n_cigar;
  94                         p1 = bam1_cigar(b2); p2 = bam1_cigar(o.b2);
  95                         for (int i = 0; i < (int)b2->core.n_cigar; i++) {
  96                                 if (*p1 != *p2) return *p1 < *p2;
  97                                 ++p1; ++p2;
  98                         }
  99
 100                         return false;
 101                 }
 102         };
 103
 104         uint8_t getMAPQ(double val) {
 105                 double err = 1.0 - val;
 106                 if (err <= 1e-10) return 100;
 107                 return (uint8_t)(-10 * log10(err) + .5); // round it
 108         }
 109
 110         void push_qname(const uint8_t* qname, int l_qname, std::vector<uint8_t>& data) {
 111                 for (int i = 0; i < l_qname; i++) data.push_back(*(qname + i));
 112         }
 113
 114         void push_seq(const uint8_t* seq, int readlen, char strand, std::vector<uint8_t>& data) {
 115                 int seq_len = (readlen + 1) / 2;
 116
 117                 switch (strand) {
 118                 case '+': for (int i = 0; i < seq_len; i++) data.push_back(*(seq + i)); break;
 119                 case '-':
 120                         uint8_t code, base;
 121                         code = 0; base = 0;
 122                         for (int i = 0; i < readlen; i++) {
 123                                 switch (bam1_seqi(seq, readlen - i - 1)) {
 124                                 case 1: base = 8; break;
 125                                 case 2: base = 4; break;
 126                                 case 4: base = 2; break;
 127                                 case 8: base = 1; break;
 128                                 case 15: base = 15; break;
 129                                 default: assert(false);
 130                                 }
 131                                 code |=  base << (4 * (1 - i % 2));
 132                                 if (i % 2 == 1) { data.push_back(code); code = 0; }
 133                         }
 134
 135                         if (readlen % 2 == 1) { data.push_back(code); }
 136                         break;
 137                 default: assert(false);
 138                 }
 139         }
 140
 141         void push_qual(const uint8_t* qual, int readlen, char strand, std::vector<uint8_t>& data) {
 142                 switch (strand) {
 143                 case '+': for (int i = 0; i < readlen; i++) data.push_back(*(qual + i)); break;
 144                 case '-': for (int i = readlen - 1; i >= 0; i--) data.push_back(*(qual + i)); break;
 145                 default: assert(false);
 146                 }
 147         }
 148
 149         //convert transcript coordinate to chromosome coordinate and generate CIGAR string
 150         void tr2chr(const Transcript&, int, int, int&, int&, std::vector<uint8_t>&);
 151 };
 152
 153 //fn_list can be NULL
 154 BamWriter::BamWriter(char inpType, const char* inpF, const char* fn_list, const char* outF, const char* chr_list) {
 155         switch(inpType) {
 156         case 's': in = samopen(inpF, "r", fn_list); break;
 157         case 'b': in = samopen(inpF, "rb", fn_list); break;
 158         default: assert(false);
 159         }
 160         assert(in != 0);
 161
 162         //generate output's header
 163         bam_header_t *out_header = NULL;
 164         refmap.clear();
 165
 166         if (chr_list == NULL) {
 167                 out_header = in->header;
 168         }
 169         else {
 170                 out_header = sam_header_read2(chr_list);
 171
 172                 for (int i = 0; i < out_header->n_targets; i++) {
 173                         refmap[out_header->target_name[i]] = i;
 174                 }
 175         }
 176
 177
 178         out = samopen(outF, "wb", out_header);
 179         assert(out != 0);
 180
 181         if (chr_list != NULL) { bam_header_destroy(out_header); }
 182 }
 183
 184 BamWriter::~BamWriter() {
 185         samclose(in);
 186         samclose(out);
 187 }
 188
 189 void BamWriter::work(HitWrapper<SingleHit> wrapper, Transcripts& transcripts) {
 190         bam1_t *b;
 191         std::string cqname; // cqname : current query name
 192         std::map<SingleEndT, double> hmap;
 193         std::map<SingleEndT, double>::iterator hmapIter;
 194         SingleHit *hit;
 195
 196         int cnt = 0;
 197
 198         cqname = "";
 199         b = bam_init1();
 200         hmap.clear();
 201
 202         while (samread(in, b) >= 0) {
 203
 204                 if (verbose && cnt > 0 && cnt % 1000000 == 0) { printf("%d entries are finished!\n", cnt); }
 205                 ++cnt;
 206
 207                 if (b->core.flag & 0x0004) continue;
 208
 209                 hit = wrapper.getNextHit();
 210                 assert(hit != NULL);
 211
 212                 int sid = b->core.tid + 1;
 213                 assert(sid == hit->getSid());
 214                 const Transcript& transcript = transcripts.getTranscriptAt(sid);
 215
 216                 if (transcripts.getType() == 0) {
 217                         int pos = b->core.pos;
 218                         int readlen = b->core.l_qseq;
 219                         uint8_t *qname = b->data, *seq = bam1_seq(b), *qual = bam1_qual(b);
 220                         std::vector<uint8_t> data;
 221                         data.clear();
 222
 223                         iter = refmap.find(transcript.getSeqName());
 224                         assert(iter != refmap.end());
 225                         b->core.tid = iter->second;
 226                         b->core.qual = 255;
 227
 228                         uint16_t rstrand = b->core.flag & 0x0010; // read strand
 229                         b->core.flag -= rstrand;
 230                         rstrand = (((!rstrand && transcript.getStrand() == '+') || (rstrand && transcript.getStrand() == '-')) ? 0 : 0x0010);
 231                         b->core.flag += rstrand;
 232
 233                         push_qname(qname, b->core.l_qname, data);
 234                         int core_pos, core_n_cigar;
 235                         tr2chr(transcript, pos + 1, pos + readlen, core_pos, core_n_cigar, data);
 236                         if (core_pos < 0) b->core.tid = -1;
 237                         b->core.pos = core_pos;
 238                         b->core.n_cigar = core_n_cigar;
 239                         push_seq(seq, readlen, transcript.getStrand(), data);
 240                         push_qual(qual, readlen, transcript.getStrand(), data);
 241
 242                         free(b->data);
 243                         b->m_data = b->data_len = data.size();
 244                         b->l_aux = 0;
 245                         b->data = (uint8_t*)malloc(b->m_data);
 246                         for (int i = 0; i < b->data_len; i++) b->data[i] = data[i];
 247
 248                         b->core.bin = bam_reg2bin(b->core.pos, bam_calend(&(b->core), bam1_cigar(b)));
 249
 250                         char strand = transcript.getStrand();
 251                         bam_aux_append(b, "XS", 'A', 1, (uint8_t*)&strand);
 252
 253                 }
 254                 else {
 255                         b->m_data = b->data_len = b->data_len - b->l_aux;
 256                         b->l_aux = 0;
 257                         b->data = (uint8_t*)realloc(b->data, b->m_data);
 258                 }
 259
 260                 if (cqname != bam1_qname(b)) {
 261                         if (!hmap.empty()) {
 262                                 for (hmapIter = hmap.begin(); hmapIter != hmap.end(); hmapIter++) {
 263                                         bam1_t *tmp_b = hmapIter->first.b;
 264                                         tmp_b->core.qual = getMAPQ(hmapIter->second);
 265                                         float val = (float)hmapIter->second;
 266                                         bam_aux_append(tmp_b, "ZW", 'f', 4, (uint8_t*)&val);
 267                                         if (tmp_b->core.qual > 0) samwrite(out, tmp_b); // output only when MAPQ > 0
 268                                         bam_destroy1(tmp_b); // now hmapIter->b makes no sense
 269                                 }
 270                                 hmap.clear();
 271                         }
 272                         cqname = bam1_qname(b);
 273                 }
 274
 275                 hmapIter = hmap.find(SingleEndT(b));
 276                 if (hmapIter == hmap.end()) {
 277                         hmap[SingleEndT(bam_dup1(b))] = hit->getConPrb();
 278                 }
 279                 else {
 280                         hmapIter->second += hit->getConPrb();
 281                 }
 282         }
 283
 284         assert(wrapper.getNextHit() == NULL);
 285
 286         if (!hmap.empty()) {
 287                 for (hmapIter = hmap.begin(); hmapIter != hmap.end(); hmapIter++) {
 288                         bam1_t *tmp_b = hmapIter->first.b;
 289                         tmp_b->core.qual = getMAPQ(hmapIter->second);
 290                         float val = (float)hmapIter->second;
 291                         bam_aux_append(tmp_b, "ZW", 'f', 4, (uint8_t*)&val);
 292                         if (tmp_b->core.qual > 0) samwrite(out, tmp_b); // If MAPQ is equal to 0, do not output this alignment
 293                         bam_destroy1(tmp_b); // now hmapIter->b makes no sense
 294                 }
 295                 hmap.clear();
 296         }
 297
 298         bam_destroy1(b);
 299         if (verbose) { printf("Bam output file is generated!\n"); }
 300 }
 301
 302 void BamWriter::work(HitWrapper<PairedEndHit> wrapper, Transcripts& transcripts) {
 303         bam1_t *b, *b2;
 304         std::string cqname; // cqname : current query name
 305         std::map<PairedEndT, double> hmap;
 306         std::map<PairedEndT, double>::iterator hmapIter;
 307         PairedEndHit *hit;
 308
 309         int cnt = 0;
 310
 311         cqname = "";
 312         b = bam_init1();
 313         b2 = bam_init1();
 314         hmap.clear();
 315
 316         while (samread(in, b) >= 0 && samread(in, b2) >= 0) {
 317
 318                 if (verbose && cnt > 0 && cnt % 1000000 == 0) { printf("%d entries are finished!\n", cnt); }
 319                 ++cnt;
 320
 321                 if (!((b->core.flag & 0x0002) && (b2->core.flag & 0x0002))) continue;
 322
 323                 //swap if b is mate 2
 324                 if (b->core.flag & 0x0080) {
 325                         assert(b2->core.flag & 0x0040);
 326                         bam1_t *tmp = b;
 327                         b = b2; b2 = tmp;
 328                 }
 329
 330                 hit = wrapper.getNextHit();
 331                 assert(hit != NULL);
 332
 333                 int sid = b->core.tid + 1;
 334                 assert(sid == hit->getSid());
 335                 assert(sid == b2->core.tid + 1);
 336                 const Transcript& transcript = transcripts.getTranscriptAt(sid);
 337
 338                 if (transcripts.getType() == 0) {
 339                         int pos = b->core.pos, pos2 = b2->core.pos;
 340                         int readlen = b->core.l_qseq, readlen2 = b2->core.l_qseq;
 341                         uint8_t *qname = b->data, *seq = bam1_seq(b), *qual = bam1_qual(b);
 342                         uint8_t *qname2 = b2->data, *seq2 = bam1_seq(b2), *qual2 = bam1_qual(b2);
 343                         std::vector<uint8_t> data, data2;
 344
 345                         data.clear();
 346                         data2.clear();
 347
 348                         iter = refmap.find(transcript.getSeqName());
 349                         assert(iter != refmap.end());
 350                         b->core.tid = iter->second; b->core.mtid = iter->second;
 351                         b2->core.tid = iter->second; b2->core.mtid = iter->second;
 352
 353                         uint16_t rstrand = b->core.flag & 0x0010;
 354                         b->core.flag = b->core.flag - (b->core.flag & 0x0010) - (b->core.flag & 0x0020);
 355                         b2->core.flag = b2->core.flag - (b2->core.flag & 0x0010) - (b2->core.flag & 0x0020);
 356
 357                         uint16_t add, add2;
 358                         if ((!rstrand && transcript.getStrand() == '+') || (rstrand && transcript.getStrand() == '-')) {
 359                                 add = 0x0020; add2 = 0x0010;
 360                         }
 361                         else {
 362                                 add = 0x0010; add2 = 0x0020;
 363                         }
 364                         b->core.flag += add;
 365                         b2->core.flag += add2;
 366
 367                         b->core.qual = b2->core.qual = 255;
 368
 369                         if (transcript.getStrand() == '-') {
 370                                 b->core.isize = -b->core.isize;
 371                                 b2->core.isize = -b2->core.isize;
 372                         }
 373
 374                         push_qname(qname, b->core.l_qname, data);
 375                         push_qname(qname2, b2->core.l_qname, data2);
 376                         int core_pos, core_n_cigar;
 377                         tr2chr(transcript, pos + 1, pos + readlen, core_pos, core_n_cigar, data);
 378                         if (core_pos < 0) b->core.tid = -1;
 379                         b->core.pos = core_pos; b->core.n_cigar = core_n_cigar;
 380                         tr2chr(transcript, pos2 + 1, pos2 + readlen2, core_pos, core_n_cigar, data2);
 381                         if (core_pos < 0) b2->core.tid = -1;
 382                         b2->core.pos = core_pos; b2->core.n_cigar = core_n_cigar;
 383                         b->core.mpos = b2->core.pos;
 384                         b2->core.mpos = b->core.pos;
 385                         push_seq(seq, readlen, transcript.getStrand(), data);
 386                         push_seq(seq2, readlen2, transcript.getStrand(), data2);
 387                         push_qual(qual, readlen, transcript.getStrand(), data);
 388                         push_qual(qual2, readlen2, transcript.getStrand(), data2);
 389
 390                         free(b->data);
 391                         b->m_data = b->data_len = data.size();
 392                         b->l_aux = 0;
 393                         b->data = (uint8_t*)malloc(b->m_data);
 394                         for (int i = 0; i < b->data_len; i++) b->data[i] = data[i];
 395
 396                         free(b2->data);
 397                         b2->m_data = b2->data_len = data2.size();
 398                         b2->l_aux = 0;
 399                         b2->data = (uint8_t*)malloc(b2->m_data);
 400                         for (int i = 0; i < b2->data_len; i++) b2->data[i] = data2[i];
 401
 402                         b->core.bin = bam_reg2bin(b->core.pos, bam_calend(&(b->core), bam1_cigar(b)));
 403                         b2->core.bin = bam_reg2bin(b2->core.pos, bam_calend(&(b2->core), bam1_cigar(b2)));
 404
 405                         char strand = transcript.getStrand();
 406                         bam_aux_append(b, "XS", 'A', 1, (uint8_t*)&strand);
 407                         bam_aux_append(b2, "XS", 'A', 1, (uint8_t*)&strand);
 408                 }
 409                 else {
 410                         b->m_data = b->data_len = b->data_len - b->l_aux;
 411                         b->l_aux = 0;
 412                         b->data = (uint8_t*)realloc(b->data, b->m_data);
 413
 414                         b2->m_data = b2->data_len = b2->data_len - b2->l_aux;
 415                         b2->l_aux = 0;
 416                         b2->data = (uint8_t*)realloc(b2->data, b2->m_data);
 417                 }
 418
 419                 if (cqname != bam1_qname(b)) {
 420                         if (!hmap.empty()) {
 421                                 for (hmapIter = hmap.begin(); hmapIter != hmap.end(); hmapIter++) {
 422                                         bam1_t *tmp_b = hmapIter->first.b;
 423                                         bam1_t *tmp_b2 = hmapIter->first.b2;
 424
 425                                         tmp_b->core.qual = tmp_b2->core.qual = getMAPQ(hmapIter->second);
 426
 427                                         float val = (float)hmapIter->second;
 428                                         bam_aux_append(tmp_b, "ZW", 'f', 4, (uint8_t*)&val);
 429                                         bam_aux_append(tmp_b2, "ZW", 'f', 4, (uint8_t*)&val);
 430
 431                                         // If MAPQ is equal to 0, do not output this alignment pair
 432                                         if (tmp_b->core.qual > 0) {
 433                                           samwrite(out, tmp_b);
 434                                           samwrite(out, tmp_b2);
 435                                         }
 436
 437                                         bam_destroy1(tmp_b);
 438                                         bam_destroy1(tmp_b2);
 439                                 }
 440                                 hmap.clear();
 441                         }
 442                         cqname = bam1_qname(b);
 443                 }
 444
 445                 hmapIter = hmap.find(PairedEndT(b, b2));
 446                 if (hmapIter == hmap.end()) {
 447                         hmap[PairedEndT(bam_dup1(b), bam_dup1(b2))] = hit->getConPrb();
 448                 }
 449                 else {
 450                         hmapIter->second += hit->getConPrb();
 451                 }
 452         }
 453
 454         assert(wrapper.getNextHit() == NULL);
 455
 456         if (!hmap.empty()) {
 457                 for (hmapIter = hmap.begin(); hmapIter != hmap.end(); hmapIter++) {
 458                         bam1_t *tmp_b = hmapIter->first.b;
 459                         bam1_t *tmp_b2 = hmapIter->first.b2;
 460
 461                         tmp_b->core.qual = tmp_b2->core.qual = getMAPQ(hmapIter->second);
 462
 463                         float val = (float)hmapIter->second;
 464                         bam_aux_append(tmp_b, "ZW", 'f', 4, (uint8_t*)&val);
 465                         bam_aux_append(tmp_b2, "ZW", 'f', 4, (uint8_t*)&val);
 466
 467                         if (tmp_b->core.qual > 0) {
 468                           samwrite(out, tmp_b);
 469                           samwrite(out, tmp_b2);
 470                         }
 471
 472                         bam_destroy1(tmp_b);
 473                         bam_destroy1(tmp_b2);
 474                 }
 475                 hmap.clear();
 476         }
 477
 478         bam_destroy1(b);
 479         bam_destroy1(b2);
 480
 481         if (verbose) { printf("Bam output file is generated!\n"); }
 482 }
 483
 484 void BamWriter::tr2chr(const Transcript& transcript, int sp, int ep, int& pos, int& n_cigar, std::vector<uint8_t>& data) {
 485         int length = transcript.getLength();
 486         char strand = transcript.getStrand();
 487         const std::vector<Interval>& structure = transcript.getStructure();
 488
 489         int s, i;
 490         int oldlen, curlen;
 491
 492         uint32_t operation;
 493         uint8_t *p;
 494
 495         n_cigar = 0;
 496         s = structure.size();
 497
 498         if (strand == '-') {
 499                 int tmp = sp;
 500                 sp = length - ep + 1;
 501                 ep = length - tmp + 1;
 502         }
 503
 504         if (ep < 1 || sp > length) { // a read which align to polyA tails totally!
 505           pos = (sp > length ? structure[s - 1].end : structure[0].start - 1); // 0 based
 506
 507           n_cigar = 1;
 508           operation = (ep - sp + 1) << BAM_CIGAR_SHIFT | BAM_CINS; //BAM_CSOFT_CLIP;
 509           p = (uint8_t*)(&operation);
 510           for (int j = 0; j < 4; j++) data.push_back(*(p + j));
 511
 512           return;
 513         }
 514
 515         if (sp < 1) {
 516                 n_cigar++;
 517                 operation = (1 - sp) << BAM_CIGAR_SHIFT | BAM_CINS; //BAM_CSOFT_CLIP;
 518                 p = (uint8_t*)(&operation);
 519                 for (int j = 0; j < 4; j++) data.push_back(*(p + j));
 520                 sp = 1;
 521         }
 522
 523         oldlen = curlen = 0;
 524
 525         for (i = 0; i < s; i++) {
 526                 oldlen = curlen;
 527                 curlen += structure[i].end - structure[i].start + 1;
 528                 if (curlen >= sp) break;
 529         }
 530         assert(i < s);
 531         pos = structure[i].start + (sp - oldlen - 1) - 1; // 0 based
 532
 533         while (curlen < ep && i < s) {
 534                 n_cigar++;
 535                 operation = (curlen - sp + 1) << BAM_CIGAR_SHIFT | BAM_CMATCH;
 536                 p = (uint8_t*)(&operation);
 537                 for (int j = 0; j < 4; j++) data.push_back(*(p + j));
 538
 539                 ++i;
 540                 if (i >= s) continue;
 541                 n_cigar++;
 542                 operation = (structure[i].start - structure[i - 1].end - 1) << BAM_CIGAR_SHIFT | BAM_CREF_SKIP;
 543                 p = (uint8_t*)(&operation);
 544                 for (int j = 0; j < 4; j++) data.push_back(*(p + j));
 545
 546                 oldlen = curlen;
 547                 sp = oldlen + 1;
 548                 curlen += structure[i].end - structure[i].start + 1;
 549         }
 550
 551         if (i >= s) {
 552                 n_cigar++;
 553                 operation = (ep - length) << BAM_CIGAR_SHIFT | BAM_CINS; //BAM_CSOFT_CLIP;
 554                 p = (uint8_t*)(&operation);
 555                 for (int j = 0; j < 4; j++) data.push_back(*(p + j));
 556         }
 557         else {
 558                 n_cigar++;
 559                 operation = (ep - sp + 1) << BAM_CIGAR_SHIFT | BAM_CMATCH;
 560                 p = (uint8_t*)(&operation);
 561                 for (int j = 0; j < 4; j++) data.push_back(*(p + j));
 562         }
 563 }
 564
 565 #endif /* BAMWRITER_H_ */