Mendelian transmission

import numpy as np
import scipy
import pandas
import h5py
%matplotlib inline
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style('white')
sns.set_style('ticks')
import allel; print('scikit-allel', allel.__version__)

samples = pandas.read_csv('data/phase1.AR3.1/samples/cross.samples.meta.txt', sep='\t')
samples.head()

samples.cross.value_counts()

callset = h5py.File('data/phase1.AR3/variation/crosses/ar3/hdf5/ag1000g.crosses.phase1.ar3sites.3R.h5',
                    mode='r')
callset

genotypes = allel.GenotypeChunkedArray(callset['3R/calldata/genotype'])
genotypes

# locate the indices of the samples within the callset
sample_indices = samples[samples.cross == '29-2'].index.values.tolist()
sample_indices

# do an allele count to find segregating variants
ac = genotypes.count_alleles(max_allele=3, subpop=sample_indices)[:]
ac

# locate segregating SNPs
loc_seg = ac.is_segregating()
loc_seg

# how many SNPs are segregating within the cross?
np.count_nonzero(loc_seg)

# locate SNPs that passed all quality filters
loc_pass = callset['3R/variants/FILTER_PASS'][:]
loc_pass

# perform the subset and load the results into memory uncompressed
genotypes_cross_pass = genotypes.subset(loc_seg & loc_pass, sample_indices)[:]
genotypes_cross_pass

genotypes_cross_phased = allel.phase_by_transmission(genotypes_cross_pass, window_size=100)
genotypes_cross_phased

# locate variants where all genotype calls were phased
loc_phased_all = np.all(genotypes_cross_phased.is_phased, axis=1)
loc_phased_all

genotypes_cross_phased_all = genotypes_cross_phased[loc_phased_all]
genotypes_cross_phased_all

# pull out mother's genotypes from the first column
genotypes_mother = genotypes_cross_phased_all[:, 0]
# convert to haplotype array
haplotypes_mother = genotypes_mother.to_haplotypes()
# pull out maternal haplotypes from the progeny
haplotypes_progeny_maternal = allel.HaplotypeArray(genotypes_cross_phased_all[:, 2:, 0])
# stack mother's haplotypes alongside haplotypes she transmitted to her progeny
haplotypes_maternal = haplotypes_mother.concatenate(haplotypes_progeny_maternal, axis=1)
haplotypes_maternal

painting_maternal = allel.paint_transmission(haplotypes_mother, haplotypes_progeny_maternal)
painting_maternal

help(allel.paint_transmission)

genotypes_father = genotypes_cross_phased_all[:, 1]
haplotypes_father = genotypes_father.to_haplotypes()
haplotypes_progeny_paternal = allel.HaplotypeArray(genotypes_cross_phased_all[:, 2:, 1])
haplotypes_paternal = haplotypes_father.concatenate(haplotypes_progeny_paternal, axis=1)
haplotypes_paternal

painting_paternal = allel.paint_transmission(haplotypes_father, haplotypes_progeny_paternal)
painting_paternal

# create a colour palette
palette = sns.color_palette('colorblind')
sns.palplot(palette)

def plot_transmission(painting, title=None):
    
    # set figure height depending on number of haplotypes
    fig, ax = plt.subplots(figsize=(12, .2 * painting.shape[1]))
    
    # map painting codes onto colours
    cmap = mpl.colors.ListedColormap([
        'grey',      # 0 = undetermined
        palette[0],  # 1 = allele inherited from first parental haplotype
        palette[2],  # 2 = allele inherited from second parental haplotype
        palette[1],  # 3 = reference allele, also carried by both parental haplotypes
        palette[3],  # 4 = non-reference allele, also carried by both parental haplotypes
        'black',     # 5 = non-parental allele (i.e., Mendelian error)
        palette[5],  # 6 = either or both parental alleles missing
        'white',     # 7 = missing allele
    ])
    
    # plot painting
    ax.pcolormesh(painting.T, cmap=cmap, vmin=0, vmax=7)
    
    # tidy up axes
    ax.set_yticks(np.arange(painting.shape[1]) + .5)
    ax.set_yticklabels(range(painting.shape[1]))
    ax.set_ylabel('Progeny haplotypes')
    ax.set_xlabel('Variants')
    if title:
        ax.set_title(title)

plot_transmission(painting_maternal[::10], title='Maternal haplotypes')

plot_transmission(painting_maternal[genotypes_mother.is_het()][::10], 
                  title='Maternal haplotypes (sites heterozygous in mother)')

plot_transmission(painting_paternal[::10], title='Paternal haplotypes')

plot_transmission(painting_paternal[genotypes_father.is_het()][::10], 
                  title='Paternal haplotypes (sites heterozygous in father)')

# load variant physical positions
pos = callset['3R/variants/POS'][:]
pos_cross_pass = pos[loc_seg & loc_pass]

def plot_density(pos, window_size=100000, title=None):
    
    fig, ax = plt.subplots(figsize=(12, 2))
    sns.despine(ax=ax, offset=5)
    y, windows = allel.windowed_count(pos, size=window_size)
    x = np.mean(windows, axis=1)
    ax.plot(x, y/window_size)
    ax.set_ylabel('Density (bp$^{-1}$)')
    ax.set_xlabel('Position (bp)')
    if title:
        ax.set_title(title)
    

plot_density(pos_cross_pass[genotypes_cross_pass[:, 0].is_het()],
             title='Heterozygosity in mother')

plot_density(pos_cross_pass[genotypes_cross_pass[:, 1].is_het()],
             title='Heterozygosity in father')

genotypes_cross = genotypes.subset(loc_seg, sample_indices)
genotypes_cross

genotypes_parents = genotypes_cross[:, :2]
genotypes_progeny = genotypes_cross[:, 2:]
mendel_errors = allel.mendel_errors(genotypes_parents, genotypes_progeny)
mendel_errors

mendel_errors.shape

mendel_errors_by_sample = np.sum(mendel_errors, axis=0)
mendel_errors_by_sample

# plot errors by sample
fig, ax = plt.subplots(figsize=(10, 5))
sns.despine(ax=ax, offset=5)
left = list(range(len(mendel_errors_by_sample)))
height = mendel_errors_by_sample
ax.bar(left, height)
ax.set_xticks(left)
ax.set_xlabel('Progeny')
ax.set_ylabel('No. errors')
ax.set_title('Mendel errors by individual');

mendel_errors_by_variant = np.sum(mendel_errors, axis=1)
mendel_errors_by_variant

pos_cross = pos[loc_seg]

fig, ax = plt.subplots(figsize=(12, 2))
sns.despine(ax=ax, offset=5)
y, windows, _ = allel.windowed_statistic(pos_cross, values=mendel_errors_by_variant, statistic=np.sum, size=100000)
x = np.mean(windows, axis=1)
ax.plot(x, y)
ax.set_xlabel('Position (bp)')
ax.set_ylabel('No. errors')
ax.set_title('Mendel errors by genome position');

# load some variant annotations
vtbl = allel.VariantChunkedTable(callset['3R/variants'], names=['POS', 'QD', 'MQ'])
# subset annotations to sites segregating in the cross
vtbl_cross = vtbl.compress(loc_seg)[:]
vtbl_cross

def plot_variant_me(f, bins):

    fig, ax = plt.subplots(figsize=(7, 5))
    sns.despine(ax=ax, offset=5)
    x = vtbl_cross[f]

    # plot a histogram
    ax.hist(x, bins=bins)
    ax.set_xlabel(f)
    ax.set_ylabel('No. variants')

    # plot Mendel errors
    ax = ax.twinx()
    sns.despine(ax=ax, bottom=True, left=True, right=False, offset=10)
    values = mendel_errors_by_variant
    with np.errstate(over='ignore'):
        # binned_statistic generates an annoying overflow warning which we can ignore
        y1, _, _ = scipy.stats.binned_statistic(x, values, statistic=np.mean, bins=bins)
    bx = (bins[1:] + bins[:-1]) / 2
    ax.plot(bx, y1, color='k', label='Mendel errors')
    ax.set_ylabel('Mendel error rate')
    ax.legend(loc='upper right')

plot_variant_me('QD', bins=np.arange(0, 40, 1))

plot_variant_me('MQ', bins=np.arange(0, 60, 1))

def plot_joint_variant_me(f1, f2, gridsize=20, mincnt=20, extent=None):
    fig, ax = plt.subplots()
    sns.despine(ax=ax, offset=5)
    x = vtbl_cross[f1]
    y = vtbl_cross[f2]
    C = mendel_errors_by_variant
    im = ax.hexbin(x, y, C=C, reduce_C_function=np.mean, mincnt=mincnt, extent=extent,
                   gridsize=gridsize, cmap='jet')
    fig.colorbar(im)
    ax.set_xlabel(f1)
    ax.set_ylabel(f2)
    ax.set_title('Variant %s versus %s and Mendel error rate' % (f1, f2))

plot_joint_variant_me('QD', 'MQ', extent=(0, 40, 0, 60))