added standing variation options

broadinstitute · Oct 14, 2017 · 95c4acc · 95c4acc
1 parent a497e26
commit 95c4acc
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 16 deletions.
diff --git a/cosi/historical.cc b/cosi/historical.cc
@@ -415,23 +415,37 @@ class Event_Sweep: public HistEvents::Event {
 // A selective sweep .
 class Event_MSweep: public HistEvents::Event {
 public:
+	struct standing_variation_tag { };
+
+
 	 Event_MSweep( HistEvents *histEvents_, istream& is ): Event( histEvents_, is ) {
 		 is >> sweepPop >> gen >> selCoeff >> selPos >> final_sel_freq;
+		 selBegPop = sweepPop;
+		 selBegGen = gen;
 		 if ( !( final_sel_freq.getMin() && final_sel_freq.getMax() &&
 						 0. <= final_sel_freq.getMin() && final_sel_freq.getMin() <= 1. ) )
 				BOOST_THROW_EXCEPTION( cosi_hist_event_error() << error_msg( "invalid final freq range" ) );
 	 }
 
+	Event_MSweep( HistEvents *histEvents_, istream& is, standing_variation_tag ): Event( histEvents_, is ) {
+		 is >> sweepPop >> gen >> selCoeff >> selPos >> final_sel_freq >> selBegPop >> selBegGen;
+		 if ( !( final_sel_freq.getMin() && final_sel_freq.getMax() &&
+						 0. <= final_sel_freq.getMin() && final_sel_freq.getMin() <= 1. ) )
+				BOOST_THROW_EXCEPTION( cosi_hist_event_error() << error_msg( "invalid final freq range" ) );
+	 }
+
+
 	 // Method: typeStr
 	 // Return the string denoting this type of historical event in the <parameter file>.
 	 // Historical events are specified in the parameter file by lines of the form
 	 // pop_event <eventType> <eventParams>
 	 // This method specifies the eventType.
 	 static const char *typeStr() { return "sweep_mult"; }
+	 static const char *typeStr_standing_variation() { return "sweep_mult_standing"; }
 	 virtual eventKind_t getEventKind() const { return E_SWEEP; }
 
 	 virtual void addToBaseModel( BaseModel& m ) const {
-		 setSweepInfo( m, gen, selCoeff, selPos, sweepPop, final_sel_freq  );
+		 setSweepInfo( m, gen, selCoeff, selPos, sweepPop, final_sel_freq, selBegPop, selBegGen  );
 	 }
 
 	 virtual ~Event_MSweep();
@@ -440,7 +454,7 @@ class Event_MSweep: public HistEvents::Event {
 
 private:
 	 // Field: sweepPop
-	 // The population in which the sweep happens.
+	 // The population in which the selected allele is born
 	 popid sweepPop;
 
 	 // Field: selCoeff
@@ -454,6 +468,14 @@ class Event_MSweep: public HistEvents::Event {
 	 // Field: final_sel_freq
 	 // The frequency of the derived allele at the end of the sweep.
 	 util::ValRange<freq_t> final_sel_freq;
+
+	 // Field: selBegGen
+	 // Generation when selection begins in `selBegPop`.
+	 genid selBegGen;
+
+	 // Field: selBegPop
+	 // Pop in which selection begins at time `selBegGen`
+	 popid selBegPop;
 };  // class Event_MSweep
 
 
@@ -839,6 +861,8 @@ HistEvents::EventP HistEvents::parseEvent( const char *buffer ) {
 		else if ( typestr == sweep2::Event_SweepOnePop2_typeStr() ) event = sweep2::make_shared_Event_SweepOnePop2( this, is );
 		else if ( typestr == sweep3::Event_SweepOnePop3_typeStr() ) event = sweep3::make_shared_Event_SweepOnePop3( this, is );
 		else if (  typestr == Event_MSweep::typeStr() ) event.reset( new Event_MSweep( this, is ) );
+		else if (  typestr == Event_MSweep::typeStr_standing_variation() ) 
+			event.reset( new Event_MSweep( this, is, Event_MSweep::standing_variation_tag() ) );
 		else chkCond( False, "could not parse event %s", buffer );
 	} catch( ios::failure e ) {
 		chkCond( False, "could not parse event %s", buffer );

diff --git a/cosi/msweep.cc b/cosi/msweep.cc
@@ -60,19 +60,23 @@ struct SweepInfo {
 // *** Field: selPop - the population in which the selected allele is born.
 	 popid selPop;
 	 util::ValRange<freq_t> final_sel_freq;
+// *** Field: selBegPop - the population in which selection begins
+	popid selBegPop;
+// *** Field: selBegGen - the time at which selection begins
+	genid selBegGen;
 
 	 SweepInfo(): selGen( NULL_GEN ), selCoeff( 0.0 ), selPos( 0.0 ),
-								selPop( NULL_POPID ) { }
+								selPop( NULL_POPID ), selBegPop( NULL_POPID ), selBegGen( NULL_GEN ) { }
 	 SweepInfo( genid selGen_, double selCoeff_, loc_t selPos_, popid selPop_,
-							util::ValRange<freq_t> final_sel_freq_ ):
+							util::ValRange<freq_t> final_sel_freq_, popid selBegPop_, genid selBegGen_ ):
 		 selGen( selGen_ ), selCoeff( selCoeff_ ), selPos( selPos_ ), selPop( selPop_ ),
-		 final_sel_freq( final_sel_freq_ ) { }
+		 final_sel_freq( final_sel_freq_ ), selBegPop( selBegPop_ ), selBegGen( selBegGen_ ) { }
 };  // struct SweepInfo
 
 void setSweepInfo( BaseModel& baseModel,
 									 genid selGen, double selCoeff, loc_t selPos, popid selPop,
-									 util::ValRange<freq_t> final_sel_freq ) {
-	baseModel.sweepInfo = boost::make_shared<SweepInfo>( selGen, selCoeff, selPos, selPop, final_sel_freq );
+									 util::ValRange<freq_t> final_sel_freq, popid selBegPop, genid selBegGen ) {
+	baseModel.sweepInfo = boost::make_shared<SweepInfo>( selGen, selCoeff, selPos, selPop, final_sel_freq, selBegPop, selBegGen );
 }
 
 // ** class MSweep - implementation of selective sweep in multiple populations
@@ -145,6 +149,7 @@ class MSweep: public Module {
 			 } else
 					this->mtraj = this->simulateTrajFwd( baseModel, fits, sweepInfo.selGen,
 																							 begFreqs, endFreqs,
+																							 sweepInfo.selBegPop, sweepInfo.selBegGen, sweepInfo.selCoeff,
 																							 *randGen, maxAttempts );
 
 			 if(0){
@@ -394,6 +399,7 @@ class MSweep: public Module {
 										std::map<popid, std::map<genotype_t,double> > pop2fits,
 										genid begGen, std::map<popid,freq_t> begFreqs,
 										std::map<popid, util::ValRange<freq_t> > endFreqs,
+										popid selBegPop, genid selBegGen, double selCoeff,
 										URNG& urng,
 										size_t maxAttempts = 1000000 ) {
 		 cosi_using5( util::at, std::map, boost::assign::insert, boost::adaptors::map_keys,
@@ -455,14 +461,16 @@ class MSweep: public Module {
 					 { // if !trajFailed
 						 // for each pop, determine n_A and n_a after random mating and selection, but before migration
 						 std::map<popid, freq_t> b4mig_p_A;
-						 cosi_for_map( pop, fits, pop2fits ) {
-								 double fit_AA = at( fits, GT_AA ), fit_Aa = at( fits, GT_Aa ), fit_aa = at( fits, GT_aa );
-								 freq_t p_A = freqs[ pop ];
-								 freq_t p_a = 1. - p_A;
-								 double amt_A = p_A * p_A * fit_AA + p_A * p_a * fit_Aa;
-								 double amt_a = p_a * p_a * fit_aa + p_A * p_a * fit_Aa;
-								 b4mig_p_A[ pop ] = amt_A / ( amt_A + amt_a );
-								 if ( tr ) PRINT9( pop, fit_AA, fit_Aa, fit_aa, p_A, p_a, amt_A, amt_a, b4mig_p_A[pop] );
+						 cosi_for_map_keys( pop, pop2fits ) {
+							 double s = ( (pop == selBegPop) && (gen <= selBegGen) ) ? selCoeff : 0.;
+							 double fit_AA = 1.0 + s, fit_Aa = 1.0 + 0.5*s, fit_aa = 1.0;
+							 //	 double fit_AA = at( fits, GT_AA ), fit_Aa = at( fits, GT_Aa ), fit_aa = at( fits, GT_aa );
+							 freq_t p_A = freqs[ pop ];
+							 freq_t p_a = 1. - p_A;
+							 double amt_A = p_A * p_A * fit_AA + p_A * p_a * fit_Aa;
+							 double amt_a = p_a * p_a * fit_aa + p_A * p_a * fit_Aa;
+							 b4mig_p_A[ pop ] = amt_A / ( amt_A + amt_a );
+							 if ( tr ) PRINT9( pop, fit_AA, fit_Aa, fit_aa, p_A, p_a, amt_A, amt_a, b4mig_p_A[pop] );
 						 } cosi_end_for;
 
 						 // now model migration.

diff --git a/cosi/msweep.h b/cosi/msweep.h
@@ -29,7 +29,7 @@ void addSelMut( MSweepP, MutlistP );
 
 void setSweepInfo( BaseModel& baseModel,
 									 genid selGen, double selCoeff, loc_t selPos, popid selPop,
-									 util::ValRange<freq_t> final_sel_freq );
+									 util::ValRange<freq_t> final_sel_freq, popid selBegPop, genid selBegGen );
 
 }  // namespace cosi