samtools/bcftools/em.c.pysam.c

   1 #include "pysam.h"
   2
   3 #include <stdlib.h>
   4 #include <string.h>
   5 #include <math.h>
   6 #include "bcf.h"
   7 #include "kmin.h"
   8
   9 static double g_q2p[256];
  10
  11 #define ITER_MAX 50
  12 #define ITER_TRY 10
  13 #define EPS 1e-5
  14
  15 extern double kf_gammaq(double, double);
  16
  17 /*
  18         Generic routines
  19  */
  20 // get the 3 genotype likelihoods
  21 static double *get_pdg3(const bcf1_t *b)
  22 {
  23         double *pdg;
  24         const uint8_t *PL = 0;
  25         int i, PL_len = 0;
  26         // initialize g_q2p if necessary
  27         if (g_q2p[0] == 0.)
  28                 for (i = 0; i < 256; ++i)
  29                         g_q2p[i] = pow(10., -i / 10.);
  30         // set PL and PL_len
  31         for (i = 0; i < b->n_gi; ++i) {
  32                 if (b->gi[i].fmt == bcf_str2int("PL", 2)) {
  33                         PL = (const uint8_t*)b->gi[i].data;
  34                         PL_len = b->gi[i].len;
  35                         break;
  36                 }
  37         }
  38         if (i == b->n_gi) return 0; // no PL
  39         // fill pdg
  40         pdg = malloc(3 * b->n_smpl * sizeof(double));
  41         for (i = 0; i < b->n_smpl; ++i) {
  42                 const uint8_t *pi = PL + i * PL_len;
  43                 double *p = pdg + i * 3;
  44                 p[0] = g_q2p[pi[2]]; p[1] = g_q2p[pi[1]]; p[2] = g_q2p[pi[0]];
  45         }
  46         return pdg;
  47 }
  48
  49 // estimate site allele frequency in a very naive and inaccurate way
  50 static double est_freq(int n, const double *pdg)
  51 {
  52         int i, gcnt[3], tmp1;
  53         // get a rough estimate of the genotype frequency
  54         gcnt[0] = gcnt[1] = gcnt[2] = 0;
  55         for (i = 0; i < n; ++i) {
  56                 const double *p = pdg + i * 3;
  57                 if (p[0] != 1. || p[1] != 1. || p[2] != 1.) {
  58                         int which = p[0] > p[1]? 0 : 1;
  59                         which = p[which] > p[2]? which : 2;
  60                         ++gcnt[which];
  61                 }
  62         }
  63         tmp1 = gcnt[0] + gcnt[1] + gcnt[2];
  64         return (tmp1 == 0)? -1.0 : (.5 * gcnt[1] + gcnt[2]) / tmp1;
  65 }
  66
  67 /*
  68         Single-locus EM
  69  */
  70
  71 typedef struct {
  72         int beg, end;
  73         const double *pdg;
  74 } minaux1_t;
  75
  76 static double prob1(double f, void *data)
  77 {
  78         minaux1_t *a = (minaux1_t*)data;
  79         double p = 1., l = 0., f3[3];
  80         int i;
  81 //      printf("brent %lg\n", f);
  82         if (f < 0 || f > 1) return 1e300;
  83         f3[0] = (1.-f)*(1.-f); f3[1] = 2.*f*(1.-f); f3[2] = f*f;
  84         for (i = a->beg; i < a->end; ++i) {
  85                 const double *pdg = a->pdg + i * 3;
  86                 p *= pdg[0] * f3[0] + pdg[1] * f3[1] + pdg[2] * f3[2];
  87                 if (p < 1e-200) l -= log(p), p = 1.;
  88         }
  89         return l - log(p);
  90 }
  91
  92 // one EM iteration for allele frequency estimate
  93 static double freq_iter(double *f, const double *_pdg, int beg, int end)
  94 {
  95         double f0 = *f, f3[3], err;
  96         int i;
  97 //      printf("em %lg\n", *f);
  98         f3[0] = (1.-f0)*(1.-f0); f3[1] = 2.*f0*(1.-f0); f3[2] = f0*f0;
  99         for (i = beg, f0 = 0.; i < end; ++i) {
 100                 const double *pdg = _pdg + i * 3;
 101                 f0 += (pdg[1] * f3[1] + 2. * pdg[2] * f3[2])
 102                         / (pdg[0] * f3[0] + pdg[1] * f3[1] + pdg[2] * f3[2]);
 103         }
 104         f0 /= (end - beg) * 2;
 105         err = fabs(f0 - *f);
 106         *f = f0;
 107         return err;
 108 }
 109
 110 /* The following function combines EM and Brent's method. When the signal from
 111  * the data is strong, EM is faster but sometimes, EM may converge very slowly.
 112  * When this happens, we switch to Brent's method. The idea is learned from
 113  * Rasmus Nielsen.
 114  */
 115 static double freqml(double f0, int beg, int end, const double *pdg)
 116 {
 117         int i;
 118         double f;
 119         for (i = 0, f = f0; i < ITER_TRY; ++i)
 120                 if (freq_iter(&f, pdg, beg, end) < EPS) break;
 121         if (i == ITER_TRY) { // haven't converged yet; try Brent's method
 122                 minaux1_t a;
 123                 a.beg = beg; a.end = end; a.pdg = pdg;
 124                 kmin_brent(prob1, f0 == f? .5*f0 : f0, f, (void*)&a, EPS, &f);
 125         }
 126         return f;
 127 }
 128
 129 // one EM iteration for genotype frequency estimate
 130 static double g3_iter(double g[3], const double *_pdg, int beg, int end)
 131 {
 132         double err, gg[3];
 133         int i;
 134         gg[0] = gg[1] = gg[2] = 0.;
 135 //      printf("%lg,%lg,%lg\n", g[0], g[1], g[2]);
 136         for (i = beg; i < end; ++i) {
 137                 double sum, tmp[3];
 138                 const double *pdg = _pdg + i * 3;
 139                 tmp[0] = pdg[0] * g[0]; tmp[1] = pdg[1] * g[1]; tmp[2] = pdg[2] * g[2];
 140                 sum = (tmp[0] + tmp[1] + tmp[2]) * (end - beg);
 141                 gg[0] += tmp[0] / sum; gg[1] += tmp[1] / sum; gg[2] += tmp[2] / sum;
 142         }
 143         err = fabs(gg[0] - g[0]) > fabs(gg[1] - g[1])? fabs(gg[0] - g[0]) : fabs(gg[1] - g[1]);
 144         err = err > fabs(gg[2] - g[2])? err : fabs(gg[2] - g[2]);
 145         g[0] = gg[0]; g[1] = gg[1]; g[2] = gg[2];
 146         return err;
 147 }
 148
 149 // perform likelihood ratio test
 150 static double lk_ratio_test(int n, int n1, const double *pdg, double f3[3][3])
 151 {
 152         double r;
 153         int i;
 154         for (i = 0, r = 1.; i < n1; ++i) {
 155                 const double *p = pdg + i * 3;
 156                 r *= (p[0] * f3[1][0] + p[1] * f3[1][1] + p[2] * f3[1][2])
 157                         / (p[0] * f3[0][0] + p[1] * f3[0][1] + p[2] * f3[0][2]);
 158         }
 159         for (; i < n; ++i) {
 160                 const double *p = pdg + i * 3;
 161                 r *= (p[0] * f3[2][0] + p[1] * f3[2][1] + p[2] * f3[2][2])
 162                         / (p[0] * f3[0][0] + p[1] * f3[0][1] + p[2] * f3[0][2]);
 163         }
 164         return r;
 165 }
 166
 167 // x[0]: ref frequency
 168 // x[1..3]: alt-alt, alt-ref, ref-ref frequenc
 169 // x[4]: HWE P-value
 170 // x[5..6]: group1 freq, group2 freq
 171 // x[7]: 1-degree P-value
 172 // x[8]: 2-degree P-value
 173 int bcf_em1(const bcf1_t *b, int n1, int flag, double x[9])
 174 {
 175         double *pdg;
 176         int i, n, n2;
 177         if (b->n_alleles < 2) return -1; // one allele only
 178         // initialization
 179         if (n1 < 0 || n1 > b->n_smpl) n1 = 0;
 180         if (flag & 1<<7) flag |= 7<<5; // compute group freq if LRT is required
 181         if (flag & 0xf<<1) flag |= 0xf<<1;
 182         n = b->n_smpl; n2 = n - n1;
 183         pdg = get_pdg3(b);
 184         if (pdg == 0) return -1;
 185         for (i = 0; i < 9; ++i) x[i] = -1.;
 186         {
 187                 if ((x[0] = est_freq(n, pdg)) < 0.) {
 188                         free(pdg);
 189                         return -1; // no data
 190                 }
 191                 x[0] = freqml(x[0], 0, n, pdg);
 192         }
 193         if (flag & (0xf<<1|1<<8)) { // estimate the genotype frequency and test HWE
 194                 double *g = x + 1, f3[3], r;
 195                 f3[0] = g[0] = (1 - x[0]) * (1 - x[0]);
 196                 f3[1] = g[1] = 2 * x[0] * (1 - x[0]);
 197                 f3[2] = g[2] = x[0] * x[0];
 198                 for (i = 0; i < ITER_MAX; ++i)
 199                         if (g3_iter(g, pdg, 0, n) < EPS) break;
 200                 // Hardy-Weinberg equilibrium (HWE)
 201                 for (i = 0, r = 1.; i < n; ++i) {
 202                         double *p = pdg + i * 3;
 203                         r *= (p[0] * g[0] + p[1] * g[1] + p[2] * g[2]) / (p[0] * f3[0] + p[1] * f3[1] + p[2] * f3[2]);
 204                 }
 205                 x[4] = kf_gammaq(.5, log(r));
 206         }
 207         if ((flag & 7<<5) && n1 > 0 && n1 < n) { // group frequency
 208                 x[5] = freqml(x[0], 0, n1, pdg);
 209                 x[6] = freqml(x[0], n1, n, pdg);
 210         }
 211         if ((flag & 1<<7) && n1 > 0 && n1 < n) { // 1-degree P-value
 212                 double f[3], f3[3][3];
 213                 f[0] = x[0]; f[1] = x[5]; f[2] = x[6];
 214                 for (i = 0; i < 3; ++i)
 215                         f3[i][0] = (1-f[i])*(1-f[i]), f3[i][1] = 2*f[i]*(1-f[i]), f3[i][2] = f[i]*f[i];
 216                 x[7] = kf_gammaq(.5, log(lk_ratio_test(n, n1, pdg, f3)));
 217         }
 218         if ((flag & 1<<8) && n1 > 0 && n1 < n) { // 2-degree P-value
 219                 double g[3][3];
 220                 for (i = 0; i < 3; ++i) memcpy(g[i], x + 1, 3 * sizeof(double));
 221                 for (i = 0; i < ITER_MAX; ++i)
 222                         if (g3_iter(g[1], pdg, 0, n1) < EPS) break;
 223                 for (i = 0; i < ITER_MAX; ++i)
 224                         if (g3_iter(g[2], pdg, n1, n) < EPS) break;
 225                 x[8] = kf_gammaq(1., log(lk_ratio_test(n, n1, pdg, g)));
 226         }
 227         // free
 228         free(pdg);
 229         return 0;
 230 }
 231
 232 /*
 233         Two-locus EM (LD)
 234  */
 235
 236 #define _G1(h, k) ((h>>1&1) + (k>>1&1))
 237 #define _G2(h, k) ((h&1) + (k&1))
 238
 239 // 0: the previous site; 1: the current site
 240 static int pair_freq_iter(int n, double *pdg[2], double f[4])
 241 {
 242         double ff[4];
 243         int i, k, h;
 244 //      printf("%lf,%lf,%lf,%lf\n", f[0], f[1], f[2], f[3]);
 245         memset(ff, 0, 4 * sizeof(double));
 246         for (i = 0; i < n; ++i) {
 247                 double *p[2], sum, tmp;
 248                 p[0] = pdg[0] + i * 3; p[1] = pdg[1] + i * 3;
 249                 for (k = 0, sum = 0.; k < 4; ++k)
 250                         for (h = 0; h < 4; ++h)
 251                                 sum += f[k] * f[h] * p[0][_G1(k,h)] * p[1][_G2(k,h)];
 252                 for (k = 0; k < 4; ++k) {
 253                         tmp = f[0] * (p[0][_G1(0,k)] * p[1][_G2(0,k)] + p[0][_G1(k,0)] * p[1][_G2(k,0)])
 254                                 + f[1] * (p[0][_G1(1,k)] * p[1][_G2(1,k)] + p[0][_G1(k,1)] * p[1][_G2(k,1)])
 255                                 + f[2] * (p[0][_G1(2,k)] * p[1][_G2(2,k)] + p[0][_G1(k,2)] * p[1][_G2(k,2)])
 256                                 + f[3] * (p[0][_G1(3,k)] * p[1][_G2(3,k)] + p[0][_G1(k,3)] * p[1][_G2(k,3)]);
 257                         ff[k] += f[k] * tmp / sum;
 258                 }
 259         }
 260         for (k = 0; k < 4; ++k) f[k] = ff[k] / (2 * n);
 261         return 0;
 262 }
 263
 264 double bcf_pair_freq(const bcf1_t *b0, const bcf1_t *b1, double f[4])
 265 {
 266         const bcf1_t *b[2];
 267         int i, j, n_smpl;
 268         double *pdg[2], flast[4], r, f0[2];
 269         // initialize others
 270         if (b0->n_smpl != b1->n_smpl) return -1; // different number of samples
 271         n_smpl = b0->n_smpl;
 272         b[0] = b0; b[1] = b1;
 273         f[0] = f[1] = f[2] = f[3] = -1.;
 274         if (b[0]->n_alleles < 2 || b[1]->n_alleles < 2) return -1; // one allele only
 275         pdg[0] = get_pdg3(b0); pdg[1] = get_pdg3(b1);
 276         if (pdg[0] == 0 || pdg[1] == 0) {
 277                 free(pdg[0]); free(pdg[1]);
 278                 return -1;
 279         }
 280         // set the initial value
 281         f0[0] = est_freq(n_smpl, pdg[0]);
 282         f0[1] = est_freq(n_smpl, pdg[1]);
 283         f[0] = (1 - f0[0]) * (1 - f0[1]); f[3] = f0[0] * f0[1];
 284         f[1] = (1 - f0[0]) * f0[1]; f[2] = f0[0] * (1 - f0[1]);
 285         // iteration
 286         for (j = 0; j < ITER_MAX; ++j) {
 287                 double eps = 0;
 288                 memcpy(flast, f, 4 * sizeof(double));
 289                 pair_freq_iter(n_smpl, pdg, f);
 290                 for (i = 0; i < 4; ++i) {
 291                         double x = fabs(f[i] - flast[i]);
 292                         if (x > eps) eps = x;
 293                 }
 294                 if (eps < EPS) break;
 295         }
 296         // free
 297         free(pdg[0]); free(pdg[1]);
 298         { // calculate r^2
 299                 double p[2], q[2], D;
 300                 p[0] = f[0] + f[1]; q[0] = 1 - p[0];
 301                 p[1] = f[0] + f[2]; q[1] = 1 - p[1];
 302                 D = f[0] * f[3] - f[1] * f[2];
 303                 r = sqrt(D * D / (p[0] * p[1] * q[0] * q[1]));
 304 //              printf("R(%lf,%lf,%lf,%lf)=%lf\n", f[0], f[1], f[2], f[3], r);
 305                 if (isnan(r)) r = -1.;
 306         }
 307         return r;
 308 }