]> git.donarmstrong.com Git - rsem.git/blob - BamWriter.h
RSEM Source Codes
[rsem.git] / 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         out = samopen(outF, "wb", out_header);
178         assert(out != 0);
179
180         if (chr_list != NULL) { bam_header_destroy(out_header); }
181 }
182
183 BamWriter::~BamWriter() {
184         samclose(in);
185         samclose(out);
186 }
187
188 void BamWriter::work(HitWrapper<SingleHit> wrapper, Transcripts& transcripts) {
189         bam1_t *b;
190         std::string cqname; // cqname : current query name
191         std::map<SingleEndT, double> hmap;
192         std::map<SingleEndT, double>::iterator hmapIter;
193         SingleHit *hit;
194
195         int cnt = 0;
196
197         cqname = "";
198         b = bam_init1();
199         hmap.clear();
200
201         while (samread(in, b) >= 0) {
202
203                 if (verbose && cnt > 0 && cnt % 1000000 == 0) { printf("%d entries are finished!\n", cnt); }
204                 ++cnt;
205
206                 if (b->core.flag & 0x0004) continue;
207
208                 hit = wrapper.getNextHit();
209                 assert(hit != NULL);
210
211                 int sid = b->core.tid + 1;
212                 assert(sid == hit->getSid());
213                 const Transcript& transcript = transcripts.getTranscriptAt(sid);
214
215                 if (transcripts.getType() == 0) {
216                         int pos = b->core.pos;
217                         int readlen = b->core.l_qseq;
218                         uint8_t *qname = b->data, *seq = bam1_seq(b), *qual = bam1_qual(b);
219                         std::vector<uint8_t> data;
220                         data.clear();
221
222                         iter = refmap.find(transcript.getSeqName());
223                         assert(iter != refmap.end());
224                         b->core.tid = iter->second;
225                         b->core.qual = 255;
226
227                         uint16_t rstrand = b->core.flag & 0x0010; // read strand
228                         b->core.flag -= rstrand;
229                         rstrand = (!rstrand && transcript.getStrand() == '+' || rstrand && transcript.getStrand() == '-' ? 0 : 0x0010);
230                         b->core.flag += rstrand;
231
232                         push_qname(qname, b->core.l_qname, data);
233                         int core_pos, core_n_cigar;
234                         tr2chr(transcript, pos + 1, pos + readlen, core_pos, core_n_cigar, data);
235                         if (core_pos < 0) b->core.tid = -1;
236                         b->core.pos = core_pos;
237                         b->core.n_cigar = core_n_cigar;
238                         push_seq(seq, readlen, transcript.getStrand(), data);
239                         push_qual(qual, readlen, transcript.getStrand(), data);
240
241                         free(b->data);
242                         b->m_data = b->data_len = data.size() + 7; // 7 extra bytes for ZW tag
243                         b->l_aux = 7;
244                         b->data = (uint8_t*)malloc(b->m_data);
245                         for (int i = 0; i < b->data_len; i++) b->data[i] = data[i];
246
247                         b->core.bin = bam_reg2bin(b->core.pos, bam_calend(&(b->core), bam1_cigar(b)));
248                 }
249                 else {
250                         b->m_data = b->data_len = b->data_len - b->l_aux + 7; // 7 extra bytes for ZW tag
251                         b->l_aux = 7;
252                         b->data = (uint8_t*)realloc(b->data, b->m_data);
253                 }
254
255
256                 if (cqname != bam1_qname(b)) {
257                         if (!hmap.empty()) {
258                                 for (hmapIter = hmap.begin(); hmapIter != hmap.end(); hmapIter++) {
259                                         bam1_t *tmp_b = hmapIter->first.b;
260                                         tmp_b->core.qual = getMAPQ(hmapIter->second);
261                                         uint8_t *p = bam1_aux(tmp_b);
262                                         *p = 'Z'; ++p; *p = 'W'; ++p; *p = 'f'; ++p;
263                                         float val = (float)hmapIter->second;
264                                         memcpy(p, &val, 4);
265                                         samwrite(out, tmp_b);
266                                         bam_destroy1(tmp_b); // now hmapIter->b makes no sense
267                                 }
268                                 hmap.clear();
269                         }
270                         cqname = bam1_qname(b);
271                 }
272
273                 hmapIter = hmap.find(SingleEndT(b));
274                 if (hmapIter == hmap.end()) {
275                         hmap[SingleEndT(bam_dup1(b))] = hit->getConPrb();
276                 }
277                 else {
278                         hmapIter->second += hit->getConPrb();
279                 }
280         }
281
282         assert(wrapper.getNextHit() == NULL);
283
284         if (!hmap.empty()) {
285                 for (hmapIter = hmap.begin(); hmapIter != hmap.end(); hmapIter++) {
286                         bam1_t *tmp_b = hmapIter->first.b;
287                         tmp_b->core.qual = getMAPQ(hmapIter->second);
288                         uint8_t *p = bam1_aux(tmp_b);
289                         *p = 'Z'; ++p; *p = 'W'; ++p; *p = 'f'; ++p;
290                         float val = (float)hmapIter->second;
291                         memcpy(p, &val, 4);
292                         samwrite(out, tmp_b);
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                         //Do I really need this? The insert size uses transcript coordinates
370                         if (transcript.getStrand() == '-') {
371                                 b->core.isize = -b->core.isize;
372                                 b2->core.isize = -b2->core.isize;
373                         }
374
375                         push_qname(qname, b->core.l_qname, data);
376                         push_qname(qname2, b2->core.l_qname, data2);
377                         int core_pos, core_n_cigar;
378                         tr2chr(transcript, pos + 1, pos + readlen, core_pos, core_n_cigar, data);
379                         if (core_pos < 0) b->core.tid = -1;
380                         b->core.pos = core_pos; b->core.n_cigar = core_n_cigar;
381                         tr2chr(transcript, pos2 + 1, pos2 + readlen2, core_pos, core_n_cigar, data2);
382                         if (core_pos < 0) b2->core.tid = -1;
383                         b2->core.pos = core_pos; b2->core.n_cigar = core_n_cigar;
384                         b->core.mpos = b2->core.pos;
385                         b2->core.mpos = b->core.pos;
386                         push_seq(seq, readlen, transcript.getStrand(), data);
387                         push_seq(seq2, readlen2, transcript.getStrand(), data2);
388                         push_qual(qual, readlen, transcript.getStrand(), data);
389                         push_qual(qual2, readlen2, transcript.getStrand(), data2);
390
391                         free(b->data);
392                         b->m_data = b->data_len = data.size() + 7; // 7 extra bytes for ZW tag
393                         b->l_aux = 7;
394                         b->data = (uint8_t*)malloc(b->m_data);
395                         for (int i = 0; i < b->data_len; i++) b->data[i] = data[i];
396
397                         free(b2->data);
398                         b2->m_data = b2->data_len = data2.size() + 7; // 7 extra bytes for ZW tag
399                         b2->l_aux = 7;
400                         b2->data = (uint8_t*)malloc(b2->m_data);
401                         for (int i = 0; i < b2->data_len; i++) b2->data[i] = data2[i];
402
403                         b->core.bin = bam_reg2bin(b->core.pos, bam_calend(&(b->core), bam1_cigar(b)));
404                         b2->core.bin = bam_reg2bin(b2->core.pos, bam_calend(&(b2->core), bam1_cigar(b2)));
405                 }
406                 else {
407                         b->m_data = b->data_len = b->data_len - b->l_aux + 7; // 7 extra bytes for ZW tag
408                         b->l_aux = 7;
409                         b->data = (uint8_t*)realloc(b->data, b->m_data);
410
411                         b2->m_data = b2->data_len = b2->data_len - b2->l_aux + 7; // 7 extra bytes for ZW tag
412                         b2->l_aux = 7;
413                         b2->data = (uint8_t*)realloc(b2->data, b2->m_data);
414                 }
415
416                 if (cqname != bam1_qname(b)) {
417                         if (!hmap.empty()) {
418                                 for (hmapIter = hmap.begin(); hmapIter != hmap.end(); hmapIter++) {
419                                         bam1_t *tmp_b = hmapIter->first.b;
420                                         bam1_t *tmp_b2 = hmapIter->first.b2;
421
422                                         tmp_b->core.qual = tmp_b2->core.qual = getMAPQ(hmapIter->second);
423
424                                         uint8_t *p = bam1_aux(tmp_b), *p2 = bam1_aux(tmp_b2);
425                                         *p = 'Z'; ++p; *p = 'W'; ++p; *p = 'f'; ++p;
426                                         *p2 = 'Z'; ++p2; *p2 = 'W'; ++p2; *p2 = 'f'; ++p2;
427
428                                         float val = (float)hmapIter->second;
429                                         memcpy(p, &val, 4);
430                                         memcpy(p2, &val, 4);
431
432                                         samwrite(out, tmp_b);
433                                         samwrite(out, tmp_b2);
434
435                                         bam_destroy1(tmp_b);
436                                         bam_destroy1(tmp_b2);
437                                 }
438                                 hmap.clear();
439                         }
440                         cqname = bam1_qname(b);
441                 }
442
443                 hmapIter = hmap.find(PairedEndT(b, b2));
444                 if (hmapIter == hmap.end()) {
445                         hmap[PairedEndT(bam_dup1(b), bam_dup1(b2))] = hit->getConPrb();
446                 }
447                 else {
448                         hmapIter->second += hit->getConPrb();
449                 }
450         }
451
452         assert(wrapper.getNextHit() == NULL);
453
454         if (!hmap.empty()) {
455                 for (hmapIter = hmap.begin(); hmapIter != hmap.end(); hmapIter++) {
456                         bam1_t *tmp_b = hmapIter->first.b;
457                         bam1_t *tmp_b2 = hmapIter->first.b2;
458
459                         tmp_b->core.qual = tmp_b2->core.qual = getMAPQ(hmapIter->second);
460
461                         uint8_t *p = bam1_aux(tmp_b), *p2 = bam1_aux(tmp_b2);
462                         *p = 'Z'; ++p; *p = 'W'; ++p; *p = 'f'; ++p;
463                         *p2 = 'Z'; ++p2; *p2 = 'W'; ++p2; *p2 = 'f'; ++p2;
464
465                         float val = (float)hmapIter->second;
466                         memcpy(p, &val, 4);
467                         memcpy(p2, &val, 4);
468
469                         samwrite(out, tmp_b);
470                         samwrite(out, tmp_b2);
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!"); }
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_ */