]> git.donarmstrong.com Git - rsem.git/blob - BamConverter.h
The order of @SQ tags in SAM/BAM files can be arbitrary now
[rsem.git] / BamConverter.h
1 #ifndef BAMCONVERTER_H_
2 #define BAMCONVERTER_H_
3
4 #include<cstdio>
5 #include<cstring>
6 #include<cassert>
7 #include<string>
8 #include<map>
9
10 #include <stdint.h>
11 #include "sam/bam.h"
12 #include "sam/sam.h"
13 #include "sam_rsem_aux.h"
14 #include "sam_rsem_cvt.h"
15
16 #include "utils.h"
17 #include "my_assert.h"
18 #include "bc_aux.h"
19 #include "Transcript.h"
20 #include "Transcripts.h"
21
22 class BamConverter {
23 public:
24         BamConverter(const char*, const char*, const char*, Transcripts&);
25         ~BamConverter();
26
27         void process();
28 private:
29         samfile_t *in, *out;
30         Transcripts& transcripts;
31
32         std::map<std::string, int> refmap;
33         std::map<std::string, int>::iterator iter;
34
35         CollapseMap collapseMap;
36
37         void convert(bam1_t*, const Transcript&);
38
39         void writeCollapsedLines();
40         void flipSeq(uint8_t*, int);
41         void flipQual(uint8_t*, int);
42         void addXSTag(bam1_t*, const Transcript&);
43 };
44
45 BamConverter::BamConverter(const char* inpF, const char* outF, const char* chr_list, Transcripts& transcripts)
46         : transcripts(transcripts)
47 {
48         general_assert(transcripts.getType() == 0, "Genome information is not provided! RSEM cannot convert the transcript bam file!");
49
50         in = samopen(inpF, "rb", NULL);
51         assert(in != 0);
52
53         transcripts.buildMappings(in->header->n_targets, in->header->target_name);
54
55         bam_header_t *out_header = sam_header_read2(chr_list);
56         refmap.clear();
57         for (int i = 0; i < out_header->n_targets; i++) {
58                 refmap[out_header->target_name[i]] = i;
59         }
60
61         append_header_text(out_header, in->header->text, in->header->l_text);
62
63         out = samopen(outF, "wb", out_header);
64         assert(out != 0);
65
66         bam_header_destroy(out_header);
67 }
68
69 BamConverter::~BamConverter() {
70         samclose(in);
71         samclose(out);
72 }
73
74 void BamConverter::process() {
75         bam1_t *b, *b2;
76         std::string cqname;
77         bool isPaired = false;
78
79         int cnt = 0;
80
81         cqname = "";
82         b = bam_init1(); b2 = bam_init1();
83
84         while (samread(in, b) >= 0) {
85                 ++cnt;
86                 isPaired = (b->core.flag & 0x0001) > 0;
87                 if (isPaired) {
88                         assert(samread(in, b2) >= 0 && (b2->core.flag & 0x0001) && b->core.tid == b2->core.tid);
89                         assert((b->core.flag & 0x0040) && (b2->core.flag & 0x0080)); // for collapsing
90                         ++cnt;
91                 }
92
93                 if (cnt % 1000000 == 0) { printf("."); fflush(stdout); }
94
95                 // at least one segment is not properly mapped
96                 if ((b->core.flag & 0x0004) || (isPaired && (b2->core.flag & 0x0004))) continue;
97
98                 const Transcript& transcript = transcripts.getTranscriptViaEid(b->core.tid + 1);
99
100                 convert(b, transcript);
101                 if (isPaired) {
102                         convert(b2, transcript);
103                         b->core.mpos = b2->core.pos;
104                         b2->core.mpos = b->core.pos;
105                 }
106
107                 if (cqname != bam1_qname(b)) {
108                         writeCollapsedLines();
109                         cqname = bam1_qname(b);
110                         collapseMap.init(isPaired);
111                 }
112
113                 collapseMap.insert(b, b2, bam_aux2f(bam_aux_get(b, "ZW")));
114         }
115
116         writeCollapsedLines();
117
118         bam_destroy1(b);
119         bam_destroy1(b2);
120
121         if (cnt >= 1000000) printf("\n");
122 }
123
124 void BamConverter::convert(bam1_t* b, const Transcript& transcript) {
125         int pos = b->core.pos;
126         int readlen = b->core.l_qseq;
127
128         general_assert(readlen > 0, "One alignment line has SEQ field as *. RSEM does not support this currently!");
129
130         iter = refmap.find(transcript.getSeqName());
131         assert(iter != refmap.end());
132         b->core.tid = iter->second;
133         if (b->core.flag & 0x0001) { b->core.mtid = b->core.tid; }
134         b->core.qual = 255; // set to not available temporarily
135
136         if (transcript.getStrand() == '-') {
137                 b->core.flag ^= 0x0010;
138                 if (b->core.flag & 0x0001) {
139                         b->core.flag ^= 0x0020;
140                         b->core.isize = -b->core.isize;
141                 }
142                 flipSeq(bam1_seq(b), readlen);
143                 flipQual(bam1_qual(b), readlen);
144         }
145
146         std::vector<uint32_t> data;
147         data.clear();
148
149         int core_pos, core_n_cigar;
150         tr2chr(transcript, pos + 1, pos + readlen, core_pos, core_n_cigar, data);
151         assert(core_pos >= 0);
152
153         int rest_len = b->data_len - b->core.l_qname - b->core.n_cigar * 4;
154         b->data_len = b->core.l_qname + core_n_cigar * 4 + rest_len;
155         expand_data_size(b);
156         uint8_t* pt = b->data + b->core.l_qname;
157         memmove(pt + core_n_cigar * 4, pt + b->core.n_cigar * 4, rest_len);
158         for (int i = 0; i < core_n_cigar; i++) { memmove(pt, &data[i], 4); pt += 4; }
159
160         b->core.pos = core_pos;
161         b->core.n_cigar = core_n_cigar;
162         b->core.bin = bam_reg2bin(b->core.pos, bam_calend(&(b->core), bam1_cigar(b)));
163
164         addXSTag(b, transcript); // check if need to add XS tag, if need, add it
165 }
166
167 inline void BamConverter::writeCollapsedLines() {
168         bam1_t *tmp_b = NULL,*tmp_b2 = NULL;
169         float prb;
170         bool isPaired;
171
172         if (!collapseMap.empty(isPaired)) {
173                 while (collapseMap.next(tmp_b, tmp_b2, prb)) {
174                         memcpy(bam_aux_get(tmp_b, "ZW") + 1, (uint8_t*)&(prb), bam_aux_type2size('f'));
175                         tmp_b->core.qual = getMAPQ(prb);
176                         if (tmp_b->core.qual > 0) {
177                                 samwrite(out, tmp_b);
178                                 if (isPaired) {
179                                         memcpy(bam_aux_get(tmp_b2, "ZW") + 1, (uint8_t*)&(prb), bam_aux_type2size('f'));
180                                         tmp_b2->core.qual = tmp_b->core.qual;
181                                         samwrite(out, tmp_b2);
182                                 }
183                         }
184                         bam_destroy1(tmp_b);
185                         if (isPaired) bam_destroy1(tmp_b2);
186
187                 }
188         }
189 }
190
191 inline void BamConverter::flipSeq(uint8_t* s, int readlen) {
192         uint8_t code, base;
193         std::vector<uint8_t> seq;
194
195         code = 0; base = 0;
196         seq.clear();
197         for (int i = 0; i < readlen; i++) {
198                 switch (bam1_seqi(s, readlen - i - 1)) {
199                 case 1: base = 8; break;
200                 case 2: base = 4; break;
201                 case 4: base = 2; break;
202                 case 8: base = 1; break;
203                 case 15: base = 15; break;
204                 default: assert(false);
205                 }
206                 code |=  base << (4 * (1 - i % 2));
207                 if (i % 2 == 1) { seq.push_back(code); code = 0; }
208         }
209         if (readlen % 2 == 1) { seq.push_back(code); }
210
211         for (int i = 0; i < (int)seq.size(); i++) s[i] = seq[i];
212 }
213
214 inline void BamConverter::flipQual(uint8_t* q, int readlen) {
215         int32_t mid = readlen / 2;
216         uint8_t tmp;
217         for (int i = 0; i < mid; i++) {
218                 tmp = q[i]; q[i] = q[readlen - i - 1]; q[readlen -i -1] = tmp;
219         }
220 }
221
222 inline void BamConverter::addXSTag(bam1_t* b, const Transcript& transcript) {
223         uint32_t* p = bam1_cigar(b);
224         bool hasN = false;
225         for (int i = 0; i < (int)b->core.n_cigar; i++)
226                 if ((*(p + i) & BAM_CIGAR_MASK) == BAM_CREF_SKIP) { hasN = true; break; }
227         if (!hasN) return;
228         char strand = transcript.getStrand();
229         bam_aux_append(b, "XS", 'A', 1, (uint8_t*)&strand);
230 }
231
232 #endif /* BAMCONVERTER_H_ */