from __future__ import division
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
%matplotlib inline
#Total genetic variance
def var_g(p):
return pow(p * (1 - p) * (2 - p), 2) + p * pow(1-p, 4) * (2 - p)
#Additive genetic variance
def var_a(p):
return 2 * p * pow(1 - p, 3)
#Calculate allele prevalence from the fraction starred in the population
def prev_from_frac_pop(pop_frac):
return 1 - np.sqrt(1 - pop_frac)
#Look at variance with changing population starred
pvals = np.arange(0, 1, 0.01)
plt.plot(pvals, var_g(prev_from_frac_pop(pvals)), 'k')
plt.plot(pvals, var_a(prev_from_frac_pop(pvals)), 'b')
plt.figure()
#And look at the narrow-sense heritability versus population starred
plt.plot(pvals, var_a(prev_from_frac_pop(pvals)) / var_g(prev_from_frac_pop(pvals)), 'k')
prev = prev_from_frac_pop(0.5)
print prev
print var_a(prev) / var_g(prev)
#Note these are fixed - don't change without adjusting the fractions below!
pop_frac = 0.5
p = prev_from_frac_pop(0.5) #positive allele prevalence
#Go ahead and change these to adjust the simulation
h2i = 0.0 #"true" narrow-sense heritability of income-havingness
d = 2.0 #magnitude of discrimination
total_pop = 50000 #Size of population to simulate
#Data for plain-plain couplings
means = np.array([-d, -d])
covariance_matrix = np.array([[1, h2i], [h2i, 1]])
data_plain_plain = np.random.multivariate_normal(means, covariance_matrix, size=total_pop // 4)
plt.scatter(data_plain_plain[:,0], data_plain_plain[:,1], lw=0, s=1)
print np.corrcoef(data_plain_plain.transpose())
#Data for plain-star couplings (1/4 of offspring plain, remaining starred)
means = np.array([-d/2, 0]) #We'll adjust 1/4 of the offspring downward after
covariance_matrix = np.array([[1, h2i], [h2i, 1]])
data_plain_star = np.random.multivariate_normal(means, covariance_matrix, size=total_pop // 2)
data_plain_star[:total_pop // 8, 1] -= d #1/4 of the offspring are discriminated against
plt.scatter(data_plain_star[:,0], data_plain_star[:,1], lw=0, s=1)
print np.corrcoef(data_plain_star.transpose())
#Data for star-star couplings (only 1/16 of offspring plain, remaining starred)
means = np.array([0, 0])
covariance_matrix = np.array([[1, h2i], [h2i, 1]])
data_star_star = np.random.multivariate_normal(means, covariance_matrix, size=total_pop // 4)
data_star_star[:total_pop // 64, 1] -= d #1/16 of the offspring are discriminated against
plt.scatter(data_star_star[:,0], data_star_star[:,1], lw=0, s=1)
print np.corrcoef(data_star_star.transpose())
#Merge the data and get the overall regression
data_all = np.vstack([data_plain_plain, data_plain_star, data_star_star])
plt.scatter(data_all[:,0], data_all[:,1], lw=0, s=1)
print np.corrcoef(data_all.transpose())