Tennessen_SoftSweeps.slim

//Mariana Harris
//mharris94@ucla.edu
//March 2021

//Tennesen et al. demographic model with selection in the European population
//A beneficial mutation is added introduceMut generations ago and the simulation is restarted if the mutation is lost
//We modified the stdpopsim code for Two population out-of-Africa demographic model to include positive selection
//stdpopsim: Adrion et al.(2020) A community-maintained standard library of population genetic models, eLife 2020;9:e54967; doi: https://doi.org/10.7554/eLife.54967

initialize() {
	if (!exists("verbosity"))
		defineConstant("verbosity", 1);
	
	if (exists("slimgui")) {
		defineConstant("burn_in", 0.0);
		defineConstant("K", 5.0);
		defineConstant("introduceMut", 1000);
		defineConstant("selec",0.1);
		defineConstant("sampleSize", 177);
	}
	
	//define constants for simulation
	defineConstant("simID", getSeed());
	defineConstant("Q", 1); // Scaling factor to speed up simulation.
	defineConstant("generation_time", 28);
	defineConstant("mutation_rate", Q * 1.25e-08);
	defineConstant("recomb_rate",5e-09);
	defineConstant("chromosome_length", 5e5);
	defineConstant("file_modern", "tmp/VCF_modern.csv");
	defineConstant("file_mesolithic", "tmp/VCF_mesolithic.csv");
	defineConstant("file_modern_dats", "out/modern_dats.txt");
	defineConstant("file_mesolithic_dats", "out/mesolithic_dats.txt");
	defineConstant("file_250", "tmp/VCF_250.csv");
	defineConstant("file_100", "tmp/VCF_100.csv");
	defineConstant("file_250_dats", "out/s250_dats.txt");
	defineConstant("file_100_dats", "out/s100_dats.txt");
	defineConstant("pop_names", c("AFR", "EUR"));
	
	// Time of epoch boundaries, in years before present.
	// The first epoch spans from INF to _T[0].
	defineConstant("_T", c(148000, 51000, 23000, 5000, 0));
	
	// Population sizes in each epoch.
	_N = array(c(
		// INF:_T[0], _T[0]:_T[1], etc.
		c(7310, 14474, 14474, 14474, 15622), // AFR
		c(0, 0, 1861, 1017, 10150) // EUR
		), c(5, 2));
	defineConstant("num_epochs", length(_T));
	defineConstant("num_populations", ncol(_N));
	
	// Population growth rates for each epoch.
	defineConstant("growth_rates", array(c(
		// INF:_T[0], _T[0]:_T[1], etc.
		c(0.0, 0.0, 0.0, 0.0, 0.0166), // AFR
		c(0.0, 0.0, 0.0, 0.00307, 0.0195) // EUR
		), c(num_epochs, num_populations)));
	no_migration = rep(0, num_populations*num_populations);
	// Migration rates for each epoch.
	// Migrations involving a population with size=0 are ignored.
	// XXX: document what the rows & cols correspond to.
	defineConstant("migration_matrices", array(c(
		
		// INF:_T[0]
		no_migration,
		
		// _T[1]:_T[2]
		no_migration,
		
		// _T[2]:_T[3]
		array(c(
		c(0, 0.00015),
		c(0.00015, 0)
		), c(num_populations, num_populations)),
		
		// _T[3]:_T[4]
		array(c(
		c(0, 2.5e-05),
		c(2.5e-05, 0)
		), c(num_populations, num_populations)),
		
		// _T[4]:_T[5]
		array(c(
		c(0, 2.5e-05),
		c(2.5e-05, 0)
		), c(num_populations, num_populations))
		
		), c(num_populations, num_populations, num_epochs)));
	
	// Population splits, one row for each event.
	defineConstant("subpopulation_splits", array(c(
		// time, newpop, size, oldpop
		c(_T[1], 1, _N[2,1], 0)
		), c(4, 1)));
	
	// Admixture pulses, one row for each pulse.
	defineConstant("admixture_pulses", c());
	
	defineConstant("N", asInteger(_N/Q));
	initializeMutationRate(mutation_rate);
	initializeMutationType("m1", 0.5, "f", 0);
	initializeMutationType("m2", 0.5, "f", selec).color = "green";;  // beneficial mutation
	//m2.convertToSubstitution = F;
	m2.mutationStackPolicy = "f";// ??
	initializeGenomicElementType("g1", m1, 1.0);
	initializeGenomicElement(g1, 0, chromosome_length-1);
	initializeRecombinationRate(recomb_rate);//

}

function (void)err(string$ s) {
	stop("ERROR: " + s);
}

function (void)warn(string$ s) {
	catn("WARNING: " + s);
}

function (void)dbg(string$ s, [integer$ debug_level = 2]) {
	if (verbosity >= debug_level) {
		catn(sim.generation + ": " + s);
	}
}

// Check that sizes aren't dangerously low or zero (e.g. due to scaling).
function (void)check_size(integer$ pop, integer$ size, integer$ g) {
	if (size == 0) {
		err("The population size of p"+pop+" ("+pop_names[pop]+") is zero " +
			"at generation "+g+".");
	} else if (size < 50) {
		warn("p"+pop+" ("+pop_names[pop]+") has only "+size+" individuals " +
			"alive at generation "+g+".");
	}
}

// Return the epoch index for generation g.
function (integer)epoch(integer G, integer $g) {
	for (i in 0:(num_epochs-1)) {
		if (g < G[i]) {
			return i;
		}
	}
	return num_epochs - 1;
}

// Return the population size of pop at generation g.
function (integer)pop_size_at(integer G, integer$ pop, integer$ g) {
	e = epoch(G, g);
	N0 = N[e,pop];
	r = Q * growth_rates[e,pop];
	if (r == 0) {
		N_g = N0;
	} else {
		g_diff = g - G[e-1];
		N_g = asInteger(round(N0*exp(r*g_diff)));
	}
	return N_g;
}

// Return the number of generations that separate t0 and t1.
function (integer)gdiff(numeric$ t0, numeric t1) {
	return asInteger(round((t0-t1)/generation_time/Q));
}
// Add mutation halph way through the chromosome
function (void)addMut(void) {
	// save the state of the simulation
	sim.outputFull("/tmp/slim_" + simID + ".txt");
	
	target = sample(p1.genomes, 1); // sample from european population
	
	target.addNewDrawnMutation(m2, asInteger(chromosome_length/2)-1);
}
function (void)checkLoss(void) {
	counts = p1.genomes.countOfMutationsOfType(m2);
	freq = mean(counts > 0);
	muts=size(sim.mutationsOfType(m2))+sum(sim.substitutions.mutationType==m2);
	//if ((freq==0 & sum(sim.substitutions.mutationType==m2)==0) | (muts==1 & sum(sim.substitutions.mutationType==m2)==0)){
	if ((freq==0 & sum(sim.substitutions.mutationType==m2)==0)){
		// go back to generation
		sim.outputFull("/tmp/slim_" + simID + ".txt");
		
		// start a newly seeded run
		setSeed(rdunif(1, 0, asInteger(2^62) - 1));
		
		// re-introduce the sweep mutation
		for(i in 1:K){
			target = sample(p1.genomes, 1);
			target.addNewDrawnMutation(m2, asInteger(chromosome_length/2)-1);
		}
	
	
	}

}


// Output VCF file (modern sample) and end simulation.
function (void)outModern(integer  size) {
	muts=size(sim.mutationsOfType(m2))+sum(sim.substitutions.mutationType==m2);
	line = asString(muts) + "\t" + asString(K)+ "\t" +asString(selec);
	writeFile(file_modern_dats,line,append=T);
	p1.outputVCFSample(size,filePath=file_modern);
	sim.simulationFinished();
}

function (void)outMesolithic(integer size) {
	p1.outputVCFSample(size,filePath=file_mesolithic);
}

function (void)outputSample(integer size,string$ path, string$ path_dats) {
	muts=size(sim.mutationsOfType(m2))+sum(sim.substitutions.mutationType==m2);
	line = asString(muts) + "\t" + asString(K)+ "\t" +asString(selec);
	writeFile(path_dats,line,append=T);
	p1.outputVCFSample(size,filePath=path);
}



// Add mutation half way through the chromosome
function (void)addRecurrentMut(void) {
	num_muts=rpois(1,THETA_A/2);
	if(num_muts>0){
		for(i in 1:num_muts){
			target = sample(p1.genomes, 1); // sample from european population
			target.addNewDrawnMutation(m2, asInteger(chromosome_length/2)-1);
		}
	}
}

function (void)addSGVMut(void) {
	sim.outputFull("/tmp/slim_" + simID + ".txt");
	for(i in 1:K){
		target = sample(p1.genomes, 1);
		target.addNewDrawnMutation(m2, asInteger(chromosome_length/2)-1);
	}

}

1 {
	/*
     * Create initial populations and migration rates.
     */
	
	// Initial populations.
	for (i in 0:(num_populations-1)) {
		if (N[0,i] > 0) {
			check_size(i, N[0,i], sim.generation);
			dbg("sim.addSubpop("+i+", "+N[0,i]+");");
			sim.addSubpop(i, N[0,i]);
		}
	}
	
	if (length(sim.subpopulations) == 0) {
		err("No populations with non-zero size in generation 1.");
	}
	
	// Initial migration rates.
	i = 0;
	for (j in 0:(num_populations-1)) {
		for (k in 0:(num_populations-1)) {
			if (j==k | N[i,j] == 0 | N[i,k] == 0) {
				next;
			}
			
			m = Q * migration_matrices[k,j,i];
			p = sim.subpopulations[j];
			dbg("p"+j+".setMigrationRates("+k+", "+m+");");
			p.setMigrationRates(k, m);
		}
	}
	
	
	// The end of the burn-in is the starting generation, and corresponds to
	// time T_start. All remaining events are relative to this generation.
	N_max = max(N[0,0:(num_populations-1)]);
	G_start = sim.generation + asInteger(round(burn_in * N_max));
	T_start = max(_T);
	G = G_start + gdiff(T_start, _T);
	G_end = max(G);
	
	
	/*
     * Register events occurring at time T_start or more recently.
     */
	
	// Split events.
	if (length(subpopulation_splits) > 0 ) {
		for (i in 0:(ncol(subpopulation_splits)-1)) {
			g = G_start + gdiff(T_start, subpopulation_splits[0,i]);
			newpop = drop(subpopulation_splits[1,i]);
			size = asInteger(subpopulation_splits[2,i] / Q);
			oldpop = subpopulation_splits[3,i];
			check_size(newpop, size, g);
			sim.registerLateEvent(NULL,
				"{dbg(self.source); " +
				"sim.addSubpopSplit("+newpop+","+size+","+oldpop+");}",
				g, g);
		}
	}
	
	// Population size changes.
	if (num_epochs > 1) {
		for (i in 1:(num_epochs-1)) {
			g = G[i-1];
			for (j in 0:(num_populations-1)) {
				// Change population size if this hasn't already been taken
				// care of by sim.addSubpop() or sim.addSubpopSplit().
				if (N[i,j] != N[i-1,j] & N[i-1,j] != 0) {
					check_size(j, N[i,j], g);
					sim.registerLateEvent(NULL,
						"{dbg(self.source); " +
						"p"+j+".setSubpopulationSize("+N[i,j]+");}",
						g, g);
				}
				
				if (growth_rates[i,j] != 0) {
					growth_phase_start = g+1;
					if (i == num_epochs-1) {
						growth_phase_end = G[i];
					} else {
						// We already registered a size change at generation G[i].
						growth_phase_end = G[i] - 1;
					}
					
					if (growth_phase_start >= growth_phase_end) {
						// Some demographic models have duplicate epoch times,
						// which should be ignored.
						next;
					}
					
					N_growth_phase_end = pop_size_at(G, j, growth_phase_end);
					check_size(j, N_growth_phase_end, growth_phase_end);
					
					N0 = N[i,j];
					r = Q * growth_rates[i,j];
					sim.registerLateEvent(NULL,
						"{" +
						"dbg(self.source); " +
						"gx=sim.generation-"+g+"; " +
						"size=asInteger(round("+N0+"*exp("+r+"*gx))); " +
						"p"+j+".setSubpopulationSize(size);" +
						"}",
						growth_phase_start, growth_phase_end);
				}
			}
		}
		
		// Migration rates.
		for (i in 1:(num_epochs-1)) {
			for (j in 0:(num_populations-1)) {
				for (k in 0:(num_populations-1)) {
					if (j==k | N[i,j] == 0 | N[i,k] == 0) {
						next;
					}
					
					m_last = Q * migration_matrices[k,j,i-1];
					m = Q * migration_matrices[k,j,i];
					if (m == m_last) {
						// Do nothing if the migration rate hasn't changed.
						next;
					}
					g = G[i-1];
					sim.registerLateEvent(NULL,
						"{dbg(self.source); " +
						"p"+j+".setMigrationRates("+k+", "+m+");}",
						g, g);
				}
			}
		}
	}
	
	// Admixture pulses.
	if (length(admixture_pulses) > 0 ) {
		for (i in 0:(ncol(admixture_pulses)-1)) {
			g = G_start + gdiff(T_start, admixture_pulses[0,i]);
			dest = admixture_pulses[1,i];
			src = admixture_pulses[2,i];
			rate = admixture_pulses[3,i];
			sim.registerLateEvent(NULL,
				"{dbg(self.source); " +
				"p"+dest+".setMigrationRates("+src+", "+rate+");}",
				g, g);
			sim.registerLateEvent(NULL,
				"{dbg(self.source); " +
				"p"+dest+".setMigrationRates("+src+", 0);}",
				g+1, g+1);
		}
	}
	
	
	// start recurrent muts
	genSelec = asInteger(G_end-introduceMut);
	sim.registerLateEvent(NULL, "{addSGVMut();}", genSelec, genSelec);
	
	sim.registerLateEvent(NULL, "{checkLoss();}", genSelec, G_end);
	
	
	
	
	// output VCF at gen G_end-400 (Mesolithic Samples)
	sim.registerLateEvent(NULL, "{dbg(self.source); outputSample(sampleSize,file_mesolithic,file_mesolithic_dats);}", G_end-400, G_end-400);
	
	// output VCF at gen G_end-250
	sim.registerLateEvent(NULL, "{dbg(self.source); outputSample(sampleSize,file_250,file_250_dats);}", G_end-250, G_end-250);
	
	// output VCF at gen G_end-100 
	sim.registerLateEvent(NULL, "{dbg(self.source); outputSample(sampleSize,file_100,file_100_dats);}", G_end-100, G_end-100);
	
	// output VCF at gen G_end-40 (Modern Samples)
	sim.registerLateEvent(NULL, "{dbg(self.source); outModern(sampleSize);}", G_end-40, G_end-40);
	
	
	
	
	if (G_start > sim.generation) {
		dbg("Starting burn-in...");
	}

}