am_ld.R

# Num individuals
N = 10e3
# Num markers
M = 10

# Assortative Mating
AM_cor = 0.75
# Minor allele frequency
maf = 0.5
# Heritability
var_genetic = 0.5
var_env = 1 - var_genetic
# Generations to mate for
n_gen = 21

seeds = 50

clr = c("#3182bd33","#3182bd")

par(mfrow=c(2,2))
# Try a range of heritability params
for ( var_genetic in c(0.2,0.8) ) {
# Try a range of assortative mating params
for ( AM_cor in c(0.25,0.75) ) {
  plot( 0 , 0 , lwd=3 , type="n" , xlab="Generations", ylab="Mean pairwise LD (r)" , main=paste("Mate Correlation = ",AM_cor,"\nHeritability = ",var_genetic,sep='') , las=1 , xlim=c(0,n_gen-1) , ylim=c(0,0.05) )
  mus = matrix(NA,nrow=n_gen,ncol=seeds)
  
  for ( s in 1:seeds ) {
    # ---
    # sample per-variant effect sizes (adjusted for 2*p*q SNP variance)
    betas = rnorm( M , 0 , sqrt(var_genetic/(M*2*maf*(1-maf)) ) )
    
    # sample 0/1 haplotypes
    mate1_hap_a = matrix( rbinom(N*M,1,maf) , nrow=N , ncol=M )
    mate1_hap_b = matrix( rbinom(N*M,1,maf) , nrow=N , ncol=M )
    
    # sample a random mate
    mate2_hap_a = matrix( rbinom(N*M,1,maf) , nrow=N , ncol=M )
    mate2_hap_b = matrix( rbinom(N*M,1,maf) , nrow=N , ncol=M )
    
    # compute genetic value for direct phenotype
    gv_mate1 = ( mate1_hap_a + mate1_hap_b ) %*% betas
    gv_mate2 = ( mate2_hap_a + mate2_hap_b ) %*% betas
    genetic_mean = mean(c(gv_mate1,gv_mate2))
    
    for ( gens in 1:n_gen ) {
    
      genos = scale(mate1_hap_a + mate1_hap_b)
      genos[ , betas < 0 ] = -1*genos[ , betas < 0 ]
      ldmat = (t(genos) %*% genos / (N-1))
      corrs = ldmat[lower.tri(ldmat,diag=F)]
      mus[gens,s] = mean(corrs,na.rm=T)
      
      # compute genetic value for direct phenotype
      gv_mate1 = ( mate1_hap_a + mate1_hap_b ) %*% betas - genetic_mean
      gv_mate2 = ( mate2_hap_a + mate2_hap_b ) %*% betas - genetic_mean
      pheno_mate1 = gv_mate1 + rnorm(N,0,sqrt(var_env))
      pheno_mate2 = gv_mate2 + rnorm(N,0,sqrt(var_env))
      
      # rank the phenotypes and find similar ranks
      if ( AM_cor != 0 ) {
        ord1 = order(pheno_mate1)
        ord2 = order( AM_cor * scale(pheno_mate2) + rnorm(N,0,sqrt(1-AM_cor^2)))
      } else {
        ord1 = sample(1:N)
        ord2 = sample(1:N)
      }
      # reorder to induce assortment
      mate1_hap_a = mate1_hap_a[ord1,]
      mate1_hap_b = mate1_hap_b[ord1,]
      mate2_hap_a = mate2_hap_a[ord2,]
      mate2_hap_b = mate2_hap_b[ord2,]
      pheno_mate1 = pheno_mate1[ord1]
      pheno_mate2 = pheno_mate2[ord2]
      
      ## cat( "AM_corr", cor( pheno_mate1 , pheno_mate2 , use="complete" ) , '\n' , sep='\t' )
      
      # --- mate and generate kids
      
      child1_hap_a = mate1_hap_a
      sampled_vars = matrix( as.logical(rbinom(N*M,1,0.5)) , nrow=N , ncol=M )
      child1_hap_a[ !sampled_vars ] = mate1_hap_b[ !sampled_vars ]
      
      child1_hap_b = mate2_hap_a
      sampled_vars = matrix( as.logical(rbinom(N*M,1,0.5)) , nrow=N , ncol=M )
      child1_hap_b[ !sampled_vars ] = mate2_hap_b[ !sampled_vars ]
    
      # --- find new mate and generate kids
      if ( AM_cor != 0 ) {
        ord1 = order(pheno_mate1)
        ord2 = order( AM_cor * scale(pheno_mate2) + rnorm(N,0,sqrt(1-AM_cor^2)))
      } else {
        ord1 = sample(1:N)
        ord2 = sample(1:N)
      }
      
      mate1_hap_a = mate1_hap_a[ord1,]
      mate1_hap_b = mate1_hap_b[ord1,]
      mate2_hap_a = mate2_hap_a[ord2,]
      mate2_hap_b = mate2_hap_b[ord2,]
      pheno_mate1 = pheno_mate1[ord1]
      pheno_mate2 = pheno_mate2[ord2]
    
      child2_hap_a = mate1_hap_a
      sampled_vars = matrix( as.logical(rbinom(N*M,1,0.5)) , nrow=N , ncol=M )
      child2_hap_a[ !sampled_vars ] = mate1_hap_b[ !sampled_vars ]
      
      child2_hap_b = mate2_hap_a
      sampled_vars = matrix( as.logical(rbinom(N*M,1,0.5)) , nrow=N , ncol=M )
      child2_hap_b[ !sampled_vars ] = mate2_hap_b[ !sampled_vars ]
      
      # update to the next generation
      mate1_hap_a = child1_hap_a
      mate1_hap_b = child1_hap_b
      mate2_hap_a = child2_hap_a
      mate2_hap_a = child2_hap_b
    }
    lines( 0:(n_gen-1) , mus[,s] , col=clr[1] )
  }
  lines( 0:(n_gen-1) , apply(mus,1,mean) , col=clr[2],lwd=3 )

}}