def ld(args):
"""
%prog ld map
Calculate pairwise linkage disequilibrium given MSTmap.
"""
import numpy as np
from random import sample
from jcvi.algorithms.matrix import symmetrize
p = OptionParser(ld.__doc__)
p.add_option("--subsample", default=500, type="int",
help="Subsample markers to speed up [default: %default]")
opts, args, iopts = p.set_image_options(args, figsize="8x8")
if len(args) != 1:
sys.exit(not p.print_help())
mstmap, = args
subsample = opts.subsample
data = MSTMap(mstmap)
# Take random subsample while keeping marker order
if subsample < data.nmarkers:
data = [data[x] for x in \
sorted(sample(xrange(len(data)), subsample))]
markerbedfile = mstmap + ".subsample.bed"
ldmatrix = mstmap + ".subsample.matrix"
if need_update(mstmap, (markerbedfile, ldmatrix)):
nmarkers = len(data)
fw = open(markerbedfile, "w")
print >> fw, "\n".join(x.bedline for x in data)
logging.debug("Write marker set of size {0} to file `{1}`."\
.format(nmarkers, markerbedfile))
M = np.zeros((nmarkers, nmarkers), dtype=float)
for i, j in combinations(range(nmarkers), 2):
a = data[i]
b = data[j]
M[i, j] = calc_ldscore(a.genotype, b.genotype)
M = symmetrize(M)
logging.debug("Write LD matrix to file `{0}`.".format(ldmatrix))
M.tofile(ldmatrix)
else:
nmarkers = len(Bed(markerbedfile))
M = np.fromfile(ldmatrix, dtype="float").reshape(nmarkers, nmarkers)
logging.debug("LD matrix `{0}` exists ({1}x{1})."\
.format(ldmatrix, nmarkers))
from jcvi.graphics.base import plt, savefig, Rectangle, draw_cmap
plt.rcParams["axes.linewidth"] = 0
fig = plt.figure(1, (iopts.w, iopts.h))
root = fig.add_axes([0, 0, 1, 1])
ax = fig.add_axes([.1, .1, .8, .8]) # the heatmap
ax.matshow(M, cmap=iopts.cmap)
# Plot chromosomes breaks
bed = Bed(markerbedfile)
xsize = len(bed)
extent = (0, nmarkers)
chr_labels = []
ignore_size = 20
for (seqid, beg, end) in bed.get_breaks():
ignore = abs(end - beg) < ignore_size
pos = (beg + end) / 2
chr_labels.append((seqid, pos, ignore))
if ignore:
continue
ax.plot((end, end), extent, "w-", lw=1)
ax.plot(extent, (end, end), "w-", lw=1)
# Plot chromosome labels
for label, pos, ignore in chr_labels:
pos = .1 + pos * .8 / xsize
if not ignore:
root.text(pos, .91, label,
ha="center", va="bottom", rotation=45, color="grey")
root.text(.09, pos, label,
ha="right", va="center", color="grey")
ax.set_xlim(extent)
ax.set_ylim(extent)
ax.set_axis_off()
draw_cmap(root, "Pairwise LD (r2)", 0, 1, cmap=default_cm)
root.add_patch(Rectangle((.1, .1), .8, .8, fill=False, ec="k", lw=2))
m = mstmap.split(".")[0]
root.text(.5, .06, "Linkage Disequilibrium between {0} markers".format(m), ha="center")
root.set_xlim(0, 1)
root.set_ylim(0, 1)
root.set_axis_off()
image_name = m + ".subsample" + "." + iopts.format
savefig(image_name, dpi=iopts.dpi, iopts=iopts)