sam/misc/vcfutils.lua

   1 #!/usr/bin/env luajit
   2
   3 -----------------------------------
   4 -- BEGIN: routines from klib.lua --
   5 -----------------------------------
   6
   7 -- Description: getopt() translated from the BSD getopt(); compatible with the default Unix getopt()
   8 --[[ Example:
   9         for o, a in os.getopt(arg, 'a:b') do
  10                 print(o, a)
  11         end
  12 ]]--
  13 function os.getopt(args, ostr)
  14         local arg, place = nil, 0;
  15         return function ()
  16                 if place == 0 then -- update scanning pointer
  17                         place = 1
  18                         if #args == 0 or args[1]:sub(1, 1) ~= '-' then place = 0; return nil end
  19                         if #args[1] >= 2 then
  20                                 place = place + 1
  21                                 if args[1]:sub(2, 2) == '-' then -- found "--"
  22                                         table.remove(args, 1);
  23                                         place = 0
  24                                         return nil;
  25                                 end
  26                         end
  27                 end
  28                 local optopt = place <= #args[1] and args[1]:sub(place, place) or nil
  29                 place = place + 1;
  30                 local oli = optopt and ostr:find(optopt) or nil
  31                 if optopt == ':' or oli == nil then -- unknown option
  32                         if optopt == '-' then return nil end
  33                         if place > #args[1] then
  34                                 table.remove(args, 1);
  35                                 place = 0;
  36                         end
  37                         return '?';
  38                 end
  39                 oli = oli + 1;
  40                 if ostr:sub(oli, oli) ~= ':' then -- do not need argument
  41                         arg = nil;
  42                         if place > #args[1] then
  43                                 table.remove(args, 1);
  44                                 place = 0;
  45                         end
  46                 else -- need an argument
  47                         if place <= #args[1] then  -- no white space
  48                                 arg = args[1]:sub(place);
  49                         else
  50                                 table.remove(args, 1);
  51                                 if #args == 0 then -- an option requiring argument is the last one
  52                                         place = 0;
  53                                         if ostr:sub(1, 1) == ':' then return ':' end
  54                                         return '?';
  55                                 else arg = args[1] end
  56                         end
  57                         table.remove(args, 1);
  58                         place = 0;
  59                 end
  60                 return optopt, arg;
  61         end
  62 end
  63
  64 -- Description: string split
  65 function string:split(sep, n)
  66         local a, start = {}, 1;
  67         sep = sep or "%s+";
  68         repeat
  69                 local b, e = self:find(sep, start);
  70                 if b == nil then
  71                         table.insert(a, self:sub(start));
  72                         break
  73                 end
  74                 a[#a+1] = self:sub(start, b - 1);
  75                 start = e + 1;
  76                 if n and #a == n then
  77                         table.insert(a, self:sub(start));
  78                         break
  79                 end
  80         until start > #self;
  81         return a;
  82 end
  83
  84 -- Description: smart file open
  85 function io.xopen(fn, mode)
  86         mode = mode or 'r';
  87         if fn == nil then return io.stdin;
  88         elseif fn == '-' then return (mode == 'r' and io.stdin) or io.stdout;
  89         elseif fn:sub(-3) == '.gz' then return (mode == 'r' and io.popen('gzip -dc ' .. fn, 'r')) or io.popen('gzip > ' .. fn, 'w');
  90         elseif fn:sub(-4) == '.bz2' then return (mode == 'r' and io.popen('bzip2 -dc ' .. fn, 'r')) or io.popen('bgzip2 > ' .. fn, 'w');
  91         else return io.open(fn, mode) end
  92 end
  93
  94 -- Description: log gamma function
  95 -- Required by: math.lbinom()
  96 -- Reference: AS245, 2nd algorithm, http://lib.stat.cmu.edu/apstat/245
  97 function math.lgamma(z)
  98         local x;
  99         x = 0.1659470187408462e-06     / (z+7);
 100         x = x + 0.9934937113930748e-05 / (z+6);
 101         x = x - 0.1385710331296526     / (z+5);
 102         x = x + 12.50734324009056      / (z+4);
 103         x = x - 176.6150291498386      / (z+3);
 104         x = x + 771.3234287757674      / (z+2);
 105         x = x - 1259.139216722289      / (z+1);
 106         x = x + 676.5203681218835      / z;
 107         x = x + 0.9999999999995183;
 108         return math.log(x) - 5.58106146679532777 - z + (z-0.5) * math.log(z+6.5);
 109 end
 110
 111 -- Description: regularized incomplete gamma function
 112 -- Dependent on: math.lgamma()
 113 --[[
 114   Formulas are taken from Wiki, with additional input from Numerical
 115   Recipes in C (for modified Lentz's algorithm) and AS245
 116   (http://lib.stat.cmu.edu/apstat/245).
 117
 118   A good online calculator is available at:
 119
 120     http://www.danielsoper.com/statcalc/calc23.aspx
 121
 122   It calculates upper incomplete gamma function, which equals
 123   math.igamma(s,z,true)*math.exp(math.lgamma(s))
 124 ]]--
 125 function math.igamma(s, z, complement)
 126
 127         local function _kf_gammap(s, z)
 128                 local sum, x = 1, 1;
 129                 for k = 1, 100 do
 130                         x = x * z / (s + k);
 131                         sum = sum + x;
 132                         if x / sum < 1e-14 then break end
 133                 end
 134                 return math.exp(s * math.log(z) - z - math.lgamma(s + 1.) + math.log(sum));
 135         end
 136
 137         local function _kf_gammaq(s, z)
 138                 local C, D, f, TINY;
 139                 f = 1. + z - s; C = f; D = 0.; TINY = 1e-290;
 140                 -- Modified Lentz's algorithm for computing continued fraction. See Numerical Recipes in C, 2nd edition, section 5.2
 141                 for j = 1, 100 do
 142                         local d;
 143                         local a, b = j * (s - j), j*2 + 1 + z - s;
 144                         D = b + a * D;
 145                         if D < TINY then D = TINY end
 146                         C = b + a / C;
 147                         if C < TINY then C = TINY end
 148                         D = 1. / D;
 149                         d = C * D;
 150                         f = f * d;
 151                         if math.abs(d - 1) < 1e-14 then break end
 152                 end
 153                 return math.exp(s * math.log(z) - z - math.lgamma(s) - math.log(f));
 154         end
 155
 156         if complement then
 157                 return ((z <= 1 or z < s) and 1 - _kf_gammap(s, z)) or _kf_gammaq(s, z);
 158         else
 159                 return ((z <= 1 or z < s) and _kf_gammap(s, z)) or (1 - _kf_gammaq(s, z));
 160         end
 161 end
 162
 163 function math.brent(func, a, b, tol)
 164         local gold1, gold2, tiny, max_iter = 1.6180339887, 0.3819660113, 1e-20, 100
 165
 166         local fa, fb = func(a, data), func(b, data)
 167         if fb > fa then -- swap, such that f(a) > f(b)
 168                 a, b, fa, fb = b, a, fb, fa
 169         end
 170         local c = b + gold1 * (b - a)
 171         local fc = func(c) -- golden section extrapolation
 172         while fb > fc do
 173                 local bound = b + 100.0 * (c - b) -- the farthest point where we want to go
 174                 local r = (b - a) * (fb - fc)
 175                 local q = (b - c) * (fb - fa)
 176                 if math.abs(q - r) < tiny then -- avoid 0 denominator
 177                         tmp = q > r and tiny or 0.0 - tiny
 178                 else tmp = q - r end
 179                 u = b - ((b - c) * q - (b - a) * r) / (2.0 * tmp) -- u is the parabolic extrapolation point
 180                 if (b > u and u > c) or (b < u and u < c) then -- u lies between b and c
 181                         fu = func(u)
 182                         if fu < fc then -- (b,u,c) bracket the minimum
 183                                 a, b, fa, fb = b, u, fb, fu
 184                                 break
 185                         elseif fu > fb then -- (a,b,u) bracket the minimum
 186                                 c, fc = u, fu
 187                                 break
 188                         end
 189                         u = c + gold1 * (c - b)
 190                         fu = func(u) -- golden section extrapolation
 191                 elseif (c > u and u > bound) or (c < u and u < bound) then -- u lies between c and bound
 192                         fu = func(u)
 193                         if fu < fc then -- fb > fc > fu
 194                                 b, c, u = c, u, c + gold1 * (c - b)
 195                                 fb, fc, fu = fc, fu, func(u)
 196                         else -- (b,c,u) bracket the minimum
 197                                 a, b, c = b, c, u
 198                                 fa, fb, fc = fb, fc, fu
 199                                 break
 200                         end
 201                 elseif (u > bound and bound > c) or (u < bound and bound < c) then -- u goes beyond the bound
 202                         u = bound
 203                         fu = func(u)
 204                 else -- u goes the other way around, use golden section extrapolation
 205                         u = c + gold1 * (c - b)
 206                         fu = func(u)
 207                 end
 208                 a, b, c = b, c, u
 209                 fa, fb, fc = fb, fc, fu
 210         end
 211         if a > c then a, c = c, a end -- swap
 212
 213         -- now, a<b<c, fa>fb and fb<fc, move on to Brent's algorithm
 214         local e, d = 0, 0
 215         local w, v, fw, fv
 216         w, v = b, b
 217         fw, fv = fb, fb
 218         for iter = 1, max_iter do
 219                 local mid = 0.5 * (a + c)
 220                 local tol1 = tol * math.abs(b) + tiny
 221                 local tol2 = 2.0 * tol1
 222                 if math.abs(b - mid) <= tol2 - 0.5 * (c - a) then return fb, b end -- found
 223                 if math.abs(e) > tol1 then
 224                         -- related to parabolic interpolation
 225                         local r = (b - w) * (fb - fv)
 226                         local q = (b - v) * (fb - fw)
 227                         local p = (b - v) * q - (b - w) * r
 228                         q = 2.0 * (q - r)
 229                         if q > 0.0 then p = 0.0 - p
 230                         else q = 0.0 - q end
 231                         eold, e = e, d
 232                         if math.abs(p) >= math.abs(0.5 * q * eold) or p <= q * (a - b) or p >= q * (c - b) then
 233                                 e = b >= mid and a - b or c - b
 234                                 d = gold2 * e
 235                         else
 236                                 d, u = p / q, b + d -- actual parabolic interpolation happens here
 237                                 if u - a < tol2 or c - u < tol2 then
 238                                         d = mid > b and tol1 or 0.0 - tol1
 239                                 end
 240                         end
 241                 else -- golden section interpolation
 242                         e = b >= min and a - b or c - b
 243                         d = gold2 * e
 244                 end
 245                 u = fabs(d) >= tol1 and b + d or b + (d > 0.0 and tol1 or -tol1);
 246                 fu = func(u)
 247                 if fu <= fb then -- u is the minimum point so far
 248                         if u >= b then a = b
 249                         else c = b end
 250                         v, w, b = w, b, u
 251                         fv, fw, fb = fw, fb, fu
 252                 else -- adjust (a,c) and (u,v,w)
 253                         if u < b then a = u
 254                         else c = u end
 255                         if fu <= fw or w == b then
 256                                 v, w = w, u
 257                                 fv, fw = fw, fu
 258                         elseif fu <= fv or v == b or v == w then
 259                                 v, fv = u, fu;
 260                         end
 261                 end
 262         end
 263         return fb, b
 264 end
 265
 266 matrix = {}
 267
 268 -- Description: chi^2 test for contingency tables
 269 -- Dependent on: math.igamma()
 270 function matrix.chi2(a)
 271         if #a == 2 and #a[1] == 2 then -- 2x2 table
 272                 local x, z
 273                 x = (a[1][1] + a[1][2]) * (a[2][1] + a[2][2]) * (a[1][1] + a[2][1]) * (a[1][2] + a[2][2])
 274                 if x == 0 then return 0, 1, false end
 275                 z = a[1][1] * a[2][2] - a[1][2] * a[2][1]
 276                 z = (a[1][1] + a[1][2] + a[2][1] + a[2][2]) * z * z / x
 277                 return z, math.igamma(.5, .5 * z, true), true
 278         else -- generic table
 279                 local rs, cs, n, m, N, z = {}, {}, #a, #a[1], 0, 0
 280                 for i = 1, n do rs[i] = 0 end
 281                 for j = 1, m do cs[j] = 0 end
 282                 for i = 1, n do -- compute column sum and row sum
 283                         for j = 1, m do cs[j], rs[i] = cs[j] + a[i][j], rs[i] + a[i][j] end
 284                 end
 285                 for i = 1, n do N = N + rs[i] end
 286                 for i = 1, n do -- compute the chi^2 statistics
 287                         for j = 1, m do
 288                                 local E = rs[i] * cs[j] / N;
 289                                 z = z + (a[i][j] - E) * (a[i][j] - E) / E
 290                         end
 291                 end
 292                 return z, math.igamma(.5 * (n-1) * (m-1), .5 * z, true), true;
 293         end
 294 end
 295
 296 ---------------------------------
 297 -- END: routines from klib.lua --
 298 ---------------------------------
 299
 300
 301 --------------------------
 302 -- BEGIN: misc routines --
 303 --------------------------
 304
 305 -- precompute an array for PL->probability conversion
 306 -- @param m maximum PL
 307 function algo_init_q2p(m)
 308         local q2p = {}
 309         for i = 0, m do
 310                 q2p[i] = math.pow(10, -i / 10)
 311         end
 312         return q2p
 313 end
 314
 315 -- given the haplotype frequency, compute r^2
 316 -- @param f 4 haplotype frequencies; f[] is 0-indexed.
 317 -- @return r^2
 318 function algo_r2(f)
 319         local p = { f[0] + f[1], f[0] + f[2] }
 320         local D = f[0] * f[3] - f[1] * f[2]
 321         return (p[1] == 0 or p[2] == 0 or 1-p[1] == 0 or 1-p[2] == 0) and 0 or D * D  / (p[1] * p[2] * (1 - p[1]) * (1 - p[2]))
 322 end
 323
 324 -- parse a VCF line to get PL
 325 -- @param q2p is computed by algo_init_q2p()
 326 function text_parse_pl(t, q2p, parse_GT)
 327         parse_GT = parse_GT == nil and true or false
 328         local ht, gt, pl = {}, {}, {}
 329         local s, j0 = t[9]:split(':'), 0
 330         for j = 1, #s do
 331                 if s[j] == 'PL' then j0 = j break end
 332         end
 333         local has_GT = (s[1] == 'GT' and parse_GT) and true or false
 334         for i = 10, #t do
 335                 if j0 > 0 then
 336                         local s = t[i]:split(':')
 337                         local a, b = 1, s[j0]:find(',')
 338                         pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, b - 1))]
 339                         a, b = b + 1, s[j0]:find(',', b + 1)
 340                         pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, b - 1))]
 341                         a, b = b + 1, s[j0]:find(',', b + 1)
 342                         pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, (b and b - 1) or nil))]
 343                 end
 344                 if has_GT then
 345                         if t[i]:sub(1, 1) ~= '.' then
 346                                 local g = tonumber(t[i]:sub(1, 1)) + tonumber(t[i]:sub(3, 3));
 347                                 gt[#gt+1] = 1e-6; gt[#gt+1] = 1e-6; gt[#gt+1] = 1e-6
 348                                 gt[#gt - 2 + g] = 1
 349                                 ht[#ht+1] = tonumber(t[i]:sub(1, 1)); ht[#ht+1] = tonumber(t[i]:sub(3, 3));
 350                         else
 351                                 gt[#gt+1] = 1; gt[#gt+1] = 1; gt[#gt+1] = 1
 352                                 ht[#ht+1] = -1; ht[#ht+1] = -1;
 353                         end
 354                 end
 355 --              print(t[i], pl[#pl-2], pl[#pl-1], pl[#pl], gt[#gt-2], gt[#gt-1], gt[#gt])
 356         end
 357         if #pl == 0 then pl = nil end
 358         local x = has_GT and { t[1], t[2], ht, gt, pl } or { t[1], t[2], nil, nil, pl }
 359         return x
 360 end
 361
 362 -- Infer haplotype frequency
 363 -- @param pdg  genotype likelihoods P(D|g) generated by text_parse_pl(). pdg[] is 1-indexed.
 364 -- @param eps  precision [1e-5]
 365 -- @return 2-locus haplotype frequencies, 0-indexed array
 366 function algo_hapfreq2(pdg, eps)
 367         eps = eps or 1e-5
 368         local n, f = #pdg[1] / 3, {[0]=0.25, 0.25, 0.25, 0.25}
 369         for iter = 1, 100 do
 370                 local F = {[0]=0, 0, 0, 0}
 371                 for i = 0, n - 1 do
 372                         local p1, p2 = {[0]=pdg[1][i*3+1], pdg[1][i*3+2], pdg[1][i*3+3]}, {[0]=pdg[2][i*3+1], pdg[2][i*3+2], pdg[2][i*3+3]}
 373                         local u = { [0]=
 374                                 f[0] * (f[0] * p1[0] * p2[0] + f[1] * p1[0] * p2[1] + f[2] * p1[1] * p2[0] + f[3] * p1[1] * p2[1]),
 375                                 f[1] * (f[0] * p1[0] * p2[1] + f[1] * p1[0] * p2[2] + f[2] * p1[1] * p2[1] + f[3] * p1[1] * p2[2]),
 376                                 f[2] * (f[0] * p1[1] * p2[0] + f[1] * p1[1] * p2[1] + f[2] * p1[2] * p2[0] + f[3] * p1[2] * p2[1]),
 377                                 f[3] * (f[0] * p1[1] * p2[1] + f[1] * p1[1] * p2[2] + f[2] * p1[2] * p2[1] + f[3] * p1[2] * p2[2])
 378                         }
 379                         local s = u[0] + u[1] + u[2] + u[3]
 380                         s = 1 / (s * n)
 381                         F[0] = F[0] + u[0] * s
 382                         F[1] = F[1] + u[1] * s
 383                         F[2] = F[2] + u[2] * s
 384                         F[3] = F[3] + u[3] * s
 385                 end
 386                 local e = 0
 387                 for k = 0, 3 do
 388                         e = math.abs(f[k] - F[k]) > e and math.abs(f[k] - F[k]) or e
 389                 end
 390                 for k = 0, 3 do f[k] = F[k] end
 391                 if e < eps then break end
 392 --              print(f[0], f[1], f[2], f[3])
 393         end
 394         return f
 395 end
 396
 397 ------------------------
 398 -- END: misc routines --
 399 ------------------------
 400
 401
 402 ---------------------
 403 -- BEGIN: commands --
 404 ---------------------
 405
 406 -- CMD vcf2bgl: convert PL tagged VCF to Beagle input --
 407 function cmd_vcf2bgl()
 408         if #arg == 0 then
 409                 print("\nUsage: vcf2bgl.lua <in.vcf>")
 410                 print("\nNB: This command finds PL by matching /(\\d+),(\\d+),(\\d+)/.\n");
 411                 os.exit(1)
 412         end
 413
 414         local lookup = {}
 415         for i = 0, 10000 do lookup[i] = string.format("%.4f", math.pow(10, -i/10)) end
 416
 417         local fp = io.xopen(arg[1])
 418         for l in fp:lines() do
 419                 if l:sub(1, 2) == '##' then -- meta lines; do nothing
 420                 elseif l:sub(1, 1) == '#' then -- sample lines
 421                         local t, s = l:split('\t'), {}
 422                         for i = 10, #t do s[#s+1] = t[i]; s[#s+1] = t[i]; s[#s+1] = t[i] end
 423                         print('marker', 'alleleA', 'alleleB', table.concat(s, '\t'))
 424                 else -- data line
 425                         local t = l:split('\t');
 426                         if t[5] ~= '.' and t[5]:find(",") == nil and #t[5] == 1 and #t[4] == 1 then -- biallic SNP
 427                                 local x, z = -1, {};
 428                                 if t[9]:find('PL') then
 429                                         for i = 10, #t do
 430                                                 local AA, Aa, aa = t[i]:match('(%d+),(%d+),(%d+)')
 431                                                 AA = tonumber(AA); Aa = tonumber(Aa); aa = tonumber(aa);
 432                                                 if AA ~= nil then
 433                                                         z[#z+1] = lookup[AA]; z[#z+1] = lookup[Aa]; z[#z+1] = lookup[aa];
 434                                                 else z[#z+1] = 1; z[#z+1] = 1; z[#z+1] = 1; end
 435                                         end
 436                                         print(t[1]..':'..t[2], t[4], t[5], table.concat(z, '\t'))
 437                                 elseif t[9]:find('GL') then
 438                                         print('Error: not implemented')
 439                                         os.exit(1)
 440                                 end
 441                         end
 442                 end
 443         end
 444         fp:close()
 445 end
 446
 447 -- CMD bgl2vcf: convert Beagle output to VCF
 448 function cmd_bgl2vcf()
 449         if #arg < 2 then
 450                 print('Usage: bgl2vcf.lua <in.phased> <in.gprobs>')
 451                 os.exit(1)
 452         end
 453
 454         local fpp = io.xopen(arg[1]);
 455         local fpg = io.xopen(arg[2]);
 456         for lg in fpg:lines() do
 457                 local tp, tg, a = fpp:read():split('%s'), lg:split('%s', 4), {}
 458                 if tp[1] == 'I' then
 459                         for i = 3, #tp, 2 do a[#a+1] = tp[i] end
 460                         print('#CHROM', 'POS', 'ID', 'REF', 'ALT', 'QUAL', 'FILTER', 'INFO', 'FORMAT', table.concat(a, '\t'))
 461                 else
 462                         local chr, pos = tg[1]:match('(%S+):(%d+)$')
 463                         a = {chr, pos, '.', tg[2], tg[3], 30, '.', '.', 'GT'}
 464                         for i = 3, #tp, 2 do
 465                                 a[#a+1] = ((tp[i] == tg[2] and 0) or 1) .. '|' .. ((tp[i+1] == tg[2] and 0) or 1)
 466                         end
 467                         print(table.concat(a, '\t'))
 468                 end
 469         end
 470         fpg:close(); fpp:close();
 471 end
 472
 473 -- CMD freq: count alleles in each population
 474 function cmd_freq()
 475         -- parse the command line
 476         local site_only = true; -- print site allele frequency or not
 477         for c in os.getopt(arg, 's') do
 478                 if c == 's' then site_only = false end
 479         end
 480         if #arg == 0 then
 481                 print("\nUsage: vcfutils.lua freq [-s] <in.vcf> [samples.txt]\n")
 482                 print("NB: 1) This command only considers biallelic variants.")
 483                 print("    2) Apply '-s' to get the allele frequency spectrum.")
 484                 print("    3) 'samples.txt' is TAB-delimited with each line consisting of sample and population.")
 485                 print("")
 486                 os.exit(1)
 487         end
 488
 489         -- read the sample-population pairs
 490         local pop, sample = {}, {}
 491         if #arg > 1 then
 492                 local fp = io.xopen(arg[2]);
 493                 for l in fp:lines() do
 494                         local s, p = l:match("^(%S+)%s+(%S+)"); -- sample, population pair
 495                         sample[s] = p; -- FIXME: check duplications
 496                         if pop[p] then table.insert(pop[p], s)
 497                         else pop[p] = {s} end
 498                 end
 499                 fp:close();
 500         end
 501         pop['NA'] = {}
 502
 503         -- parse VCF
 504         fp = (#arg >= 2 and io.xopen(arg[1])) or io.stdin;
 505         local col, cnt = {}, {};
 506         for k in pairs(pop) do
 507                 col[k], cnt[k] = {}, {[0]=0};
 508         end
 509         for l in fp:lines() do
 510                 if l:sub(1, 2) == '##' then -- meta lines; do nothing
 511                 elseif l:sub(1, 1) == '#' then -- the sample line
 512                         local t, del_NA = l:split('\t'), true;
 513                         for i = 10, #t do
 514                                 local k = sample[t[i]]
 515                                 if k == nil then
 516                                         k, del_NA = 'NA', false
 517                                         table.insert(pop[k], t[i])
 518                                 end
 519                                 table.insert(col[k], i);
 520                                 table.insert(cnt[k], 0);
 521                                 table.insert(cnt[k], 0);
 522                         end
 523                         if del_NA then pop['NA'], col['NA'], cnt['NA'] = nil, nil, nil end
 524                 else -- data lines
 525                         local t = l:split('\t');
 526                         if t[5] ~= '.' and t[5]:find(",") == nil then -- biallic
 527                                 if site_only == true then io.write(t[1], '\t', t[2], '\t', t[4], '\t', t[5]) end
 528                                 for k, v in pairs(col) do
 529                                         local ac, an = 0, 0;
 530                                         for i = 1, #v do
 531                                                 local a1, a2 = t[v[i]]:match("^(%d).(%d)");
 532                                                 if a1 ~= nil then ac, an = ac + a1 + a2, an + 2 end
 533                                         end
 534                                         if site_only == true then io.write('\t', k, ':', an, ':', ac) end
 535                                         if an == #cnt[k] then cnt[k][ac] = cnt[k][ac] + 1 end
 536                                 end
 537                                 if site_only == true then io.write('\n') end
 538                         end
 539                 end
 540         end
 541         fp:close();
 542
 543         -- print
 544         if site_only == false then
 545                 for k, v in pairs(cnt) do
 546                         io.write(k .. "\t" .. #v);
 547                         for i = 0, #v do io.write("\t" .. v[i]) end
 548                         io.write('\n');
 549                 end
 550         end
 551 end
 552
 553 function cmd_vcf2chi2()
 554         if #arg < 3 then
 555                 print("Usage: vcfutils.lua vcf2chi2 <in.vcf> <group1.list> <group2.list>");
 556                 os.exit(1)
 557         end
 558
 559         local g = {};
 560
 561         -- read the list of groups
 562         local fp = io.xopen(arg[2]);
 563         for l in fp:lines() do local x = l:match("^(%S+)"); g[x] = 1 end -- FIXME: check duplicate
 564         fp:close()
 565         fp = io.xopen(arg[3]);
 566         for l in fp:lines() do local x = l:match("^(%S+)"); g[x] = 2 end
 567         fp:close()
 568
 569         -- process VCF
 570         fp = io.xopen(arg[1])
 571         local h = {{}, {}}
 572         for l in fp:lines() do
 573                 if l:sub(1, 2) == '##' then print(l) -- meta lines; do nothing
 574                 elseif l:sub(1, 1) == '#' then -- sample lines
 575                         local t = l:split('\t');
 576                         for i = 10, #t do
 577                                 if g[t[i]] == 1 then table.insert(h[1], i)
 578                                 elseif g[t[i]] == 2 then table.insert(h[2], i) end
 579                         end
 580                         while #t > 8 do table.remove(t) end
 581                         print(table.concat(t, "\t"))
 582                 else -- data line
 583                         local t = l:split('\t');
 584                         if t[5] ~= '.' and t[5]:find(",") == nil then -- biallic
 585                                 local a = {{0, 0}, {0, 0}}
 586                                 for i = 1, 2 do
 587                                         for _, k in pairs(h[i]) do
 588                                                 if t[k]:find("^0.0") then a[i][1] = a[i][1] + 2
 589                                                 elseif t[k]:find("^1.1") then a[i][2] = a[i][2] + 2
 590                                                 elseif t[k]:find("^0.1") or t[k]:find("^1.0") then
 591                                                         a[i][1], a[i][2] = a[i][1] + 1, a[i][2] + 1
 592                                                 end
 593                                         end
 594                                 end
 595                                 local chi2, p, succ = matrix.chi2(a);
 596                                 while #t > 8 do table.remove(t) end
 597                                 --print(a[1][1], a[1][2], a[2][1], a[2][2], chi2, p);
 598                                 if succ then print(table.concat(t, "\t") .. ";PCHI2=" .. string.format("%.3g", p)
 599                                                 .. string.format(';AF1=%.4g;AF2=%.4g,%.4g', (a[1][2]+a[2][2]) / (a[1][1]+a[1][2]+a[2][1]+a[2][2]),
 600                                                 a[1][2]/(a[1][1]+a[1][2]), a[2][2]/(a[2][1]+a[2][2])))
 601                                 else print(table.concat(t, "\t")) end
 602                         end
 603                 end
 604         end
 605         fp:close()
 606 end
 607
 608 -- CMD: compute r^2
 609 function cmd_r2()
 610         local w, is_ht, is_gt = 1, false, false
 611         for o, a in os.getopt(arg, 'w:hg') do
 612                 if o == 'w' then w = tonumber(a)
 613                 elseif o == 'h' then is_ht, is_gt = true, true
 614                 elseif o == 'g' then is_gt = true
 615                 end
 616         end
 617         if #arg == 0 then
 618                 print("Usage: vcfutils.lua r2 [-hg] [-w 1] <in.vcf>")
 619                 os.exit(1)
 620         end
 621         local stack, fp, q2p = {}, io.xopen(arg[1]), algo_init_q2p(1023)
 622         for l in fp:lines() do
 623                 if l:sub(1, 1) ~= '#' then
 624                         local t = l:split('\t')
 625                         local x = text_parse_pl(t, q2p)
 626                         if #t[5] == 1 and t[5] ~= '.' then -- biallelic
 627                                 local r2 = {}
 628                                 for k = 1, w do
 629                                         if is_gt == false then -- use PL
 630                                                 if stack[k] then
 631                                                         local pdg = { stack[k][5], x[5] }
 632                                                         r2[#r2+1] = algo_r2(algo_hapfreq2(pdg))
 633                                                 else r2[#r2+1] = 0 end
 634                                         elseif is_ht == false then -- use unphased GT
 635                                                 if stack[k] then
 636                                                         local pdg = { stack[k][4], x[4] }
 637                                                         r2[#r2+1] = algo_r2(algo_hapfreq2(pdg))
 638                                                 else r2[#r2+1] = 0 end
 639                                         else -- use phased GT
 640                                                 if stack[k] then
 641                                                         local f, ht = { [0]=0, 0, 0, 0 }, { stack[k][3], x[3] }
 642                                                         for i = 1, #ht[1] do
 643                                                                 local j = ht[1][i] * 2 + ht[2][i]
 644                                                                 f[j] = f[j] + 1
 645                                                         end
 646                                                         local sum = f[0] + f[1] + f[2] + f[3]
 647                                                         for k = 0, 3 do f[k] = f[k] / sum end
 648                                                         r2[#r2+1] = algo_r2(f)
 649                                                 else r2[#r2+1] = 0 end
 650                                         end
 651                                 end
 652                                 for k = 1, #r2 do
 653                                         r2[k] = string.format('%.3f', r2[k])
 654                                 end
 655                                 print(x[1], x[2], table.concat(r2, '\t'))
 656                                 if #stack == w then table.remove(stack, 1) end
 657                                 stack[#stack+1] = x
 658                         end
 659                 end
 660         end
 661         fp:close()
 662 end
 663
 664 -------------------
 665 -- END: commands --
 666 -------------------
 667
 668
 669 -------------------
 670 -- MAIN FUNCTION --
 671 -------------------
 672
 673 if #arg == 0 then
 674         print("\nUsage:   vcfutils.lua <command> <arguments>\n")
 675         print("Command: freq        count biallelic alleles in each population")
 676         print("         r2          compute r^2")
 677         print("         vcf2chi2    compute 1-degree chi-square between two groups of samples")
 678         print("         vcf2bgl     convert PL annotated VCF to Beagle input")
 679         print("         bgl2vcf     convert Beagle input to VCF")
 680         print("")
 681         os.exit(1)
 682 end
 683
 684 local cmd = arg[1]
 685 table.remove(arg, 1)
 686 if cmd == 'vcf2bgl' then cmd_vcf2bgl()
 687 elseif cmd == 'bgl2vcf' then cmd_bgl2vcf()
 688 elseif cmd == 'freq' then cmd_freq()
 689 elseif cmd == 'r2' then cmd_r2()
 690 elseif cmd == 'vcf2chi2' then cmd_vcf2chi2()
 691 else
 692         print('ERROR: unknown command "' .. cmd .. '"')
 693         os.exit(1)
 694 end