Genotype Combinations

genotype_combinations.py

import pandas as pd
from collections import Counter
from itertools import combinations_with_replacement


population = [
    'pp',
    'pp',
    'qq',
    'pq',
    'pq',
    'qp',
    'pp'
]
print('Population: {}'.format(population))

alleles = 'pq'
print('Possible Alleles {}'.format(alleles))

counts = Counter(population)
print(counts)

print('Possible Allele Combinations are..')

observed_table = pd.DataFrame()

for genotype in combinations_with_replacement(alleles, 2):
    genotype_string = genotype[0] + genotype[1]

    if genotype[0] == genotype[1]:
        # Homo therefore order doesn't matter
        homo_count = counts[genotype_string]
        print('{} {}'.format(genotype_string, homo_count))
        observed_table.set_value(0, genotype_string, homo_count)
        # df.set_value() is ~100x faster than other individual setting methods
        #observed_table[genotype_string] = homo_count
    else:
        # Hetero therefore order does matter
        hetero_count = counts[genotype_string]
        hetero_count += counts[genotype_string[::-1]]   # I'm using extended slice syntax to reverse the string [::-1]
        print('{} {}'.format(genotype_string, hetero_count))
        observed_table.set_value(0, genotype_string, hetero_count)
        #observed_table[genotype_string] = hetero_count

print(observed_table)