shared_multi_haploid.py

'''
sample_diploid_recombinant_selection_expectation()
10000 experiments
segment count, mean improvement over mean, standard deviation of improvement
1 3.74 0.43
2 5.28 0.42
3 6.03 0.42
4 6.46 0.42
5 6.75 0.41
6 6.97 0.41
7 7.14 0.42
8 7.28 0.41
9 7.39 0.41
10 7.47 0.41
'''

from genome import *


def sample_diploid_recombinant_selection_expectation():
   for segment_count in range(1, 11):
      values = []
      for x in range(100):
      #for x in range(10000):
         couple = Couple(segment_count = segment_count)
         values.append(sum(couple.parents[parent].chromosome_pairs[chr].max_recombination for parent in [0,1] for chr in range(1,24)) - couple.embryo_mean)
      print(segment_count, round(np.mean(values),2), round(np.std(values),2))

SEGMENT_COUNT=5 

def sample_diploid_replacement(external_donors = 0, max_replacements = 0):
   couple = Couple(segment_count = SEGMENT_COUNT)
   total_sum = 0
   for parent in [0,1]:
      max_recombinations = [couple.parents[parent].chromosome_pairs[chr].max_recombination for chr in range(1,24)] 

      improvement = 0
      if external_donors>0:
         donors = [DiploidGenome(segment_count=SEGMENT_COUNT) for _ in range(external_donors)]
         donor_maxes = [max([donor.chromosome_pairs[index].max_recombination for donor in donors]) for index in range(1,24)]

         replace_improvements = [donor_maxes[i] - max_recombinations[i] for i in range(23)]
         replace_improvements.sort(reverse=True )
         improvements_to_make = replace_improvements[:max_replacements]
         if improvements_to_make[-1]<0:
            print(f'did not replace all {max_replacements} with {external_donors} external donors')
         improvement = sum(max(x,0) for x in improvements_to_make)

      total_sum += sum(max_recombinations) + improvement
   return total_sum - couple.embryo_mean


def sample_diploid_replacement_expectations(donors=0, k=0, samples = 10):
   values = []
   for x in range(samples):
      values.append(sample_diploid_replacement(external_donors=donors, max_replacements=k))
   return (donors,  k, np.mean(values), np.std(values))

def get_replacement_expectations(donor_list, k_list, samples):
   results = {}
   for donors in donor_list:
      results[donors] = {k:sample_diploid_replacement_expectations(samples = samples, donors = donors, k=k)
                        for k in k_list}
      print('got samples for ', donors, 'donors')
   return results


donor_list = [0,1,2,4, 10, 23]
k_list = list(range(1, 8)) + [10, 12]
#SAMPLES = 500 
SAMPLES = 10 
results = get_replacement_expectations(donor_list, k_list, SAMPLES)

data = [ [round(results[donors][k][2], 1)
         for donors in donor_list]
        for k in k_list]

# then you can use the data in matplotlib, e.g. something like 
#ax.table(cellText=data,colLabels=donor_list,rowLabels=k_list, loc='center')