如何简化以下群体遗传学模拟代码?

2024-07-01 07:02:13 发布

您现在位置:Python中文网/ 问答频道 /正文

一个单倍体与所有二倍体之间随机交配的模拟,除了与自身含有相同单倍体的那些。它工作,但似乎超级冗余-例如,字典是一个更好的方式在这种情况下?这是我的密码:

import numpy as np

import itertools

N = 100                                                     # number of haploids in the population
numGen = 5                                                  # number of generations to simulate
gen = {'R0':0.1, 'R1':0.2, 'R2':0.25, 'R3':0.15, 'R4':0.3}      # haplotypes w/ frequencies
gens = []

for i in range(numGen):
    Pstls = [*gen]                                              # put keys into list
    freq = np.round(np.random.multinomial(N, list(gen.values()))/N, 4)
    gen = dict(zip(Pstls, freq))
    mating = []
    for i in itertools.combinations(Pstls,2):                   # random mating
        mating.append(i)                                        # [('R0', 'R1'), ('R0', 'R2'), ('R0', 'R3'), ('R0', 'R4'), ('R1', 'R2'), \
                                                                # ('R1', 'R3'), ('R1', 'R4'), ('R2', 'R3'), ('R2', 'R4'), ('R3', 'R4')]

    haps = []
    haps_freq = []
    for i in range(len(gen)):
        for j in range(len(mating)):
            if Pstls[i] not in mating[j]:                       # new haplotype frequencies w/ duplicates
                hf0 = round(gen[Pstls[i]] * gen[mating[j][0]], 4)
                hf1 = round(gen[Pstls[i]] * gen[mating[j][1]], 4)
                hap0 = Pstls[i] + ',' + mating[j][0]
                hap1 = Pstls[i] + ',' + mating[j][1]
                haps.append(hap0)
                haps.append(hap1)
                haps_freq.append(hf0)
                haps_freq.append(hf1)

    newHaps = dict(zip(haps,haps_freq))
    keys = []                                                   # preparatory steps for the final dictionary
    for key in newHaps.keys():
        keys.append(key.split(','))
    newHaps_freq = []
    for i in range(len(mating)):
        for j in range(len(keys)):
            if keys[j][0] in mating[i] and keys[j][1] in mating[i]:
                keys[j][0] = mating[i][0]
                keys[j][1] = mating[i][1]

    keys_new = []
    for i in range(len(keys)):                                  # remove duplicates
        if keys[i] not in keys_new:
            keys_new.append(keys[i])

    for i in range(len(keys_new)):
        newHaps_freq.append(round(newHaps[','.join(keys_new[i])], 4))   # diploid frequencies in the next generation
    nextGen = dict(zip(mating,newHaps_freq))
    print(nextGen)

    haps_new = []
    for i in range(len(Pstls)):
        newHap = 0
        for j in range(len(mating)):
            if Pstls[i] in mating[j]:
                newHap += nextGen[mating[j]]
        haps_new.append(round(newHap, 4))
    gen = dict(zip(Pstls, haps_new))
    print(gen)
    gens.append(list(gen.values()))

print(gens)

Tags: innewforlenrangekeysgenfreq

热门问题