Imported Upstream version 0.6
[pysam.git] / samtools / bcftools / bcfutils.c.pysam.c
1 #include "pysam.h"
2
3 #include <string.h>
4 #include <math.h>
5 #include "bcf.h"
6 #include "kstring.h"
7 #include "khash.h"
8 KHASH_MAP_INIT_STR(str2id, int)
9
10 #ifdef _WIN32
11 #define srand48(x) srand(x)
12 #define drand48() ((double)rand() / RAND_MAX)
13 #endif
14
15 // FIXME: valgrind report a memory leak in this function. Probably it does not get deallocated...
16 void *bcf_build_refhash(bcf_hdr_t *h)
17 {
18         khash_t(str2id) *hash;
19         int i, ret;
20         hash = kh_init(str2id);
21         for (i = 0; i < h->n_ref; ++i) {
22                 khint_t k;
23                 k = kh_put(str2id, hash, h->ns[i], &ret); // FIXME: check ret
24                 kh_val(hash, k) = i;
25         }
26         return hash;
27 }
28
29 void *bcf_str2id_init()
30 {
31         return kh_init(str2id);
32 }
33
34 void bcf_str2id_destroy(void *_hash)
35 {
36         khash_t(str2id) *hash = (khash_t(str2id)*)_hash;
37         if (hash) kh_destroy(str2id, hash); // Note that strings are not freed.
38 }
39
40 void bcf_str2id_thorough_destroy(void *_hash)
41 {
42         khash_t(str2id) *hash = (khash_t(str2id)*)_hash;
43         khint_t k;
44         if (hash == 0) return;
45         for (k = 0; k < kh_end(hash); ++k)
46                 if (kh_exist(hash, k)) free((char*)kh_key(hash, k));
47         kh_destroy(str2id, hash);
48 }
49
50 int bcf_str2id(void *_hash, const char *str)
51 {
52         khash_t(str2id) *hash = (khash_t(str2id)*)_hash;
53         khint_t k;
54         if (!hash) return -1;
55         k = kh_get(str2id, hash, str);
56         return k == kh_end(hash)? -1 : kh_val(hash, k);
57 }
58
59 int bcf_str2id_add(void *_hash, const char *str)
60 {
61         khint_t k;
62         int ret;
63         khash_t(str2id) *hash = (khash_t(str2id)*)_hash;
64         if (!hash) return -1;
65         k = kh_put(str2id, hash, str, &ret);
66         if (ret == 0) return kh_val(hash, k);
67         kh_val(hash, k) = kh_size(hash) - 1;
68         return kh_val(hash, k);
69 }
70
71 int bcf_shrink_alt(bcf1_t *b, int n)
72 {
73         char *p;
74         int i, j, k, n_smpl = b->n_smpl;
75         if (b->n_alleles <= n) return -1;
76         // update ALT
77         if (n > 1) {
78                 for (p = b->alt, k = 1; *p; ++p)
79                         if (*p == ',' && ++k == n) break;
80                 *p = '\0';
81         } else p = b->alt, *p = '\0';
82         ++p;
83         memmove(p, b->flt, b->str + b->l_str - b->flt);
84         b->l_str -= b->flt - p;
85         // update PL
86         for (i = 0; i < b->n_gi; ++i) {
87                 bcf_ginfo_t *g = b->gi + i;
88                 if (g->fmt == bcf_str2int("PL", 2)) {
89                         int l, x = b->n_alleles * (b->n_alleles + 1) / 2;
90                         uint8_t *d = (uint8_t*)g->data;
91                         g->len = n * (n + 1) / 2;
92                         for (l = k = 0; l < n_smpl; ++l) {
93                                 uint8_t *dl = d + l * x;
94                                 for (j = 0; j < g->len; ++j) d[k++] = dl[j];
95                         }
96                 } // FIXME: to add GL
97         }
98         b->n_alleles = n;
99         bcf_sync(b);
100         return 0;
101 }
102
103 int bcf_gl2pl(bcf1_t *b)
104 {
105         char *p;
106         int i, n_smpl = b->n_smpl;
107         bcf_ginfo_t *g;
108         float *d0;
109         uint8_t *d1;
110         if (strstr(b->fmt, "PL")) return -1;
111         if ((p = strstr(b->fmt, "GL")) == 0) return -1;
112         *p = 'P';
113         for (i = 0; i < b->n_gi; ++i)
114                 if (b->gi[i].fmt == bcf_str2int("GL", 2))
115                         break;
116         g = b->gi + i;
117         g->fmt = bcf_str2int("PL", 2);
118         g->len /= 4; // 4 == sizeof(float)
119         d0 = (float*)g->data; d1 = (uint8_t*)g->data;
120         for (i = 0; i < n_smpl * g->len; ++i) {
121                 int x = (int)(-10. * d0[i] + .499);
122                 if (x > 255) x = 255;
123                 if (x < 0) x = 0;
124                 d1[i] = x;
125         }
126         return 0;
127 }
128 /* FIXME: this function will fail given AB:GTX:GT. BCFtools never
129  * produces such FMT, but others may do. */
130 int bcf_fix_gt(bcf1_t *b)
131 {
132         char *s;
133         int i;
134         uint32_t tmp;
135         bcf_ginfo_t gt;
136         // check the presence of the GT FMT
137         if ((s = strstr(b->fmt, ":GT")) == 0) return 0; // no GT or GT is already the first
138         if (s[3] != '\0' && s[3] != ':') return 0; // :GTX in fact
139         tmp = bcf_str2int("GT", 2);
140         for (i = 0; i < b->n_gi; ++i)
141                 if (b->gi[i].fmt == tmp) break;
142         if (i == b->n_gi) return 0; // no GT in b->gi; probably a bug...
143         gt = b->gi[i];
144         // move GT to the first
145         for (; i > 0; --i) b->gi[i] = b->gi[i-1];
146         b->gi[0] = gt;
147         memmove(b->fmt + 3, b->fmt, s + 1 - b->fmt);
148         b->fmt[0] = 'G'; b->fmt[1] = 'T'; b->fmt[2] = ':';
149         return 0;
150 }
151
152 int bcf_fix_pl(bcf1_t *b)
153 {
154         int i;
155         uint32_t tmp;
156         uint8_t *PL, *swap;
157         bcf_ginfo_t *gi;
158         // pinpoint PL
159         tmp = bcf_str2int("PL", 2);
160         for (i = 0; i < b->n_gi; ++i)
161                 if (b->gi[i].fmt == tmp) break;
162         if (i == b->n_gi) return 0;
163         // prepare
164         gi = b->gi + i;
165         PL = (uint8_t*)gi->data;
166         swap = alloca(gi->len);
167         // loop through individuals
168         for (i = 0; i < b->n_smpl; ++i) {
169                 int k, l, x;
170                 uint8_t *PLi = PL + i * gi->len;
171                 memcpy(swap, PLi, gi->len);
172                 for (k = x = 0; k < b->n_alleles; ++k)
173                         for (l = k; l < b->n_alleles; ++l)
174                                 PLi[l*(l+1)/2 + k] = swap[x++];
175         }
176         return 0;
177 }
178
179 int bcf_smpl_covered(const bcf1_t *b)
180 {
181         int i, j, n = 0;
182         uint32_t tmp;
183         bcf_ginfo_t *gi;
184         // pinpoint PL
185         tmp = bcf_str2int("PL", 2);
186         for (i = 0; i < b->n_gi; ++i)
187                 if (b->gi[i].fmt == tmp) break;
188         if (i == b->n_gi) return 0;
189         // count how many samples having PL!=[0..0]
190         gi = b->gi + i;
191         for (i = 0; i < b->n_smpl; ++i) {
192                 uint8_t *PLi = ((uint8_t*)gi->data) + i * gi->len;
193                 for (j = 0; j < gi->len; ++j)
194                         if (PLi[j]) break;
195                 if (j < gi->len) ++n;
196         }
197         return n;
198 }
199
200 static void *locate_field(const bcf1_t *b, const char *fmt, int l)
201 {
202         int i;
203         uint32_t tmp;
204         tmp = bcf_str2int(fmt, l);
205         for (i = 0; i < b->n_gi; ++i)
206                 if (b->gi[i].fmt == tmp) break;
207         return i == b->n_gi? 0 : b->gi[i].data;
208 }
209
210 int bcf_anno_max(bcf1_t *b)
211 {
212         int k, max_gq, max_sp, n_het;
213         kstring_t str;
214         uint8_t *gt, *gq;
215         int32_t *sp;
216         max_gq = max_sp = n_het = 0;
217         gt = locate_field(b, "GT", 2);
218         if (gt == 0) return -1;
219         gq = locate_field(b, "GQ", 2);
220         sp = locate_field(b, "SP", 2);
221         if (sp)
222                 for (k = 0; k < b->n_smpl; ++k)
223                         if (gt[k]&0x3f)
224                                 max_sp = max_sp > (int)sp[k]? max_sp : sp[k];
225         if (gq)
226                 for (k = 0; k < b->n_smpl; ++k)
227                         if (gt[k]&0x3f)
228                                 max_gq = max_gq > (int)gq[k]? max_gq : gq[k];
229         for (k = 0; k < b->n_smpl; ++k) {
230                 int a1, a2;
231                 a1 = gt[k]&7; a2 = gt[k]>>3&7;
232                 if ((!a1 && a2) || (!a2 && a1)) { // a het
233                         if (gq == 0) ++n_het;
234                         else if (gq[k] >= 20) ++n_het;
235                 }
236         }
237         if (n_het) max_sp -= (int)(4.343 * log(n_het) + .499);
238         if (max_sp < 0) max_sp = 0;
239         memset(&str, 0, sizeof(kstring_t));
240         if (*b->info) kputc(';', &str);
241         ksprintf(&str, "MXSP=%d;MXGQ=%d", max_sp, max_gq);
242         bcf_append_info(b, str.s, str.l);
243         free(str.s);
244         return 0;
245 }
246
247 // FIXME: only data are shuffled; the header is NOT
248 int bcf_shuffle(bcf1_t *b, int seed)
249 {
250         int i, j, *a;
251         if (seed > 0) srand48(seed);
252         a = malloc(b->n_smpl * sizeof(int));
253         for (i = 0; i < b->n_smpl; ++i) a[i] = i;
254         for (i = b->n_smpl; i > 1; --i) {
255                 int tmp;
256                 j = (int)(drand48() * i);
257                 tmp = a[j]; a[j] = a[i-1]; a[i-1] = tmp;
258         }
259         for (j = 0; j < b->n_gi; ++j) {
260                 bcf_ginfo_t *gi = b->gi + j;
261                 uint8_t *swap, *data = (uint8_t*)gi->data;
262                 swap = malloc(gi->len * b->n_smpl);
263                 for (i = 0; i < b->n_smpl; ++i)
264                         memcpy(swap + gi->len * a[i], data + gi->len * i, gi->len);
265                 free(gi->data);
266                 gi->data = swap;
267         }
268         free(a);
269         return 0;
270 }
271
272 bcf_hdr_t *bcf_hdr_subsam(const bcf_hdr_t *h0, int n, char *const* samples, int *list)
273 {
274         int i, ret, j;
275         khint_t k;
276         bcf_hdr_t *h;
277         khash_t(str2id) *hash;
278         kstring_t s;
279         s.l = s.m = 0; s.s = 0;
280         hash = kh_init(str2id);
281         for (i = 0; i < h0->n_smpl; ++i) {
282                 k = kh_put(str2id, hash, h0->sns[i], &ret);
283                 kh_val(hash, k) = i;
284         }
285         for (i = j = 0; i < n; ++i) {
286                 k = kh_get(str2id, hash, samples[i]);
287                 if (k != kh_end(hash)) {
288                         list[j++] = kh_val(hash, k);
289                         kputs(samples[i], &s); kputc('\0', &s);
290                 }
291         }
292         if (j < n) fprintf(pysamerr, "<%s> %d samples in the list but not in BCF.", __func__, n - j);
293         kh_destroy(str2id, hash);
294         h = calloc(1, sizeof(bcf_hdr_t));
295         *h = *h0;
296         h->ns = 0; h->sns = 0;
297         h->name = malloc(h->l_nm); memcpy(h->name, h0->name, h->l_nm);
298         h->txt = calloc(1, h->l_txt + 1); memcpy(h->txt, h0->txt, h->l_txt);
299         h->l_smpl = s.l; h->sname = s.s;
300         bcf_hdr_sync(h);
301         return h;
302 }
303
304 int bcf_subsam(int n_smpl, int *list, bcf1_t *b)
305 {
306         int i, j;
307         for (j = 0; j < b->n_gi; ++j) {
308                 bcf_ginfo_t *gi = b->gi + j;
309                 uint8_t *swap;
310                 swap = malloc(gi->len * b->n_smpl);
311                 for (i = 0; i < n_smpl; ++i)
312                         memcpy(swap + i * gi->len, (uint8_t*)gi->data + list[i] * gi->len, gi->len);
313                 free(gi->data);
314                 gi->data = swap;
315         }
316         b->n_smpl = n_smpl;
317         return 0;
318 }
319
320 static int8_t nt4_table[128] = {
321         4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4, 
322         4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4, 
323         4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4 /*'-'*/, 4, 4,
324         4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4,  4, 4, 4, 4, 
325         4, 0, 4, 1,  4, 4, 4, 2,  4, 4, 4, 4,  4, 4, 4, 4, 
326         4, 4, 4, 4,  3, 4, 4, 4, -1, 4, 4, 4,  4, 4, 4, 4, 
327         4, 0, 4, 1,  4, 4, 4, 2,  4, 4, 4, 4,  4, 4, 4, 4, 
328         4, 4, 4, 4,  3, 4, 4, 4, -1, 4, 4, 4,  4, 4, 4, 4
329 };
330
331 int bcf_gl10(const bcf1_t *b, uint8_t *gl)
332 {
333         int a[4], k, l, map[4], k1, j, i;
334         const bcf_ginfo_t *PL;
335         char *s;
336         if (b->ref[1] != 0 || b->n_alleles > 4) return -1; // ref is not a single base or >4 alleles
337         for (i = 0; i < b->n_gi; ++i)
338                 if (b->gi[i].fmt == bcf_str2int("PL", 2)) break;
339         if (i == b->n_gi) return -1; // no PL
340         PL = b->gi + i;
341         a[0] = nt4_table[(int)b->ref[0]];
342         if (a[0] > 3 || a[0] < 0) return -1; // ref is not A/C/G/T
343         a[1] = a[2] = a[3] = -2; // -1 has a special meaning
344         if (b->alt[0] == 0) return -1; // no alternate allele
345         map[0] = map[1] = map[2] = map[3] = -2;
346         map[a[0]] = 0;
347         for (k = 0, s = b->alt, k1 = -1; k < 3 && *s; ++k, s += 2) {
348                 if (s[1] != ',' && s[1] != 0) return -1; // ALT is not single base
349                 a[k+1] = nt4_table[(int)*s];
350                 if (a[k+1] >= 0) map[a[k+1]] = k+1;
351                 else k1 = k + 1;
352                 if (s[1] == 0) break; // the end of the ALT string
353         }
354         for (k = 0; k < 4; ++k)
355                 if (map[k] < 0) map[k] = k1;
356         for (i = 0; i < b->n_smpl; ++i) {
357                 const uint8_t *p = PL->data + i * PL->len; // the PL for the i-th individual
358                 uint8_t *g = gl + 10 * i;
359                 for (k = j = 0; k < 4; ++k) {
360                         for (l = k; l < 4; ++l) {
361                                 int t, x = map[k], y = map[l];
362                                 if (x > y) t = x, x = y, y = t; // make sure x is the smaller
363                                 g[j++] = p[y * (y+1) / 2 + x];
364                         }
365                 }
366         }
367         return 0;
368 }
369
370 int bcf_gl10_indel(const bcf1_t *b, uint8_t *gl)
371 {
372         int k, l, j, i;
373         const bcf_ginfo_t *PL;
374         if (b->alt[0] == 0) return -1; // no alternate allele
375         for (i = 0; i < b->n_gi; ++i)
376                 if (b->gi[i].fmt == bcf_str2int("PL", 2)) break;
377         if (i == b->n_gi) return -1; // no PL
378         PL = b->gi + i;
379         for (i = 0; i < b->n_smpl; ++i) {
380                 const uint8_t *p = PL->data + i * PL->len; // the PL for the i-th individual
381                 uint8_t *g = gl + 10 * i;
382                 for (k = j = 0; k < 4; ++k) {
383                         for (l = k; l < 4; ++l) {
384                                 int t, x = k, y = l;
385                                 if (x > y) t = x, x = y, y = t; // make sure x is the smaller
386                                 x = y * (y+1) / 2 + x;
387                                 g[j++] = x < PL->len? p[x] : 255;
388                         }
389                 }
390         }
391         return 0;
392 }