WF_multilocus_cluster.slim

initialize(){
	
	// csv constants	
	defineCfgParam("make_csv", 0); // whether to create a csv file summarizing each generation
	defineCfgParam("path", paste0(getwd(), "/")); 	// directory for data storage (folder must be created in advance)	
	
	// general constants
	defineCfgParam("population_size", 1e4);
	defineCfgParam("release_size", 200); // represents a 1% drive allele frequency
	defineCfgParam("recomb_rate", 0.5); // between target sites only (drive locus is always unlinked)
	defineCfgParam("shared_cas9", 0); // whether cas9 is shared between the drive site and target sites
	defineCfgParam("baseline_cleavage_rate", 0.8); // baseline target site cleavage rate
	defineCfgParam("saturation_factor", 0.0); // set to 0.0 for no saturation
	
	// constants for the drive locus
	defineCfgParam("resistance_rate", 0.01);  // functional resistance at drive locus
	defineCfgParam("drive_coeff", 0.1);  // selection coefficient of drive allele
	defineCfgParam("drive_dom", 0.5);  // dominance of drive allele
	defineCfgParam("drive_cleavage_rate", 1.0); // drive_cut_rate if shared_cas9 = 0
	
	// constants for distant target loci
	defineCfgParam("num_target_loci", 5);  // either linked or unlinked
	defineCfgParam("func_resist_rate", 0.01);  // functional alleles at target loci	
	defineCfgParam("broken_coeff", 0.43);  // fitness cost of each disrupted target
	defineCfgParam("sd_broken_coeff", 0);  // sd of rnorm for fitness cost
	defineCfgParam("broken_dom", 0);  // dominance of disrupted allele at a target site
	
	// constants for tracking output
	defineCfgParam("recovery_val", 0.2); // genetic load signifying recovery
	defineCfgParam("end_simulation_at_genetic_load_cutoff", 0); // whether to end the simulation when the genetic load falls below this
	defineConstant("PLOT_HOMZ_COUNTS", F); // whether to liveplot the count of homozygous disrupted target sites each generation
	defineConstant("PRINT_COMMON_HAPLOTYPES", F); // prints the most common haplotypes every generation 
	defineConstant("NUM_HAPLOTYPES_PRINT", 10); // when PRINT_COMMON_HAPLOTYPES = T, only print the NUM_HAPLOTYPES_PRINT most common haplotypes
	
	// constants based on other defined constants
	defineConstant("cleavage_rate", ifelse(shared_cas9==1, baseline_cleavage_rate/((num_target_loci + 1)^(saturation_factor)), baseline_cleavage_rate/((num_target_loci)^(saturation_factor))));
	defineConstant("disruption_rate", cleavage_rate);  // creating non-functional alleles at target loci
	defineConstant("drive_cut_rate", ifelse(shared_cas9==1, cleavage_rate, drive_cleavage_rate)); // cutting the wild-type allele at drive locus
	
	defineGlobal("genetic_load_below_val", F);
	defineGlobal("genetic_load_rose_above_val", F);
	defineGlobal("time_to_recovery", -1);
	defineGlobal("genetic_load_track", c());
	
	
	catn("==============================");
	catn(paste("TARGET_CUT_RATE::", disruption_rate, "DRIVE_CUT_RATE::", drive_cut_rate));
	catn("==============================");
	
	catn(paste("SEED_NUMBER::", getSeed()));
	
	// setup mutation types
	initializeMutationRate(0);  // no outside mutations
	wt = initializeMutationType("m1", 1, "f", 0);  // wild-type allele
	d = initializeMutationType("m2", drive_dom, "f", -1*drive_coeff);  // gene drive
	r = initializeMutationType("m3", 1, "f", 0);  // resistance allele
	wt2 = initializeMutationType("m4", 1, "f", 0);  // wild-type at other loci
	b = initializeMutationType("m5", 1, "f", 0);  // broken allele
	f = initializeMutationType("m6", 1, "f", 0);  // functional allele
	
	// setup the genome
	initializeGenomicElementType("g1", c(wt, d, r), c(1, 1, 1));
	initializeGenomicElementType("g2", c(wt2, b, f), c(1, 1, 1));
	initializeGenomicElement(g1, 0, 0);
	if (num_target_loci > 0) {initializeGenomicElement(g2, 1, num_target_loci);}
	if (num_target_loci > 1) {initializeRecombinationRate(c(0.5, recomb_rate), c(1, num_target_loci));}
	else {initializeRecombinationRate(0.5);}  // 0.5 = unlinked loci
	
	// do not allow mutation stacking or substitution
	variations = c(wt, d, r, wt2, b, f);
	variations.mutationStackGroup = 1;
	variations.mutationStackPolicy = "l";
	variations.convertToSubstitution = F;
	
	// setup the output file
	if (make_csv){
		if (num_target_loci > 0){
			// gen_marker = 1 when it's the generation when the genetic load < 20%
			file = "gen_num, mean_fit, genetic_load, wt_rate, drive_rate, resist_rate, wt2_rate, broken_rate, func_rate, gen_marker";
			
			// Also record the count of individuals with 0 homz disrupted, 1 homz disrupted,..., all homz disrupted sites
			possible_homz_sites = 0:num_target_loci;
			headers = c();
			for (count in possible_homz_sites){
				col = paste0("frac_inds_with_", count,"_homz_disrupted");
				headers = c(headers, col);
			}
			
			file = file + "," + paste(headers, sep = ", ") + "," + "mean_num_disrupted_on_haplotype";
		
		} else {
			file = "gen_num, mean_fit, genetic_load, wt_rate, drive_rate, resist_rate";
		}
		defineGlobal("output", file);
		
		// determine the file name
		date_time = paste0(c(rev(strsplit(date(), sep="-")), strsplit(time(), sep=":")));
		seed = getSeed();
		defineConstant("file_name", path + date_time + "_" + seed + "_output.csv");
		catn("FILENAME = " + file_name);
		
		// Create file here and then append to it
		if (!writeFile(file_name, output)){stop("Error initiating output file.");}
		
		// create a parameters file
		file_text = "population_size: " + population_size + "\nrelease_size: " + release_size + "\nrecomb_rate: " + recomb_rate + "\nshared_cas9: " + shared_cas9 + "\nbaseline_cleavage_rate: " + baseline_cleavage_rate + "\nsaturation_factor: " + saturation_factor + "\nresistance_rate: " + resistance_rate + "\ndrive_coeff: " + drive_coeff + "\ndrive_dom: " + drive_dom + "\ndrive_cleavage_rate: " + drive_cleavage_rate + "\nnum_target_loci: " + num_target_loci + "\nfunc_resist_rate: " + func_resist_rate + "\nbroken_coeff: " + broken_coeff + "\nsd_broken_coeff: " + sd_broken_coeff + "\nbroken_dom: " + broken_dom + "\nrecovery_val: " + recovery_val + "\ncleavage_rate: " + cleavage_rate + "\ndisruption_rate: " + disruption_rate + "\ndrive_cut_rate: " + drive_cut_rate;
		if (!writeFile(path + date_time + "_" + seed + "_parameters.txt", file_text)){
			stop("Error creating parameters file.");}
	}
}

function (void) defineCfgParam(string$ name, lifs value) {
	
	if (!exists(name))
		defineConstant(name, value);
}

function (string)haplotypeString(o<Haplosome>$ genome){
	// Converts an individual's haplosome to a string
	
	// drive site key:
	// d_ = drive
	// w_ = wild-type at drive locus
	// r_ = resistant at drive locus
	
	// target sites key:
	// b = disrupted
	// w = wild-type (undisrupted)
	// r = resistant
	
	
	contains_drive = genome.countOfMutationsOfType(m2) > 0;
	contains_wt =  genome.countOfMutationsOfType(m1) > 0;
	contains_r = genome.countOfMutationsOfType(m3) > 0;
	if (contains_drive)
		d_letter = "d_";
	else if (contains_wt)
		d_letter = "w_";
	else if (contains_r)
		d_letter = "r_";
	
	pos_disrupted = genome.positionsOfMutationsOfType(m5);
	pos_wildtype = genome.positionsOfMutationsOfType(m4);
	pos_r1 = genome.positionsOfMutationsOfType(m6);
	parallel_disrupted = rep("b", size(pos_disrupted));
	parallel_wildtype = rep("w", size(pos_wildtype));
	parallel_r1 = rep("r", size(pos_r1));
	
	combined_dis_wt_r1 = c(pos_disrupted, pos_wildtype, pos_r1);
	parallel_dis_wt_r1 = c(parallel_disrupted, parallel_wildtype, parallel_r1);
	
	ordering = order(combined_dis_wt_r1);
	
	haplotype_order = parallel_dis_wt_r1[ordering];
	
	haplotype = paste0(haplotype_order);
	haplotype = paste0(d_letter, haplotype);
	
	return haplotype;
}


function (void)reportHaplotypeCounts([integer max_num_report = -1]){
	// Lists the max_num_report most common haplotypes and their counts
	// if max_num_report is -1, reports ALL of them
	// see haplotypeString for key
	
	all = p1.individuals;
	haplos = c();
	
	for (ind in all){
		haplo_str1 = haplotypeString(ind.haploidGenome1);
		haplo_str2 = haplotypeString(ind.haploidGenome2);
		
		haplos = c(haplos, haplo_str1, haplo_str2);
	}
	
	unique_haplotypes = unique(haplos);
	num_unique_haplotypes = size(unique_haplotypes);
	haplotype_counts = sapply(unique_haplotypes, "sum(haplos == applyValue);");
	
	ordering = order(haplotype_counts, ascending = F); // greatest to least
	
	common_haplotypes = unique_haplotypes[ordering];
	common_haplotypes_counts = haplotype_counts[ordering];
	
	if ((max_num_report != -1) & (num_unique_haplotypes > max_num_report)){
		common_haplotypes = common_haplotypes[seqLen(max_num_report)];
		common_haplotypes_counts = common_haplotypes_counts[seqLen(max_num_report)];
	}
	
	catn("======================================");
	line = "HAPLOTYPES:\tCOUNT\tFREQ::\n";
	for (i in seqLen(size(common_haplotypes))){
		line = line + common_haplotypes[i] + ":\t" + asString(common_haplotypes_counts[i]) + "\t" + asString(common_haplotypes_counts[i]/(2*population_size)) + "\n";
	}
	cat(line);
	catn("======================================");
}


100: modifyChild(){

	// Determine if parent had drive
	parent1_wt = sum(parent1.haplosomes.countOfMutationsOfType(m1)) == 1;
	parent1_d = sum(parent1.haplosomes.countOfMutationsOfType(m2)) == 1;
	parent2_wt = sum(parent2.haplosomes.countOfMutationsOfType(m1)) == 1;
	parent2_d = sum(parent2.haplosomes.countOfMutationsOfType(m2)) == 1;
	
	// Look at all wild-type alleles in the child
	child1_wt = child.haploidGenome1.countOfMutationsOfType(m1) == 1;
	child2_wt = child.haploidGenome2.countOfMutationsOfType(m1) == 1;
	
	// Child is wild-type at drive locus and parent had drive: simulate homing or resistance
	if (child1_wt & parent1_wt & parent1_d){
		// cleavage
		if (runif(1) < drive_cut_rate){
			// cut becomes a resistance allele
			if (runif(1) < resistance_rate){
				child.haploidGenome1.addNewDrawnMutation(m3, 0);}  // new m3 mutation
			else {
				// HDR to introduce the drive allele
				child.haploidGenome1.addMutations(DRIVE);}}}
	
	if (child2_wt & parent2_wt & parent2_d){
		// cleavage
		if (runif(1) < drive_cut_rate){
			// cut becomes a resistance allele
			if (runif(1) < resistance_rate){
				child.haploidGenome2.addNewDrawnMutation(m3, 0);}  // new m3 mutation
			else {
				// HDR to introduce the drive allele
				child.haploidGenome2.addMutations(DRIVE);}}}
	
	// Simulate possible cleavage at the targets
	if (num_target_loci > 0){
		// if parent 1 has at least one drive allele, it can disrupt the target loci
		if (sum(parent1.haplosomes.countOfMutationsOfType(m2)) > 0){
			// chance of converting some other wt2 alleles in genome 1
			disrupt_chance = (runif(num_target_loci) < disruption_rate);
			resist_chance = (runif(num_target_loci) < func_resist_rate);
			
			// positions of the wild type 2 alleles
			wt2_pos = child.haploidGenome1.positionsOfMutationsOfType(m4);
			wt2_pos = (match(1:num_target_loci, wt2_pos) >= 0);
			
			// potential allele conversions positions
			disrupt_convert = (disrupt_chance & !resist_chance & wt2_pos);
			resist_convert = (disrupt_chance & resist_chance & wt2_pos);
			
			// make the conversions
			if (sum(disrupt_convert)){
				child.haploidGenome1.addNewDrawnMutation(m5, (1:num_target_loci)[disrupt_convert]);}
			if (sum(resist_convert)){
				child.haploidGenome1.addNewDrawnMutation(m6, (1:num_target_loci)[resist_convert]);}
		}
		
		// if parent 2 has at least one drive allele, it can disrupt the target loci		
		if (sum(parent2.haplosomes.countOfMutationsOfType(m2)) > 0){
			// chance of converting some other wt2 alleles in genome 2
			disrupt_chance = (runif(num_target_loci) < disruption_rate);
			resist_chance = (runif(num_target_loci) < func_resist_rate);
			
			// positions of the wild type 2 alleles
			wt2_pos = child.haploidGenome2.positionsOfMutationsOfType(m4);
			wt2_pos = (match(1:num_target_loci, wt2_pos) >= 0);
			
			// potential allele conversions positions
			disrupt_convert = (disrupt_chance & !resist_chance & wt2_pos);
			resist_convert = (disrupt_chance & resist_chance & wt2_pos);
			
			// make the conversions
			if (sum(disrupt_convert)){
				child.haploidGenome2.addNewDrawnMutation(m5, (1:num_target_loci)[disrupt_convert]);}
			if (sum(resist_convert)){
				child.haploidGenome2.addNewDrawnMutation(m6, (1:num_target_loci)[resist_convert]);}
		}
	}
	return T;
}

100: fitnessEffect(){

	// first make a fitness placeholder with no effect
	ind_fitness = 1.0;
	
	// count homozygous and heterozygous broken alleles
	if (num_target_loci > 0){
		positions1 = individual.haploidGenome1.positionsOfMutationsOfType(m5);
		positions2 = individual.haploidGenome2.positionsOfMutationsOfType(m5);
		homozygous = size(setIntersection(positions1, positions2));
		heterozygous = size(positions1) + size(positions2) - (2*homozygous);
		
		// create some fitness cost for homozygous deleterious mutations
		if (homozygous){
			fitness_costs = rnorm(homozygous, (1.0 - broken_coeff), sd_broken_coeff);
			fitness_costs[fitness_costs > 1] = 2 - fitness_costs[fitness_costs > 1];  // do not allow beneficial mutations (reflect)
			ind_fitness = ind_fitness*product(abs(fitness_costs));  // do not allow negative fitness either
		}
		
		// create some fitness cost for heterozygous deleterious mutations
		if (heterozygous & broken_dom){
			fitness_costs = rnorm(heterozygous, broken_coeff, sd_broken_coeff);
			fitness_costs = (1.0 - (broken_dom*fitness_costs));
			fitness_costs[fitness_costs > 1] = 2 - fitness_costs[fitness_costs > 1];  // do not allow beneficial mutations (reflect)
			ind_fitness = ind_fitness*product(abs(fitness_costs));  // do not allow negative fitness either
		}
	}
	return ind_fitness;
}

1 early(){
	// create population
	sim.addSubpop("p1", population_size);
}

1 late(){
	// give all individuals the wild-type allele
	p1.haplosomes.addNewDrawnMutation(m1, 0);
	
	// put the second wild-type allele at all the other loci
	if (num_target_loci > 0){
		p1.haplosomes.addNewDrawnMutation(m4, 1:num_target_loci);}
}

100 late(){
	// introduce the gene drive to a random sample of the population in heterozygous state (0.5% frequency)
	target = sample(p1.individuals, release_size);
	mut = target.haploidGenome1.addNewDrawnMutation(m2, 0);
	defineConstant("DRIVE", mut);
}


101: early(){
	
	// store the fitness value
	defineGlobal("fitness", mean(p1.cachedFitness(NULL)));
	defineGlobal("genetic_load", 1 - fitness);
	if (make_csv){
		defineGlobal("new_line", paste(c(sim.cycle - 100, fitness, genetic_load), sep=", "));
	}
	
	
	// Determine if the genetic load fell below 20% (after rising above this initially)
	defineGlobal("genetic_load_track", c(genetic_load_track, genetic_load));
	
	gen_marker = 0;
	if (genetic_load >= recovery_val & !genetic_load_rose_above_val){
		catn("========================================");
		catn(paste("genetic_load_above_val_AT_GEN:", sim.cycle - 100));
		catn("========================================");
		defineGlobal("genetic_load_rose_above_val", T);
	}
	
	if (genetic_load < recovery_val & genetic_load_rose_above_val & !genetic_load_below_val){
		defineGlobal("time_to_recovery", sim.cycle - 100);
		catn("========================================");
		catn(paste("GENETIC_LOAD_BELOW_VAL_AT_GEN:", time_to_recovery));
		catn("========================================");
		defineGlobal("genetic_load_below_val", T);
		gen_marker = 1;
	}
	
	// calculate frequencies of each mutation type
	wild = sum(p1.individuals.countOfMutationsOfType(m1))/(2*population_size);
	drive = sum(p1.individuals.countOfMutationsOfType(m2))/(2*population_size);
	resist = sum(p1.individuals.countOfMutationsOfType(m3))/(2*population_size);
	
	// contribute to output file
	if (make_csv){
		defineGlobal("new_line", paste(c(new_line, wild, drive, resist), sep=", "));}
	
	// same thing for mutations at other loci
	wild2 = 0.0;
	broken = 0.0;
	func = 0.0;
	all = p1.individuals;
	
	if (num_target_loci > 0){
		wild2 = sum(p1.individuals.haplosomes.countOfMutationsOfType(m4))/(2*population_size*num_target_loci);
		broken = sum(p1.individuals.haplosomes.countOfMutationsOfType(m5))/(2*population_size*num_target_loci);
		func = sum(p1.individuals.haplosomes.countOfMutationsOfType(m6))/(2*population_size*num_target_loci);
		
		possible = num_target_loci + 1;
		m_counts = matrix(rep(0.0, 3*possible), nrow = 3, ncol = possible);
		m_counts[0,] = 0:num_target_loci;
		
		// contribute to output file
		if (make_csv){
			defineGlobal("new_line", paste(c(new_line, wild2, broken, func, gen_marker), sep=", "));
		}
		
		// count the average number of homzygous target alleles
		homz_target_count = c();
		num_disrupted_on_haplotype = c();
		
		for (ind in all){
			positions1 = ind.haploidGenome1.positionsOfMutationsOfType(m5);
			positions2 = ind.haploidGenome2.positionsOfMutationsOfType(m5);
			homozygous = size(setIntersection(positions1, positions2));
			homz_target_count = c(homz_target_count, homozygous);
			
			num_disrupted_genome1 = size(positions1);
			num_disrupted_genome2 = size(positions2);
			num_disrupted_on_haplotype = c(num_disrupted_on_haplotype, num_disrupted_genome1, num_disrupted_genome2);
		}
		
		mean_num_disrupted_on_haplotype = mean(num_disrupted_on_haplotype);
		avg_homz_target_count = mean(homz_target_count);
		uniqueVals = unique(homz_target_count);
		possible_n = 0:num_target_loci;
		counts_of_each_n = sapply(possible_n, "sum(homz_target_count == applyValue);");
		m_counts[1,] = counts_of_each_n;
		n = sum(counts_of_each_n);
		if (n > 0)
			m_counts[2,] = m_counts[1,]/n;
		
		
		if (make_csv){
			defineGlobal("new_line", paste(c(new_line, m_counts[2,], mean_num_disrupted_on_haplotype), sep=", "));
		}
	
	} else {
		avg_n_drive_carriers = 0;
		avg_n_nondrive_carriers = 0;
		avg_homz_target_count = 0;
		mean_num_disrupted_on_haplotype = 0;
	}
	
	gen = sim.cycle - 100;
	
	// print summary to console
	catn(paste("GEN::", gen,"MEAN_FITNESS:", fitness,"DRIVE_RATE:", drive, "WT_AT_DRIVE_RATE:",wild, "R1_AT_DRIVE_RATE:", resist,"DISRUPTED_RATE:", broken, "WT_AT_DISRUPTED_RATE:",wild2,"R1_AT_DISRUPTED_RATE:", func,"GENETIC_LOAD:", genetic_load, "AVG_HOMZ_TARGET_COUNT:", avg_homz_target_count, "AVG_NUM_DISRUPTED_ON_HAPLOSOME:", mean_num_disrupted_on_haplotype));
	
	
	if (num_target_loci > 0){
		
		// Additional output	
		//catn("NUM_HOMZ_SITES_AND_COUNT::");
		//print(m_counts);
		
		if (PRINT_COMMON_HAPLOTYPES)
			reportHaplotypeCounts(NUM_HAPLOTYPES_PRINT);
		
		
		if (PLOT_HOMZ_COUNTS){
			plot = slimgui.createPlot(paste0("gen: ", gen,"\ngenetic load = ", 1 - fitness), xrange = c(0, possible), yrange = c(0, p1.individualCount), xlab = "Number of homozygous targets", ylab = "Count of individuals");
			plot.points(m_counts[0,], m_counts[1,], 16, size = 2.0);
		}
	}
	
	// append to output file
	if (make_csv){
		if (!writeFile(file_name, new_line, append = T)){stop("Error appending to output file.");}}
	
	
	// Conditions that end the simulation:
	
	// Drive loss when n = 0
	if (drive == 0){
		catn(paste("DRIVE_LOST_AT_GEN:", gen));
		
		if (num_target_loci == 0){
			ind = whichMax(genetic_load_track);
			max_genetic_load = genetic_load_track[ind];
			catn(paste("ENDING_GENETIC_LOAD:", genetic_load));
			catn(paste("MAXIMUM_GENETIC_LOAD:", max_genetic_load, "AT_GEN:", ind));
			sim.simulationFinished();
		}
	}
	
	// Drive fixation when n = 0
	if (drive == 1){
		catn(paste("DRIVE_FIXED_AT_GEN:", gen));
		
		if (num_target_loci == 0){
			ind = whichMax(genetic_load_track);
			max_genetic_load = genetic_load_track[ind];
			catn(paste("ENDING_GENETIC_LOAD:", genetic_load));
			catn(paste("MAXIMUM_GENETIC_LOAD:", max_genetic_load, "AT_GEN:", ind));
			sim.simulationFinished();
		}
	}
	
	// Genetic load < 20% when end_simulation_at_genetic_load_cutoff
	if (end_simulation_at_genetic_load_cutoff & genetic_load_below_val){
		ind = whichMax(genetic_load_track);
		max_genetic_load = genetic_load_track[ind];
		catn(paste("ENDING_GENETIC_LOAD:", genetic_load));
		catn(paste("MAXIMUM_GENETIC_LOAD:", max_genetic_load, "AT_GEN:", ind));
		sim.simulationFinished();
	}
	
	if (num_target_loci == 0)
		return;
	
	// Disrupted site fixation or loss
	if (broken == 0 & gen > 1){
		catn(paste("DISRUPTED_LOST_AT_GEN:",gen));
	} else if (broken == 1){
		catn(paste("DISRUPTED_FIXED_AT_GEN:",gen));
	}
	
	
	if ((broken == 0 | broken == 1) & (drive == 1 | drive == 0)){
		catn(paste("BOTH_FIXED_OR_LOST_AT_GEN::", gen, "broken_rate:", broken, "drive_rate:", drive));
		ind = whichMax(genetic_load_track);
		max_genetic_load = genetic_load_track[ind];
		catn(paste("ENDING_GENETIC_LOAD:", genetic_load));
		catn(paste("MAXIMUM_GENETIC_LOAD:", max_genetic_load, "AT_GEN:", ind));
		sim.simulationFinished();
	}


}



10100 early(){
	
	// end 10,000 generations after the drive was released if not already stopped
	catn("SIMULATION_TIMED_OUT");
	
	ind = whichMax(genetic_load_track);
	max_genetic_load = genetic_load_track[ind];
	catn(paste("ENDING_GENETIC_LOAD:", genetic_load));
	catn(paste("MAXIMUM_GENETIC_LOAD:", max_genetic_load, "AT_GEN:", ind));
	
	sim.simulationFinished();
}