--- /dev/null
+#!/usr/bin/env luajit
+
+-----------------------------------
+-- BEGIN: routines from klib.lua --
+-----------------------------------
+
+-- Description: getopt() translated from the BSD getopt(); compatible with the default Unix getopt()
+--[[ Example:
+ for o, a in os.getopt(arg, 'a:b') do
+ print(o, a)
+ end
+]]--
+function os.getopt(args, ostr)
+ local arg, place = nil, 0;
+ return function ()
+ if place == 0 then -- update scanning pointer
+ place = 1
+ if #args == 0 or args[1]:sub(1, 1) ~= '-' then place = 0; return nil end
+ if #args[1] >= 2 then
+ place = place + 1
+ if args[1]:sub(2, 2) == '-' then -- found "--"
+ table.remove(args, 1);
+ place = 0
+ return nil;
+ end
+ end
+ end
+ local optopt = place <= #args[1] and args[1]:sub(place, place) or nil
+ place = place + 1;
+ local oli = optopt and ostr:find(optopt) or nil
+ if optopt == ':' or oli == nil then -- unknown option
+ if optopt == '-' then return nil end
+ if place > #args[1] then
+ table.remove(args, 1);
+ place = 0;
+ end
+ return '?';
+ end
+ oli = oli + 1;
+ if ostr:sub(oli, oli) ~= ':' then -- do not need argument
+ arg = nil;
+ if place > #args[1] then
+ table.remove(args, 1);
+ place = 0;
+ end
+ else -- need an argument
+ if place <= #args[1] then -- no white space
+ arg = args[1]:sub(place);
+ else
+ table.remove(args, 1);
+ if #args == 0 then -- an option requiring argument is the last one
+ place = 0;
+ if ostr:sub(1, 1) == ':' then return ':' end
+ return '?';
+ else arg = args[1] end
+ end
+ table.remove(args, 1);
+ place = 0;
+ end
+ return optopt, arg;
+ end
+end
+
+-- Description: string split
+function string:split(sep, n)
+ local a, start = {}, 1;
+ sep = sep or "%s+";
+ repeat
+ local b, e = self:find(sep, start);
+ if b == nil then
+ table.insert(a, self:sub(start));
+ break
+ end
+ a[#a+1] = self:sub(start, b - 1);
+ start = e + 1;
+ if n and #a == n then
+ table.insert(a, self:sub(start));
+ break
+ end
+ until start > #self;
+ return a;
+end
+
+-- Description: smart file open
+function io.xopen(fn, mode)
+ mode = mode or 'r';
+ if fn == nil then return io.stdin;
+ elseif fn == '-' then return (mode == 'r' and io.stdin) or io.stdout;
+ elseif fn:sub(-3) == '.gz' then return (mode == 'r' and io.popen('gzip -dc ' .. fn, 'r')) or io.popen('gzip > ' .. fn, 'w');
+ elseif fn:sub(-4) == '.bz2' then return (mode == 'r' and io.popen('bzip2 -dc ' .. fn, 'r')) or io.popen('bgzip2 > ' .. fn, 'w');
+ else return io.open(fn, mode) end
+end
+
+-- Description: log gamma function
+-- Required by: math.lbinom()
+-- Reference: AS245, 2nd algorithm, http://lib.stat.cmu.edu/apstat/245
+function math.lgamma(z)
+ local x;
+ x = 0.1659470187408462e-06 / (z+7);
+ x = x + 0.9934937113930748e-05 / (z+6);
+ x = x - 0.1385710331296526 / (z+5);
+ x = x + 12.50734324009056 / (z+4);
+ x = x - 176.6150291498386 / (z+3);
+ x = x + 771.3234287757674 / (z+2);
+ x = x - 1259.139216722289 / (z+1);
+ x = x + 676.5203681218835 / z;
+ x = x + 0.9999999999995183;
+ return math.log(x) - 5.58106146679532777 - z + (z-0.5) * math.log(z+6.5);
+end
+
+-- Description: regularized incomplete gamma function
+-- Dependent on: math.lgamma()
+--[[
+ Formulas are taken from Wiki, with additional input from Numerical
+ Recipes in C (for modified Lentz's algorithm) and AS245
+ (http://lib.stat.cmu.edu/apstat/245).
+
+ A good online calculator is available at:
+
+ http://www.danielsoper.com/statcalc/calc23.aspx
+
+ It calculates upper incomplete gamma function, which equals
+ math.igamma(s,z,true)*math.exp(math.lgamma(s))
+]]--
+function math.igamma(s, z, complement)
+
+ local function _kf_gammap(s, z)
+ local sum, x = 1, 1;
+ for k = 1, 100 do
+ x = x * z / (s + k);
+ sum = sum + x;
+ if x / sum < 1e-14 then break end
+ end
+ return math.exp(s * math.log(z) - z - math.lgamma(s + 1.) + math.log(sum));
+ end
+
+ local function _kf_gammaq(s, z)
+ local C, D, f, TINY;
+ f = 1. + z - s; C = f; D = 0.; TINY = 1e-290;
+ -- Modified Lentz's algorithm for computing continued fraction. See Numerical Recipes in C, 2nd edition, section 5.2
+ for j = 1, 100 do
+ local d;
+ local a, b = j * (s - j), j*2 + 1 + z - s;
+ D = b + a * D;
+ if D < TINY then D = TINY end
+ C = b + a / C;
+ if C < TINY then C = TINY end
+ D = 1. / D;
+ d = C * D;
+ f = f * d;
+ if math.abs(d - 1) < 1e-14 then break end
+ end
+ return math.exp(s * math.log(z) - z - math.lgamma(s) - math.log(f));
+ end
+
+ if complement then
+ return ((z <= 1 or z < s) and 1 - _kf_gammap(s, z)) or _kf_gammaq(s, z);
+ else
+ return ((z <= 1 or z < s) and _kf_gammap(s, z)) or (1 - _kf_gammaq(s, z));
+ end
+end
+
+function math.brent(func, a, b, tol)
+ local gold1, gold2, tiny, max_iter = 1.6180339887, 0.3819660113, 1e-20, 100
+
+ local fa, fb = func(a, data), func(b, data)
+ if fb > fa then -- swap, such that f(a) > f(b)
+ a, b, fa, fb = b, a, fb, fa
+ end
+ local c = b + gold1 * (b - a)
+ local fc = func(c) -- golden section extrapolation
+ while fb > fc do
+ local bound = b + 100.0 * (c - b) -- the farthest point where we want to go
+ local r = (b - a) * (fb - fc)
+ local q = (b - c) * (fb - fa)
+ if math.abs(q - r) < tiny then -- avoid 0 denominator
+ tmp = q > r and tiny or 0.0 - tiny
+ else tmp = q - r end
+ u = b - ((b - c) * q - (b - a) * r) / (2.0 * tmp) -- u is the parabolic extrapolation point
+ if (b > u and u > c) or (b < u and u < c) then -- u lies between b and c
+ fu = func(u)
+ if fu < fc then -- (b,u,c) bracket the minimum
+ a, b, fa, fb = b, u, fb, fu
+ break
+ elseif fu > fb then -- (a,b,u) bracket the minimum
+ c, fc = u, fu
+ break
+ end
+ u = c + gold1 * (c - b)
+ fu = func(u) -- golden section extrapolation
+ elseif (c > u and u > bound) or (c < u and u < bound) then -- u lies between c and bound
+ fu = func(u)
+ if fu < fc then -- fb > fc > fu
+ b, c, u = c, u, c + gold1 * (c - b)
+ fb, fc, fu = fc, fu, func(u)
+ else -- (b,c,u) bracket the minimum
+ a, b, c = b, c, u
+ fa, fb, fc = fb, fc, fu
+ break
+ end
+ elseif (u > bound and bound > c) or (u < bound and bound < c) then -- u goes beyond the bound
+ u = bound
+ fu = func(u)
+ else -- u goes the other way around, use golden section extrapolation
+ u = c + gold1 * (c - b)
+ fu = func(u)
+ end
+ a, b, c = b, c, u
+ fa, fb, fc = fb, fc, fu
+ end
+ if a > c then a, c = c, a end -- swap
+
+ -- now, a<b<c, fa>fb and fb<fc, move on to Brent's algorithm
+ local e, d = 0, 0
+ local w, v, fw, fv
+ w, v = b, b
+ fw, fv = fb, fb
+ for iter = 1, max_iter do
+ local mid = 0.5 * (a + c)
+ local tol1 = tol * math.abs(b) + tiny
+ local tol2 = 2.0 * tol1
+ if math.abs(b - mid) <= tol2 - 0.5 * (c - a) then return fb, b end -- found
+ if math.abs(e) > tol1 then
+ -- related to parabolic interpolation
+ local r = (b - w) * (fb - fv)
+ local q = (b - v) * (fb - fw)
+ local p = (b - v) * q - (b - w) * r
+ q = 2.0 * (q - r)
+ if q > 0.0 then p = 0.0 - p
+ else q = 0.0 - q end
+ eold, e = e, d
+ if math.abs(p) >= math.abs(0.5 * q * eold) or p <= q * (a - b) or p >= q * (c - b) then
+ e = b >= mid and a - b or c - b
+ d = gold2 * e
+ else
+ d, u = p / q, b + d -- actual parabolic interpolation happens here
+ if u - a < tol2 or c - u < tol2 then
+ d = mid > b and tol1 or 0.0 - tol1
+ end
+ end
+ else -- golden section interpolation
+ e = b >= min and a - b or c - b
+ d = gold2 * e
+ end
+ u = fabs(d) >= tol1 and b + d or b + (d > 0.0 and tol1 or -tol1);
+ fu = func(u)
+ if fu <= fb then -- u is the minimum point so far
+ if u >= b then a = b
+ else c = b end
+ v, w, b = w, b, u
+ fv, fw, fb = fw, fb, fu
+ else -- adjust (a,c) and (u,v,w)
+ if u < b then a = u
+ else c = u end
+ if fu <= fw or w == b then
+ v, w = w, u
+ fv, fw = fw, fu
+ elseif fu <= fv or v == b or v == w then
+ v, fv = u, fu;
+ end
+ end
+ end
+ return fb, b
+end
+
+matrix = {}
+
+-- Description: chi^2 test for contingency tables
+-- Dependent on: math.igamma()
+function matrix.chi2(a)
+ if #a == 2 and #a[1] == 2 then -- 2x2 table
+ local x, z
+ 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])
+ if x == 0 then return 0, 1, false end
+ z = a[1][1] * a[2][2] - a[1][2] * a[2][1]
+ z = (a[1][1] + a[1][2] + a[2][1] + a[2][2]) * z * z / x
+ return z, math.igamma(.5, .5 * z, true), true
+ else -- generic table
+ local rs, cs, n, m, N, z = {}, {}, #a, #a[1], 0, 0
+ for i = 1, n do rs[i] = 0 end
+ for j = 1, m do cs[j] = 0 end
+ for i = 1, n do -- compute column sum and row sum
+ for j = 1, m do cs[j], rs[i] = cs[j] + a[i][j], rs[i] + a[i][j] end
+ end
+ for i = 1, n do N = N + rs[i] end
+ for i = 1, n do -- compute the chi^2 statistics
+ for j = 1, m do
+ local E = rs[i] * cs[j] / N;
+ z = z + (a[i][j] - E) * (a[i][j] - E) / E
+ end
+ end
+ return z, math.igamma(.5 * (n-1) * (m-1), .5 * z, true), true;
+ end
+end
+
+---------------------------------
+-- END: routines from klib.lua --
+---------------------------------
+
+
+--------------------------
+-- BEGIN: misc routines --
+--------------------------
+
+-- precompute an array for PL->probability conversion
+-- @param m maximum PL
+function algo_init_q2p(m)
+ local q2p = {}
+ for i = 0, m do
+ q2p[i] = math.pow(10, -i / 10)
+ end
+ return q2p
+end
+
+-- given the haplotype frequency, compute r^2
+-- @param f 4 haplotype frequencies; f[] is 0-indexed.
+-- @return r^2
+function algo_r2(f)
+ local p = { f[0] + f[1], f[0] + f[2] }
+ local D = f[0] * f[3] - f[1] * f[2]
+ 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]))
+end
+
+-- parse a VCF line to get PL
+-- @param q2p is computed by algo_init_q2p()
+function text_parse_pl(t, q2p, parse_GT)
+ parse_GT = parse_GT == nil and true or false
+ local ht, gt, pl = {}, {}, {}
+ local s, j0 = t[9]:split(':'), 0
+ for j = 1, #s do
+ if s[j] == 'PL' then j0 = j break end
+ end
+ local has_GT = (s[1] == 'GT' and parse_GT) and true or false
+ for i = 10, #t do
+ if j0 > 0 then
+ local s = t[i]:split(':')
+ local a, b = 1, s[j0]:find(',')
+ pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, b - 1))]
+ a, b = b + 1, s[j0]:find(',', b + 1)
+ pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, b - 1))]
+ a, b = b + 1, s[j0]:find(',', b + 1)
+ pl[#pl+1] = q2p[tonumber(s[j0]:sub(a, (b and b - 1) or nil))]
+ end
+ if has_GT then
+ if t[i]:sub(1, 1) ~= '.' then
+ local g = tonumber(t[i]:sub(1, 1)) + tonumber(t[i]:sub(3, 3));
+ gt[#gt+1] = 1e-6; gt[#gt+1] = 1e-6; gt[#gt+1] = 1e-6
+ gt[#gt - 2 + g] = 1
+ ht[#ht+1] = tonumber(t[i]:sub(1, 1)); ht[#ht+1] = tonumber(t[i]:sub(3, 3));
+ else
+ gt[#gt+1] = 1; gt[#gt+1] = 1; gt[#gt+1] = 1
+ ht[#ht+1] = -1; ht[#ht+1] = -1;
+ end
+ end
+-- print(t[i], pl[#pl-2], pl[#pl-1], pl[#pl], gt[#gt-2], gt[#gt-1], gt[#gt])
+ end
+ if #pl == 0 then pl = nil end
+ local x = has_GT and { t[1], t[2], ht, gt, pl } or { t[1], t[2], nil, nil, pl }
+ return x
+end
+
+-- Infer haplotype frequency
+-- @param pdg genotype likelihoods P(D|g) generated by text_parse_pl(). pdg[] is 1-indexed.
+-- @param eps precision [1e-5]
+-- @return 2-locus haplotype frequencies, 0-indexed array
+function algo_hapfreq2(pdg, eps)
+ eps = eps or 1e-5
+ local n, f = #pdg[1] / 3, {[0]=0.25, 0.25, 0.25, 0.25}
+ for iter = 1, 100 do
+ local F = {[0]=0, 0, 0, 0}
+ for i = 0, n - 1 do
+ 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]}
+ local u = { [0]=
+ 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]),
+ 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]),
+ 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]),
+ 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])
+ }
+ local s = u[0] + u[1] + u[2] + u[3]
+ s = 1 / (s * n)
+ F[0] = F[0] + u[0] * s
+ F[1] = F[1] + u[1] * s
+ F[2] = F[2] + u[2] * s
+ F[3] = F[3] + u[3] * s
+ end
+ local e = 0
+ for k = 0, 3 do
+ e = math.abs(f[k] - F[k]) > e and math.abs(f[k] - F[k]) or e
+ end
+ for k = 0, 3 do f[k] = F[k] end
+ if e < eps then break end
+-- print(f[0], f[1], f[2], f[3])
+ end
+ return f
+end
+
+------------------------
+-- END: misc routines --
+------------------------
+
+
+---------------------
+-- BEGIN: commands --
+---------------------
+
+-- CMD vcf2bgl: convert PL tagged VCF to Beagle input --
+function cmd_vcf2bgl()
+ if #arg == 0 then
+ print("\nUsage: vcf2bgl.lua <in.vcf>")
+ print("\nNB: This command finds PL by matching /(\\d+),(\\d+),(\\d+)/.\n");
+ os.exit(1)
+ end
+
+ local lookup = {}
+ for i = 0, 10000 do lookup[i] = string.format("%.4f", math.pow(10, -i/10)) end
+
+ local fp = io.xopen(arg[1])
+ for l in fp:lines() do
+ if l:sub(1, 2) == '##' then -- meta lines; do nothing
+ elseif l:sub(1, 1) == '#' then -- sample lines
+ local t, s = l:split('\t'), {}
+ for i = 10, #t do s[#s+1] = t[i]; s[#s+1] = t[i]; s[#s+1] = t[i] end
+ print('marker', 'alleleA', 'alleleB', table.concat(s, '\t'))
+ else -- data line
+ local t = l:split('\t');
+ if t[5] ~= '.' and t[5]:find(",") == nil and #t[5] == 1 and #t[4] == 1 then -- biallic SNP
+ local x, z = -1, {};
+ if t[9]:find('PL') then
+ for i = 10, #t do
+ local AA, Aa, aa = t[i]:match('(%d+),(%d+),(%d+)')
+ AA = tonumber(AA); Aa = tonumber(Aa); aa = tonumber(aa);
+ if AA ~= nil then
+ z[#z+1] = lookup[AA]; z[#z+1] = lookup[Aa]; z[#z+1] = lookup[aa];
+ else z[#z+1] = 1; z[#z+1] = 1; z[#z+1] = 1; end
+ end
+ print(t[1]..':'..t[2], t[4], t[5], table.concat(z, '\t'))
+ elseif t[9]:find('GL') then
+ print('Error: not implemented')
+ os.exit(1)
+ end
+ end
+ end
+ end
+ fp:close()
+end
+
+-- CMD bgl2vcf: convert Beagle output to VCF
+function cmd_bgl2vcf()
+ if #arg < 2 then
+ print('Usage: bgl2vcf.lua <in.phased> <in.gprobs>')
+ os.exit(1)
+ end
+
+ local fpp = io.xopen(arg[1]);
+ local fpg = io.xopen(arg[2]);
+ for lg in fpg:lines() do
+ local tp, tg, a = fpp:read():split('%s'), lg:split('%s', 4), {}
+ if tp[1] == 'I' then
+ for i = 3, #tp, 2 do a[#a+1] = tp[i] end
+ print('#CHROM', 'POS', 'ID', 'REF', 'ALT', 'QUAL', 'FILTER', 'INFO', 'FORMAT', table.concat(a, '\t'))
+ else
+ local chr, pos = tg[1]:match('(%S+):(%d+)$')
+ a = {chr, pos, '.', tg[2], tg[3], 30, '.', '.', 'GT'}
+ for i = 3, #tp, 2 do
+ a[#a+1] = ((tp[i] == tg[2] and 0) or 1) .. '|' .. ((tp[i+1] == tg[2] and 0) or 1)
+ end
+ print(table.concat(a, '\t'))
+ end
+ end
+ fpg:close(); fpp:close();
+end
+
+-- CMD freq: count alleles in each population
+function cmd_freq()
+ -- parse the command line
+ local site_only = true; -- print site allele frequency or not
+ for c in os.getopt(arg, 's') do
+ if c == 's' then site_only = false end
+ end
+ if #arg == 0 then
+ print("\nUsage: vcfutils.lua freq [-s] <in.vcf> [samples.txt]\n")
+ print("NB: 1) This command only considers biallelic variants.")
+ print(" 2) Apply '-s' to get the allele frequency spectrum.")
+ print(" 3) 'samples.txt' is TAB-delimited with each line consisting of sample and population.")
+ print("")
+ os.exit(1)
+ end
+
+ -- read the sample-population pairs
+ local pop, sample = {}, {}
+ if #arg > 1 then
+ local fp = io.xopen(arg[2]);
+ for l in fp:lines() do
+ local s, p = l:match("^(%S+)%s+(%S+)"); -- sample, population pair
+ sample[s] = p; -- FIXME: check duplications
+ if pop[p] then table.insert(pop[p], s)
+ else pop[p] = {s} end
+ end
+ fp:close();
+ end
+ pop['NA'] = {}
+
+ -- parse VCF
+ fp = (#arg >= 2 and io.xopen(arg[1])) or io.stdin;
+ local col, cnt = {}, {};
+ for k in pairs(pop) do
+ col[k], cnt[k] = {}, {[0]=0};
+ end
+ for l in fp:lines() do
+ if l:sub(1, 2) == '##' then -- meta lines; do nothing
+ elseif l:sub(1, 1) == '#' then -- the sample line
+ local t, del_NA = l:split('\t'), true;
+ for i = 10, #t do
+ local k = sample[t[i]]
+ if k == nil then
+ k, del_NA = 'NA', false
+ table.insert(pop[k], t[i])
+ end
+ table.insert(col[k], i);
+ table.insert(cnt[k], 0);
+ table.insert(cnt[k], 0);
+ end
+ if del_NA then pop['NA'], col['NA'], cnt['NA'] = nil, nil, nil end
+ else -- data lines
+ local t = l:split('\t');
+ if t[5] ~= '.' and t[5]:find(",") == nil then -- biallic
+ if site_only == true then io.write(t[1], '\t', t[2], '\t', t[4], '\t', t[5]) end
+ for k, v in pairs(col) do
+ local ac, an = 0, 0;
+ for i = 1, #v do
+ local a1, a2 = t[v[i]]:match("^(%d).(%d)");
+ if a1 ~= nil then ac, an = ac + a1 + a2, an + 2 end
+ end
+ if site_only == true then io.write('\t', k, ':', an, ':', ac) end
+ if an == #cnt[k] then cnt[k][ac] = cnt[k][ac] + 1 end
+ end
+ if site_only == true then io.write('\n') end
+ end
+ end
+ end
+ fp:close();
+
+ -- print
+ if site_only == false then
+ for k, v in pairs(cnt) do
+ io.write(k .. "\t" .. #v);
+ for i = 0, #v do io.write("\t" .. v[i]) end
+ io.write('\n');
+ end
+ end
+end
+
+function cmd_vcf2chi2()
+ if #arg < 3 then
+ print("Usage: vcfutils.lua vcf2chi2 <in.vcf> <group1.list> <group2.list>");
+ os.exit(1)
+ end
+
+ local g = {};
+
+ -- read the list of groups
+ local fp = io.xopen(arg[2]);
+ for l in fp:lines() do local x = l:match("^(%S+)"); g[x] = 1 end -- FIXME: check duplicate
+ fp:close()
+ fp = io.xopen(arg[3]);
+ for l in fp:lines() do local x = l:match("^(%S+)"); g[x] = 2 end
+ fp:close()
+
+ -- process VCF
+ fp = io.xopen(arg[1])
+ local h = {{}, {}}
+ for l in fp:lines() do
+ if l:sub(1, 2) == '##' then print(l) -- meta lines; do nothing
+ elseif l:sub(1, 1) == '#' then -- sample lines
+ local t = l:split('\t');
+ for i = 10, #t do
+ if g[t[i]] == 1 then table.insert(h[1], i)
+ elseif g[t[i]] == 2 then table.insert(h[2], i) end
+ end
+ while #t > 8 do table.remove(t) end
+ print(table.concat(t, "\t"))
+ else -- data line
+ local t = l:split('\t');
+ if t[5] ~= '.' and t[5]:find(",") == nil then -- biallic
+ local a = {{0, 0}, {0, 0}}
+ for i = 1, 2 do
+ for _, k in pairs(h[i]) do
+ if t[k]:find("^0.0") then a[i][1] = a[i][1] + 2
+ elseif t[k]:find("^1.1") then a[i][2] = a[i][2] + 2
+ elseif t[k]:find("^0.1") or t[k]:find("^1.0") then
+ a[i][1], a[i][2] = a[i][1] + 1, a[i][2] + 1
+ end
+ end
+ end
+ local chi2, p, succ = matrix.chi2(a);
+ while #t > 8 do table.remove(t) end
+ --print(a[1][1], a[1][2], a[2][1], a[2][2], chi2, p);
+ if succ then print(table.concat(t, "\t") .. ";PCHI2=" .. string.format("%.3g", p)
+ .. 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]),
+ a[1][2]/(a[1][1]+a[1][2]), a[2][2]/(a[2][1]+a[2][2])))
+ else print(table.concat(t, "\t")) end
+ end
+ end
+ end
+ fp:close()
+end
+
+-- CMD: compute r^2
+function cmd_r2()
+ local w, is_ht, is_gt = 1, false, false
+ for o, a in os.getopt(arg, 'w:hg') do
+ if o == 'w' then w = tonumber(a)
+ elseif o == 'h' then is_ht, is_gt = true, true
+ elseif o == 'g' then is_gt = true
+ end
+ end
+ if #arg == 0 then
+ print("Usage: vcfutils.lua r2 [-hg] [-w 1] <in.vcf>")
+ os.exit(1)
+ end
+ local stack, fp, q2p = {}, io.xopen(arg[1]), algo_init_q2p(1023)
+ for l in fp:lines() do
+ if l:sub(1, 1) ~= '#' then
+ local t = l:split('\t')
+ local x = text_parse_pl(t, q2p)
+ if #t[5] == 1 and t[5] ~= '.' then -- biallelic
+ local r2 = {}
+ for k = 1, w do
+ if is_gt == false then -- use PL
+ if stack[k] then
+ local pdg = { stack[k][5], x[5] }
+ r2[#r2+1] = algo_r2(algo_hapfreq2(pdg))
+ else r2[#r2+1] = 0 end
+ elseif is_ht == false then -- use unphased GT
+ if stack[k] then
+ local pdg = { stack[k][4], x[4] }
+ r2[#r2+1] = algo_r2(algo_hapfreq2(pdg))
+ else r2[#r2+1] = 0 end
+ else -- use phased GT
+ if stack[k] then
+ local f, ht = { [0]=0, 0, 0, 0 }, { stack[k][3], x[3] }
+ for i = 1, #ht[1] do
+ local j = ht[1][i] * 2 + ht[2][i]
+ f[j] = f[j] + 1
+ end
+ local sum = f[0] + f[1] + f[2] + f[3]
+ for k = 0, 3 do f[k] = f[k] / sum end
+ r2[#r2+1] = algo_r2(f)
+ else r2[#r2+1] = 0 end
+ end
+ end
+ for k = 1, #r2 do
+ r2[k] = string.format('%.3f', r2[k])
+ end
+ print(x[1], x[2], table.concat(r2, '\t'))
+ if #stack == w then table.remove(stack, 1) end
+ stack[#stack+1] = x
+ end
+ end
+ end
+ fp:close()
+end
+
+-------------------
+-- END: commands --
+-------------------
+
+
+-------------------
+-- MAIN FUNCTION --
+-------------------
+
+if #arg == 0 then
+ print("\nUsage: vcfutils.lua <command> <arguments>\n")
+ print("Command: freq count biallelic alleles in each population")
+ print(" r2 compute r^2")
+ print(" vcf2chi2 compute 1-degree chi-square between two groups of samples")
+ print(" vcf2bgl convert PL annotated VCF to Beagle input")
+ print(" bgl2vcf convert Beagle input to VCF")
+ print("")
+ os.exit(1)
+end
+
+local cmd = arg[1]
+table.remove(arg, 1)
+if cmd == 'vcf2bgl' then cmd_vcf2bgl()
+elseif cmd == 'bgl2vcf' then cmd_bgl2vcf()
+elseif cmd == 'freq' then cmd_freq()
+elseif cmd == 'r2' then cmd_r2()
+elseif cmd == 'vcf2chi2' then cmd_vcf2chi2()
+else
+ print('ERROR: unknown command "' .. cmd .. '"')
+ os.exit(1)
+end
projects / samtools.git / commitdiff