diff PsiCLASS-1.0.2/samtools-0.1.19/misc/vcfutils.lua @ 0:903fc43d6227 draft default tip

Uploaded

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/PsiCLASS-1.0.2/samtools-0.1.19/misc/vcfutils.lua	Fri Mar 26 16:52:45 2021 +0000
@@ -0,0 +1,694 @@
+#!/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