1 #include <cppunit/extensions/HelperMacros.h>
2 
3 #include "../../src/forest.h"
4 #include "../../src/random/constant_generator.h"
5 #include "../../src/random/mersenne_twister.h"
6 #include "../../src/event.h"
7 #include "../../src/summary_statistics/tmrca.h"
8 
9 class TestForest : public CppUnit::TestCase {
10 
11   CPPUNIT_TEST_SUITE( TestForest );
12 
13   CPPUNIT_TEST( testInitialization );
14   CPPUNIT_TEST( testGettersAndSetters );
15   CPPUNIT_TEST( testCreateExampleTree );
16   CPPUNIT_TEST( testCheckTreeLength );
17   CPPUNIT_TEST( testGetFirstNode );
18   CPPUNIT_TEST( testSamplePoint );
19   CPPUNIT_TEST( testCalcRecombinationRate );
20   CPPUNIT_TEST( testCalcRate );
21   CPPUNIT_TEST( testCalcRateWithArachicSamples );
22   CPPUNIT_TEST( testNodeIsOld );
23   CPPUNIT_TEST( testPrune );
24   CPPUNIT_TEST( testSelectFirstTime );
25   CPPUNIT_TEST( testSampleEventType );
26   CPPUNIT_TEST( testSampleEvent );
27   CPPUNIT_TEST( testGetNodeState );
28   CPPUNIT_TEST( testCut );
29   CPPUNIT_TEST( testImplementCoalescence );
30   CPPUNIT_TEST( testBuildInitialTree );
31   CPPUNIT_TEST( testImplementRecombination );
32   CPPUNIT_TEST( testImplementFixTimeEvent );
33   CPPUNIT_TEST( testPrintTree );
34   CPPUNIT_TEST( testCopyConstructor );
35   CPPUNIT_TEST( testCheckForNodeAtHeight );
36   CPPUNIT_TEST( testPrintLocusSumStats );
37   CPPUNIT_TEST( testSampleNextPosition );
38   CPPUNIT_TEST( testClear );
39 
40   CPPUNIT_TEST_SUITE_END();
41 
42  private:
43   Model *model, *model_2pop;
44   Forest *forest, *forest_2pop;
45   MersenneTwister *rg;
46 
47  public:
setUp()48   void setUp() {
49     rg = new MersenneTwister(1234);
50     model = new Model(5);
51     model_2pop = new Model(5);
52     model_2pop->set_population_number(2);
53     forest = new Forest(model, rg);
54     forest->createExampleTree();
55     forest_2pop = new Forest(model_2pop, rg);
56     forest_2pop->createExampleTree();
57   }
58 
tearDown()59   void tearDown() {
60     delete forest;
61     delete model;
62     delete forest_2pop;
63     delete model_2pop;
64     delete rg;
65   }
66 
testInitialization()67   void testInitialization() {
68     Forest test_forest = Forest(new Model(4), rg);
69     CPPUNIT_ASSERT( test_forest.model().sample_size() == 4 );
70     CPPUNIT_ASSERT( test_forest.random_generator() == rg );
71     delete test_forest.writable_model();
72   }
73 
testGettersAndSetters()74   void testGettersAndSetters() {
75     CPPUNIT_ASSERT( forest->model().sample_size() == 4 );
76   }
77 
testGetFirstNode()78   void testGetFirstNode() {
79     CPPUNIT_ASSERT( forest->nodes()->get(0)->height() == 0 );
80   }
81 
testCreateExampleTree()82   void testCreateExampleTree() {
83     CPPUNIT_ASSERT_EQUAL((size_t)2, forest_2pop->model().population_number() );
84     CPPUNIT_ASSERT_EQUAL((size_t)9, forest->nodes()->size());
85     CPPUNIT_ASSERT( forest->local_root() == forest->nodes()->get(8) );
86     CPPUNIT_ASSERT( forest->primary_root() == forest->nodes()->get(8) );
87     CPPUNIT_ASSERT( forest->getLocalTreeLength() == 24 );
88     CPPUNIT_ASSERT( forest->checkTree() );
89 
90     forest->createExampleTree();
91     CPPUNIT_ASSERT_EQUAL((size_t)9, forest->nodes()->size());
92     CPPUNIT_ASSERT( forest->local_root() == forest->nodes()->get(8) );
93     CPPUNIT_ASSERT( forest->primary_root() == forest->nodes()->get(8) );
94     CPPUNIT_ASSERT( forest->getLocalTreeLength() == 24 );
95     CPPUNIT_ASSERT( forest->checkTree() );
96   }
97 
testCheckTreeLength()98   void testCheckTreeLength() {
99     CPPUNIT_ASSERT( forest->checkTreeLength() );
100   }
101 
102 
testCalcRate()103   void testCalcRate() {
104     TimeIntervalIterator tii(forest, forest->nodes()->at(0));
105     double pop_size = 2*forest->model().population_size(0);
106     Node *node1 = new Node(0.1);
107     Node *node2 = new Node(0.2);
108 
109     forest->set_active_node(0, node1);
110     forest->set_active_node(1, node2);
111     forest->states_[0] = 0;
112     forest->states_[1] = 0;
113     CPPUNIT_ASSERT_NO_THROW( forest->calcRates(*tii) );
114     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[0] );
115     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[1] );
116     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[2] );
117 
118     forest->states_[0] = 1;
119     CPPUNIT_ASSERT_NO_THROW( forest->calcRates(*tii) );
120     CPPUNIT_ASSERT( areSame(4.0/pop_size, forest->rates_[0]) );
121     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[1] );
122     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[2] );
123 
124     forest->states_[1] = 1;
125     CPPUNIT_ASSERT_NO_THROW( forest->calcRates(*tii) );
126     CPPUNIT_ASSERT( areSame(9.0/pop_size, forest->rates_[0]) );
127     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[1] );
128     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[2] );
129 
130     // Coalescence with structure
131     forest_2pop->set_active_node(0, node1);
132     forest_2pop->set_active_node(1, node2);
133     forest_2pop->states_[0] = 1;
134     forest_2pop->states_[1] = 1;
135 
136     node1->set_population(1);
137     forest_2pop->nodes()->at(1)->set_population(1);
138 
139     TimeIntervalIterator tii2(forest_2pop, forest_2pop->nodes()->at(0));
140     CPPUNIT_ASSERT_NO_THROW( forest_2pop->calcRates(*tii2) );
141     // Only node2 can coalescence
142     CPPUNIT_ASSERT( areSame(4.0/pop_size, forest_2pop->rates_[0]) );
143     CPPUNIT_ASSERT_EQUAL( 0.0, forest_2pop->rates_[1] );
144     CPPUNIT_ASSERT_EQUAL( 0.0, forest_2pop->rates_[2] );
145 
146     std::vector<double> growth(2, 0.0);
147     growth.at(1) = 1.0;
148     forest_2pop->writable_model()->addGrowthRates(0, growth);
149     growth.at(0) = 2.0;
150     forest_2pop->writable_model()->addGrowthRates(1, growth);
151     TimeIntervalIterator tii3(forest_2pop, forest_2pop->nodes()->at(0));
152 
153     CPPUNIT_ASSERT_NO_THROW( forest_2pop->calcRates(*tii3) );
154     CPPUNIT_ASSERT( areSame(3.0/pop_size, forest_2pop->rates_[0]) );
155     CPPUNIT_ASSERT( areSame(1.0/pop_size, forest_2pop->rates_[1]) );
156     CPPUNIT_ASSERT( areSame(0.0/pop_size, forest_2pop->rates_[2]) );
157     CPPUNIT_ASSERT_EQUAL( (size_t)1, forest_2pop->active_nodes_timelines_[0] );
158     CPPUNIT_ASSERT_EQUAL( (size_t)0, forest_2pop->active_nodes_timelines_[1] );
159 
160     forest_2pop->writable_model()->increaseTime();
161     CPPUNIT_ASSERT_NO_THROW( forest_2pop->calcRates(*tii3) );
162     //CPPUNIT_ASSERT_EQUAL( 0.0/pop_size, forest_2pop->rates_[0] );
163     //CPPUNIT_ASSERT_EQUAL( 1.0/pop_size, forest_2pop->rates_[1] );
164     //CPPUNIT_ASSERT_EQUAL( 3.0/pop_size, forest_2pop->rates_[2] );
165     CPPUNIT_ASSERT_EQUAL( (size_t)1, forest_2pop->active_nodes_timelines_[0] );
166     CPPUNIT_ASSERT_EQUAL( (size_t)2, forest_2pop->active_nodes_timelines_[1] );
167 
168     node2->set_population(1);
169     CPPUNIT_ASSERT_NO_THROW( forest_2pop->calcRates(*tii3) );
170     //CPPUNIT_ASSERT_EQUAL( 0.0/pop_size, forest_2pop->rates_[0] );
171     //CPPUNIT_ASSERT( areSame(3.0/pop_size, forest_2pop->rates_[1]) );
172     //CPPUNIT_ASSERT_EQUAL( 0.0/pop_size, forest_2pop->rates_[2] );
173     CPPUNIT_ASSERT_EQUAL( (size_t)1, forest_2pop->active_nodes_timelines_[0] );
174     CPPUNIT_ASSERT_EQUAL( (size_t)1, forest_2pop->active_nodes_timelines_[1] );
175 
176     delete node1;
177     delete node2;
178   }
179 
testCalcRateWithArachicSamples()180   void testCalcRateWithArachicSamples() {
181     Node* node1 = forest->nodes()->createNode(20, 5);
182     forest->nodes()->add(node1);
183     TimeIntervalIterator tii(forest, node1);
184     double pop_size = 2*forest->model().population_size(0);
185 
186     forest->set_active_node(0, node1);
187     forest->set_active_node(1, forest->local_root());
188     forest->states_[0] = 1;
189     forest->states_[1] = 1;
190 
191     CPPUNIT_ASSERT_NO_THROW( forest->calcRates(*tii) );
192     CPPUNIT_ASSERT( areSame(1.0/pop_size, forest->rates_[0]) );
193     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[1] );
194     CPPUNIT_ASSERT_EQUAL( 0.0, forest->rates_[2] );
195   }
196 
testGetNodeState()197   void testGetNodeState() {
198     CPPUNIT_ASSERT( forest->getNodeState(forest->getNodes()->get(8), 11) == 1 );
199     CPPUNIT_ASSERT( forest->getNodeState(forest->getNodes()->get(6), 11) == 2 );
200     CPPUNIT_ASSERT( forest->getNodeState(forest->getNodes()->get(7), 11) == 1 );
201   }
202 
testPrintTree()203   void testPrintTree() {
204     //CPPUNIT_ASSERT_NO_THROW( forest->printTree() );
205   }
206 
testSelectFirstTime()207   void testSelectFirstTime() {
208     double current_time = -1.0;
209     size_t time_line = -1;
210 
211     forest->selectFirstTime(-1, 1, current_time, time_line);
212     CPPUNIT_ASSERT_EQUAL( -1.0, current_time );
213 
214     forest->selectFirstTime(5.0, 1, current_time, time_line);
215     CPPUNIT_ASSERT_EQUAL( 5.0, current_time );
216     CPPUNIT_ASSERT_EQUAL( (size_t)1, time_line );
217 
218     forest->selectFirstTime(7.0, 2, current_time, time_line);
219     CPPUNIT_ASSERT_EQUAL( 5.0, current_time );
220     CPPUNIT_ASSERT_EQUAL( (size_t)1, time_line );
221 
222     forest->selectFirstTime(-1, 2, current_time, time_line);
223     CPPUNIT_ASSERT_EQUAL( 5.0, current_time );
224     CPPUNIT_ASSERT_EQUAL( (size_t)1, time_line );
225 
226     forest->selectFirstTime(4.0, 2, current_time, time_line);
227     CPPUNIT_ASSERT_EQUAL( 4.0, current_time );
228     CPPUNIT_ASSERT_EQUAL( (size_t)2, time_line );
229   }
230 
testSampleEventType()231   void testSampleEventType() {
232     Event event;
233     Forest *forest2 = forest_2pop;
234 
235     forest2->nodes()->at(0)->set_population(1);
236     forest2->nodes()->at(1)->set_population(1);
237     forest2->writable_model()->addMigrationRates(1, std::vector<double>(4, 5.0), false, true);
238     forest2->writable_model()->addGrowthRates(0.2, std::vector<double>(2, 0.000125));
239     forest2->writable_model()->addGrowthRates(1, std::vector<double>(2, 0.0));
240     forest2->writable_model()->addGrowthRates(2, std::vector<double>(2, 2));
241     forest2->writable_model()->finalize();
242     forest2->writable_model()->resetTime();
243 
244     TimeIntervalIterator tii(forest2, forest2->nodes()->at(0));
245 
246     forest2->set_active_node(0, forest2->nodes()->at(0));
247     forest2->set_active_node(1, forest2->nodes()->at(2));
248     forest2->states_[0] = 1;
249     forest2->states_[1] = 0;
250     forest2->calcRates(*tii);
251     forest2->active_nodes_timelines_[0] = 0;
252     forest2->active_nodes_timelines_[1] = 0;
253 
254     forest2->sampleEventType(-1, 0, *tii, event);
255     CPPUNIT_ASSERT( event.isNoEvent() );
256 
257     forest2->sampleEventType(0.5, 0, *tii, event);
258     CPPUNIT_ASSERT_EQUAL( 0.5, event.time() );
259     CPPUNIT_ASSERT( event.isCoalescence() );
260 
261     // Only coalescence possible
262     for (size_t i = 0; i < 1000; ++i) {
263       forest2->sampleEventType(0.5, 0, *tii, event);
264       CPPUNIT_ASSERT( event.isCoalescence() );
265       CPPUNIT_ASSERT( event.node() == forest2->active_node(0) );
266     };
267 
268     // Coalescence of two nodes
269     // active_node 0: Pop 1, 2 Contemporaries
270     // active_node 1: Pop 0, 2 Contemporaries
271     // => 50% each
272     forest2->states_[1] = 1;
273     forest2->calcRates(*tii);
274     size_t count = 0, count2 = 0;
275     for (size_t i = 0; i < 10000; ++i) {
276       forest2->sampleEventType(0.5, 0, *tii, event);
277       CPPUNIT_ASSERT( event.isCoalescence() );
278       count += (event.node() == forest2->active_node(0));
279     };
280     CPPUNIT_ASSERT( 4900 < count && count < 5100 ); // ~5000
281 
282     // Test with Pw Coalescence
283     // active_node 0: Pop 1, 2 Contemporaries
284     // active_node 1: Pop 1, 2 Contemporaries
285     // 4/5 Prob of Coalescence, 1/5 of Pw coalescence
286     forest2->writable_model()->resetTime();
287     forest2->set_active_node(1, forest2->nodes()->at(1));
288     forest2->calcRates(*tii);
289     count = 0;
290     for (size_t i = 0; i < 10000; ++i) {
291       forest2->sampleEventType(0.5, 0, *tii, event);
292       CPPUNIT_ASSERT( event.isCoalescence() || event.isPwCoalescence() );
293       count += event.isCoalescence();
294     };
295     CPPUNIT_ASSERT( 7900 < count && count < 8100 );
296 
297     // Test with migration
298     // active_node 0: Pop 1, 2 Contemporaries
299     // active_node 1: Pop 1, 2 Contemporaries
300     // => Coal: 4/2Ne, Pw Coal: 1/2Ne
301     //    Migration: 2 * 5/4Ne
302     // => 40% Coal, 10% Pw Coal, 50% Mig
303     forest2->writable_model()->increaseTime();
304     forest2->writable_model()->increaseTime();
305     forest2->calcRates(*tii);
306 
307     count = 0;
308     for (size_t i = 0; i < 10000; ++i) {
309       forest2->sampleEventType(0.5, 0, *tii, event);
310       CPPUNIT_ASSERT( event.isCoalescence() || event.isPwCoalescence() || event.isMigration() );
311       count += event.isMigration();
312       count2 += event.isCoalescence();
313     };
314     CPPUNIT_ASSERT( 4900 < count && count < 5100 );
315     CPPUNIT_ASSERT( 3900 < count2 && count2 < 4100 );
316 
317 
318     // Coalescence and Recombination
319     // active_node 0: Pop 1, 2 Contemporaries => Coal rate: 2 / 2 * Ne = 1/Ne
320     // active_node 1: Pop 1, Recombination    => Rec rate: 10 Bases * 0.4 / 4Ne = 1/Ne
321     // => 50% Recombination, 50% Coalescence
322     forest2->writable_model()->setLocusLength(101);
323     forest2->writable_model()->setRecombinationRate(0.4, false, true);
324     forest2->set_current_base(10);
325     forest2->active_node(1)->make_nonlocal(forest2->current_rec_);
326     forest2->set_next_base(20);
327     forest2->current_rec_++;
328     forest2->states_[0] = 1;
329     forest2->states_[1] = 2;
330     forest2->writable_model()->resetTime(); // set migration to 0
331     forest2->calcRates(*tii);
332 
333     count = 0;
334     for (size_t i = 0; i < 100000; ++i) {
335       forest2->sampleEventType(0.5, 0, *tii, event);
336       CPPUNIT_ASSERT( event.isCoalescence() || event.isRecombination() );
337       count += event.isRecombination();
338     };
339     CPPUNIT_ASSERT( 49500 < count && count < 50500 );
340 
341     // Other way round
342     Node* tmp = forest2->active_node(1);
343     forest2->set_active_node(1, forest2->active_node(0));
344     forest2->set_active_node(0, tmp);
345     forest2->states_[0] = 2;
346     forest2->states_[1] = 1;
347     forest2->calcRates(*tii);
348 
349     count = 0;
350     for (size_t i = 0; i < 100000; ++i) {
351       forest2->sampleEventType(0.5, 0, *tii, event);
352       CPPUNIT_ASSERT( event.isCoalescence() || event.isRecombination() );
353       count += event.isRecombination();
354     };
355     CPPUNIT_ASSERT( 49500 < count && count < 50500 ); // True with >95%
356 
357 
358     // Double recombination
359     // => 1/3 active_node 0
360     //    2/3 active_node 1
361     forest2->states_[0] = 2;
362     forest2->states_[1] = 2;
363 
364     forest2->set_current_base(10.0);
365     forest2->set_next_base(15.0);
366     forest2->set_next_base(20.0);
367 
368     forest2->current_rec_ += 2;
369     forest2->active_node(0)->make_nonlocal(forest2->current_rec_ - 1);
370     forest2->active_node(1)->make_nonlocal(forest2->current_rec_ - 2);
371     forest2->calcRates(*tii);
372 
373     count = 0;
374     for (size_t i = 0; i < 15000; ++i) {
375       forest2->sampleEventType(0.5, 0, *tii, event);
376       CPPUNIT_ASSERT( event.isRecombination() );
377       count += (event.node() == forest2->active_node(0));
378     };
379     CPPUNIT_ASSERT( 4880 < count && count < 5120 ); // True with >96%
380 
381 
382     // Recombination with Rate 0
383     // active_node 1: Up to date = rec rate = 0;
384     // => always coalescence
385     forest2->states_[0] = 1;
386     forest2->active_node(0)->make_local();
387     forest2->active_node(1)->set_last_update(forest2->current_rec_);
388     forest2->calcRates(*tii);
389     for (size_t i = 0; i < 1000; ++i) {
390       forest2->sampleEventType(0.5, 0, *tii, event);
391       CPPUNIT_ASSERT( event.isCoalescence() );
392     };
393 
394     // Combinations With Growth
395     forest2->writable_model()->resetTime();
396     forest2->writable_model()->increaseTime();
397     forest2->writable_model()->increaseTime();
398     forest2->writable_model()->increaseTime();
399     forest2->states_[0] = 1;
400     forest2->states_[1] = 1;
401     forest2->active_node(0)->set_population(0);
402     forest2->active_node(1)->set_population(1);
403     forest2->calcRates(*tii);
404 
405     for (size_t i = 0; i < 100; ++i) {
406       CPPUNIT_ASSERT_NO_THROW( forest2->sampleEventType(0.5, 0, *tii, event) );
407       CPPUNIT_ASSERT( event.isMigration() );
408       CPPUNIT_ASSERT_NO_THROW( forest2->sampleEventType(0.5, 1, *tii, event) );
409       CPPUNIT_ASSERT( event.isCoalescence() );
410       CPPUNIT_ASSERT( event.node() == forest2->active_node(0) );
411       CPPUNIT_ASSERT_NO_THROW( forest2->sampleEventType(0.5, 2, *tii, event) );
412       CPPUNIT_ASSERT( event.isCoalescence() );
413       CPPUNIT_ASSERT( event.node() == forest2->active_node(1) );
414     };
415 
416     forest2->active_node(0)->set_population(0);
417     forest2->active_node(1)->set_population(0);
418     forest2->calcRates(*tii);
419     for (size_t i = 0; i < 100; ++i) {
420       CPPUNIT_ASSERT_NO_THROW( forest2->sampleEventType(0.5, 0, *tii, event) );
421       CPPUNIT_ASSERT( event.isMigration() );
422       CPPUNIT_ASSERT_NO_THROW( forest2->sampleEventType(0.5, 1, *tii, event) );
423       CPPUNIT_ASSERT( event.isCoalescence() || event.isPwCoalescence() );
424     };
425   }
426 
testCalcRecombinationRate()427   void testCalcRecombinationRate() {
428     forest->writable_model()->setRecombinationRate(1, false, false, 0);
429     forest->writable_model()->setRecombinationRate(2, false, false, 10);
430     forest->writable_model()->setRecombinationRate(3, false, false, 15);
431     forest->writable_model()->setRecombinationRate(4, false, false, 20);
432     forest->writable_model()->resetSequencePosition();
433 
434     Node* node = forest->nodes()->at(6);
435 
436     forest->set_next_base(7.0);
437     forest->current_rec_++;
438     CPPUNIT_ASSERT_EQUAL( 2.0, forest->calcRecombinationRate(node) );
439 
440     forest->writable_model()->increaseSequencePosition();
441     forest->set_current_base(12);
442     CPPUNIT_ASSERT_EQUAL( 5.0+2*2, forest->calcRecombinationRate(node) );
443 
444     forest->writable_model()->increaseSequencePosition();
445     forest->set_current_base(17);
446     CPPUNIT_ASSERT_EQUAL( 5.0+10.0+6.0, forest->calcRecombinationRate(node) );
447 
448     forest->writable_model()->increaseSequencePosition();
449     forest->set_current_base(22);
450     CPPUNIT_ASSERT_EQUAL( 5.0+10.0+15.0+8.0, forest->calcRecombinationRate(node) );
451   }
452 
testSampleEvent()453   void testSampleEvent() {
454     Model model = Model(5);
455     Forest forest2 = Forest(&model, rg);
456     forest2.createScaledExampleTree();
457     forest2.writable_model()->finalize();
458     forest2.writable_model()->resetTime();
459 
460     TimeIntervalIterator tii(&forest2, forest2.nodes()->at(0));
461 
462     forest2.set_active_node(0, forest2.nodes()->at(0));
463     forest2.set_active_node(1, forest2.nodes()->at(8));
464     forest2.states_[0] = 1;
465     forest2.states_[1] = 0;
466     forest2.calcRates(*tii);
467     forest2.active_nodes_timelines_[0] = 0;
468     forest2.active_nodes_timelines_[1] = 0;
469 
470     Event event;
471     double tmp_event_time = 0.0;
472     for (size_t i = 0; i < 1000; ++i) {
473       forest2.sampleEvent(*tii, tmp_event_time, event);
474       CPPUNIT_ASSERT( event.isNoEvent() || ( 0 <= event.time() && event.time() < forest2.nodes()->at(4)->height() ) );
475       CPPUNIT_ASSERT( event.isNoEvent() || event.isCoalescence() );
476     }
477 
478     ++tii;
479     forest2.calcRates(*tii);
480     for (size_t i = 0; i < 1000; ++i) {
481       forest2.sampleEvent(*tii, tmp_event_time, event);
482       CPPUNIT_ASSERT( event.isNoEvent() || ( forest2.nodes()->at(4)->height() <= event.time() && event.time() < forest2.nodes()->at(5)->height() ) );
483       CPPUNIT_ASSERT( event.isNoEvent() || event.isCoalescence() );
484     }
485   }
486 
testNodeIsOld()487   void testNodeIsOld() {
488     forest->set_current_base(5.0);
489     forest->set_next_base(15);
490     forest->current_rec_++;
491 
492     forest->writable_model()->disable_approximation();
493     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(0)) );
494     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(5)) );
495     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(6)) );
496     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(7)) );
497     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(8)) );
498 
499     forest->writable_model()->set_window_length_seq(5);
500     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(0)) );
501     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(5)) );
502     CPPUNIT_ASSERT(  forest->nodeIsOld(forest->nodes()->at(6)) );
503     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(7)) );
504     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(8)) );
505 
506     forest->writable_model()->set_window_length_seq(9.5);
507     CPPUNIT_ASSERT( forest->nodeIsOld(forest->nodes()->at(6)) );
508     forest->writable_model()->set_window_length_seq(10.5);
509     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(6)) );
510 
511     forest->createExampleTree();
512     forest->writable_model()->set_window_length_rec(2);
513     forest->set_next_base(15);
514     forest->current_rec_++;
515     forest->set_next_base(20);
516     forest->current_rec_++;
517     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(0)) );
518     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(5)) );
519     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(6)) );
520     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(7)) );
521     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(8)) );
522 
523     forest->set_next_base(25);
524     forest->current_rec_++;
525     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(0)) );
526     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(5)) );
527     CPPUNIT_ASSERT(  forest->nodeIsOld(forest->nodes()->at(6)) );
528     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(7)) );
529     CPPUNIT_ASSERT(! forest->nodeIsOld(forest->nodes()->at(8)) );
530   }
531 
testPrune()532   void testPrune() {
533     // Old node
534     forest->set_current_base(5.0);
535     forest->set_next_base(15);
536     forest->current_rec_++;
537 
538     forest->writable_model()->disable_approximation();
539     forest->writable_model()->set_window_length_seq(5);
540     CPPUNIT_ASSERT( forest->model().has_approximation() );
541     CPPUNIT_ASSERT( forest->model().has_window_seq() );
542     CPPUNIT_ASSERT( !forest->model().has_window_rec() );
543 
544     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(0)) );
545     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(1)) );
546     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(2)) );
547     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(3)) );
548     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(4)) );
549     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(5)) );
550     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(7)) );
551     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(8)) );
552 
553     CPPUNIT_ASSERT( forest->pruneNodeIfNeeded(forest->nodes()->at(6)) );
554     CPPUNIT_ASSERT( forest->nodes()->size() == 8);
555     CPPUNIT_ASSERT( forest->checkTree() );
556 
557     // Orphaned node
558     CPPUNIT_ASSERT( forest->pruneNodeIfNeeded(forest->nodes()->at(6)) );
559     CPPUNIT_ASSERT( forest->nodes()->size() == 7);
560     CPPUNIT_ASSERT( forest->checkTree() == 1 );
561 
562     // In-Between Nodes should be pruned, iff they are of same age
563     Node *parent = forest->nodes()->createNode(20),
564          *inbetween1 = forest->nodes()->createNode(19),
565          *inbetween2 = forest->nodes()->createNode(18),
566          *child = forest->nodes()->createNode(17);
567 
568     forest->set_current_base(13);
569     inbetween2->make_nonlocal(forest->current_rec_);
570     child->make_nonlocal(forest->current_rec_);
571 
572     forest->set_next_base(15);
573     forest->current_rec_++;
574     parent->make_nonlocal(forest->current_rec_);
575     inbetween1->make_nonlocal(forest->current_rec_);
576     forest->nodes()->add(parent);
577     forest->nodes()->add(inbetween1);
578     forest->nodes()->add(inbetween2);
579     forest->nodes()->add(child);
580 
581     parent->set_first_child(inbetween1);
582     inbetween1->set_first_child(inbetween2);
583     inbetween2->set_first_child(child);
584     child->set_parent(inbetween2);
585     inbetween2->set_parent(inbetween1);
586     inbetween1->set_parent(parent);
587 
588     CPPUNIT_ASSERT( !forest->nodeIsOld(inbetween2) );
589     CPPUNIT_ASSERT( forest->pruneNodeIfNeeded(inbetween2) );
590     CPPUNIT_ASSERT_EQUAL(inbetween1, child->parent());
591     CPPUNIT_ASSERT( inbetween1->parent() == parent );
592     CPPUNIT_ASSERT( parent->is_root() );
593     CPPUNIT_ASSERT( parent->first_child() == inbetween1 );
594     CPPUNIT_ASSERT( inbetween1->first_child() == child );
595     CPPUNIT_ASSERT( child->countChildren() == 0 );
596 
597     forest->nodes()->at(0)->set_parent(NULL);
598     CPPUNIT_ASSERT(! forest->pruneNodeIfNeeded(forest->nodes()->at(0)) );
599   }
600 
testBuildInitialTree()601   void testBuildInitialTree() {
602     Model model = Model(5);
603 
604     for (size_t i = 0; i < 100; ++i) {
605       Forest frst = Forest(&model, rg);
606       frst.buildInitialTree();
607       CPPUNIT_ASSERT_EQUAL(true, frst.checkTree());
608       CPPUNIT_ASSERT( frst.current_base() == 0.0 );
609       CPPUNIT_ASSERT( frst.next_base() > 0.0 );
610     }
611 
612     model = Model(2);
613     model.set_population_number(2);
614     model.addSampleSizes(0.0, std::vector<size_t>(2, 1));
615     model.addSymmetricMigration(0.0, 5.0, true, true);
616     model.finalize();
617     Forest frst = Forest(&model, rg);
618     CPPUNIT_ASSERT_NO_THROW( frst.buildInitialTree() );
619     CPPUNIT_ASSERT_EQUAL( true, frst.checkTree() );
620     size_t i = 0;
621     for (size_t j = 0; j < 4; ++j) {
622       CPPUNIT_ASSERT( frst.nodes()->at(j)->population() == 0 ||
623                       frst.nodes()->at(j)->population() == 1 );
624       i += frst.nodes()->at(j)->population();
625     }
626     CPPUNIT_ASSERT_EQUAL( (size_t)1, i );
627 
628     model = Model(0);
629     model.set_population_number(2);
630     std::vector<size_t> sample_size(2, 1);
631     sample_size.at(1) = 0;
632     model.addSampleSizes(0.0, sample_size);
633     sample_size.at(0) = 0;
634     sample_size.at(1) = 1;
635     model.addSampleSizes(1.0, sample_size);
636     model.addSymmetricMigration(0.0, 5.0, true, true);
637     model.finalize();
638     frst = Forest(&model, rg);
639     CPPUNIT_ASSERT_NO_THROW( frst.buildInitialTree() );
640     CPPUNIT_ASSERT_EQUAL( true, frst.checkTree() );
641     CPPUNIT_ASSERT_EQUAL( (size_t)0, frst.nodes()->at(0)->population() );
642     CPPUNIT_ASSERT_EQUAL( (size_t)1, frst.nodes()->at(1)->population() );
643     CPPUNIT_ASSERT_EQUAL( 0.0, frst.nodes()->at(0)->height() );
644     CPPUNIT_ASSERT_EQUAL( 1.0, frst.nodes()->at(1)->height() );
645   }
646 
testCut()647   void testCut() {
648     Node* base_node = forest->nodes()->at(4);
649     Node* new_root = forest->cut(TreePoint(base_node, 3.5, false));
650 
651     CPPUNIT_ASSERT_EQUAL((size_t)11, forest->nodes()->size());
652     CPPUNIT_ASSERT( new_root->local() );
653     CPPUNIT_ASSERT( new_root->is_root() );
654     CPPUNIT_ASSERT_EQUAL((size_t)1, new_root->countChildren() );
655     CPPUNIT_ASSERT_EQUAL(3.5, new_root->height() );
656 
657     CPPUNIT_ASSERT( base_node->parent() == new_root );
658     CPPUNIT_ASSERT( base_node->local() );
659 
660     Node* single_branch = forest->local_root()->first_child();
661     CPPUNIT_ASSERT( !single_branch->local() );
662     CPPUNIT_ASSERT_EQUAL( forest->segment_count(), single_branch->last_update() );
663     CPPUNIT_ASSERT_EQUAL( (size_t)0, single_branch->countChildren() );
664     CPPUNIT_ASSERT_EQUAL( 3.5, single_branch->height() );
665   }
666 
testImplementRecombination()667   void testImplementRecombination() {
668     Node* new_root = forest->cut(TreePoint(forest->nodes()->at(4), 3.5, false));
669     TimeIntervalIterator tii(forest, new_root);
670     forest->set_active_node(0, new_root);
671     forest->set_active_node(1, forest->local_root());
672 
673     forest->states_[0] = 2;
674     forest->states_[1] = 0;
675 
676     Event event = Event(3.7);
677     event.setToCoalescence(new_root, 0);
678     forest->tmp_event_ = event;
679 
680     forest->implementCoalescence(event, tii);
681     ++tii;
682 
683     forest->set_current_base(100);
684     event.set_time(4.5);
685     event.setToRecombination(forest->active_node(0), 0);
686     CPPUNIT_ASSERT_NO_THROW( forest->implementRecombination(event, tii) );
687 
688     CPPUNIT_ASSERT( forest->nodes()->at(9) == forest->active_node(0) ||
689                    forest->nodes()->at(10) == forest->active_node(0) );
690 
691     CPPUNIT_ASSERT( forest->active_node(0)->local() );
692     CPPUNIT_ASSERT( forest->active_node(0)->is_root() );
693     CPPUNIT_ASSERT_EQUAL( (size_t)1, forest->active_node(0)->countChildren() );
694 
695     Node* single_branch = forest->nodes()->at(9);
696     if( forest->nodes()->at(9) == forest->active_node(0) )
697       single_branch = forest->nodes()->at(10);
698 
699     CPPUNIT_ASSERT( !single_branch->local() );
700     CPPUNIT_ASSERT_EQUAL( (size_t)0, single_branch->countChildren() );
701     CPPUNIT_ASSERT_EQUAL( (size_t)1, single_branch->last_update() );
702   }
703 
testImplementCoalescence()704   void testImplementCoalescence() {
705     for (size_t i=0; i<100; ++i) {
706       forest->createExampleTree();
707 
708       Node* new_root = forest->cut(TreePoint(forest->nodes()->at(1), 1.5, false));
709       TimeIntervalIterator tii(forest, new_root);
710       CPPUNIT_ASSERT_EQUAL( (size_t)3, forest->contemporaries_.size(0) );
711       forest->set_active_node(0, new_root);
712       forest->set_active_node(1, forest->local_root());
713 
714       forest->states_[0] = 1;
715       forest->states_[1] = 0;
716 
717       forest->tmp_event_ = Event(2.0);
718       forest->tmp_event_.setToCoalescence(new_root, 0);
719 
720       forest->implementCoalescence(forest->tmp_event_, tii);
721 
722       CPPUNIT_ASSERT( forest->active_node(0) == new_root );
723       CPPUNIT_ASSERT( new_root->parent_height() == 3 ||
724                      new_root->parent_height() == 10 );
725 
726       if ( !new_root->local() ) {
727         CPPUNIT_ASSERT( new_root->countChildren() == 2 );
728         CPPUNIT_ASSERT( !(new_root->first_child()->local() && new_root->second_child()->local()) );
729         Node* child = new_root->first_child();
730         if (!new_root->second_child()->local()) child = new_root->second_child();
731         CPPUNIT_ASSERT( child->last_update() == 1 );
732         CPPUNIT_ASSERT( child->length_below() == 0 );
733 
734         CPPUNIT_ASSERT( new_root->last_update() == 1 );
735       }
736     }
737   }
738 
testImplementFixTimeEvent()739   void testImplementFixTimeEvent() {
740     Model* model2 = new Model(5);
741     model2->set_population_number(3);
742     model2->addSingleMigrationEvent(0.5, 0, 1, 1.0);
743     Forest* forest2 = new Forest(model2, rg);
744     forest2->createExampleTree();
745     Node* new_root = forest2->cut(TreePoint(forest2->nodes()->at(1), 0.5, false));
746     forest2->set_active_node(0, new_root);
747     forest2->set_active_node(1, forest2->local_root());
748     forest2->states_[0] = 1;
749     forest2->states_[1] = 0;
750     TimeIntervalIterator tii(forest2, new_root);
751 
752     CPPUNIT_ASSERT( new_root->population() == 0 );
753     forest2->implementFixedTimeEvent(tii);
754     CPPUNIT_ASSERT( new_root->is_root() );
755     CPPUNIT_ASSERT( new_root->population() == 1 );
756 
757     // Chained events
758     new_root->set_population(0);
759     model2->addSingleMigrationEvent(0.5, 1, 2, 1.0);
760     forest2->implementFixedTimeEvent(tii);
761     CPPUNIT_ASSERT( new_root->is_root() );
762     CPPUNIT_ASSERT( new_root->population() == 2 );
763 
764     // Circes do not cause problems
765     model2->addSingleMigrationEvent(0.5, 2, 0, 1.0);
766     forest2->implementFixedTimeEvent(tii);
767     CPPUNIT_ASSERT( new_root->is_root() );
768     CPPUNIT_ASSERT( new_root->population() == 0 );
769 
770     delete forest2;
771     delete model2;
772   }
773 
testSamplePoint()774   void testSamplePoint() {
775     rg->set_seed(1234);
776     forest->createScaledExampleTree();
777 
778     TreePoint point;
779     int n0 = 0, n1 = 0, n2 = 0, n3 = 0,
780         n4 = 0, n5 = 0;
781     for (int i = 0; i < 240000; ++i) {
782       point = forest->samplePoint();
783       CPPUNIT_ASSERT( point.base_node() != NULL );
784       CPPUNIT_ASSERT( point.base_node()->local() );
785       if (point.base_node() == forest->nodes()->at(0)) ++n0;
786       if (point.base_node() == forest->nodes()->at(1)) ++n1;
787       if (point.base_node() == forest->nodes()->at(2)) ++n2;
788       if (point.base_node() == forest->nodes()->at(3)) ++n3;
789       if (point.base_node() == forest->nodes()->at(4)) ++n4;
790       if (point.base_node() == forest->nodes()->at(5)) ++n5;
791     }
792 
793     //std::cout << n0 << " " << n1 << " " << n2 << " "
794     //          << n3 << " " << n4 << " " << n5 << std::endl;
795     CPPUNIT_ASSERT( 29500 <= n0 && n0 <= 30500 ); // expected 30000
796     CPPUNIT_ASSERT( 29500 <= n1 && n1 <= 30500 ); // expected 30000
797     CPPUNIT_ASSERT(  9800 <= n2 && n2 <= 10200 ); // expected 10000
798     CPPUNIT_ASSERT(  9800 <= n3 && n3 <= 10200 ); // expected 10000
799     CPPUNIT_ASSERT( 89000 <= n4 && n4 <= 91000 ); // expected 90000
800     CPPUNIT_ASSERT( 69000 <= n5 && n5 <= 71000 ); // expected 70000
801   }
802 
testCopyConstructor()803   void testCopyConstructor() {
804     forest->createScaledExampleTree();
805     forest->set_next_base(10.0);
806     CPPUNIT_ASSERT( forest->coalescence_finished_ == true );
807 
808     Forest forest2(*forest);
809 
810     CPPUNIT_ASSERT_EQUAL(forest->nodes()->size(), forest2.nodes()->size() );
811     CPPUNIT_ASSERT(forest2.model_ == forest->model_);
812     CPPUNIT_ASSERT(forest2.local_root() != NULL);
813     CPPUNIT_ASSERT(forest2.local_root() != forest->local_root());
814 
815     CPPUNIT_ASSERT_EQUAL(forest->rec_bases_.size(), forest2.rec_bases_.size());
816     CPPUNIT_ASSERT_EQUAL(forest->current_base(), forest2.current_base());
817 
818     for (auto it = forest2.nodes()->iterator(); it.good(); ++it) {
819       CPPUNIT_ASSERT( (*it)->label() <= 4 );
820       CPPUNIT_ASSERT( (*it)->parent() != NULL || (*it)->first_child() != NULL );
821     }
822 
823     CPPUNIT_ASSERT( forest2.checkLeafsOnLocalTree() );
824     CPPUNIT_ASSERT( forest2.checkTree() );
825 
826     forest2.sampleNextGenealogy();
827     CPPUNIT_ASSERT( forest2.checkTree() );
828   }
829 
testCheckForNodeAtHeight()830   void testCheckForNodeAtHeight() {
831     CPPUNIT_ASSERT( forest->checkForNodeAtHeight( 0.0 ) );
832     CPPUNIT_ASSERT( forest->checkForNodeAtHeight( 1.0 ) );
833     CPPUNIT_ASSERT( forest->checkForNodeAtHeight( 3.0 ) );
834     CPPUNIT_ASSERT( forest->checkForNodeAtHeight( 4.0 ) );
835     CPPUNIT_ASSERT( forest->checkForNodeAtHeight( 6.0 ) );
836     CPPUNIT_ASSERT( forest->checkForNodeAtHeight( 10.0 ) );
837 
838     CPPUNIT_ASSERT( !forest->checkForNodeAtHeight( 0.5 ) );
839     CPPUNIT_ASSERT( !forest->checkForNodeAtHeight( 1.5 ) );
840     CPPUNIT_ASSERT( !forest->checkForNodeAtHeight( 2.0 ) );
841     CPPUNIT_ASSERT( !forest->checkForNodeAtHeight( 9.0 ) );
842     CPPUNIT_ASSERT( !forest->checkForNodeAtHeight( 20.0 ) );
843   }
844 
testPrintLocusSumStats()845   void testPrintLocusSumStats() {
846     ostringstream output;
847     forest->writable_model()->addSummaryStatistic(std::make_shared<TMRCA>());
848     forest->set_current_base(0);
849     forest->set_next_base(forest->model().loci_length());
850     forest->calcSegmentSumStats();
851     forest->printLocusSumStats(output);
852     CPPUNIT_ASSERT( output.str() != "" );
853   }
854 
testSampleNextPosition()855   void testSampleNextPosition() {
856     forest->createScaledExampleTree();
857     forest->writable_model()->setRecombinationRate(0.0);
858     forest->writable_model()->setRecombinationRate(1.0, false, false, 3);
859     forest->sampleNextBase();
860     CPPUNIT_ASSERT_EQUAL(3.0, forest->next_base());
861     CPPUNIT_ASSERT_EQUAL(1.0, forest->model().recombination_rate());
862   }
863 
testClear()864   void testClear() {
865     forest->createScaledExampleTree();
866     forest->writable_model()->setRecombinationRate(0.0);
867     forest->writable_model()->setRecombinationRate(1.0, false, false, 3);
868     forest->writable_model()->addSymmetricMigration(1.0, 0.5);
869     forest->writable_model()->increaseSequencePosition();
870     forest->writable_model()->increaseTime();
871     CPPUNIT_ASSERT_EQUAL(3.0, forest->model().getCurrentSequencePosition());
872     CPPUNIT_ASSERT_EQUAL(1.0, forest->model().getCurrentTime());
873 
874     forest->clear();
875     CPPUNIT_ASSERT_EQUAL((size_t)0, forest->nodes()->size());
876     CPPUNIT_ASSERT_EQUAL((size_t)1, forest->rec_bases_.size());
877     CPPUNIT_ASSERT_EQUAL((size_t)0, forest->segment_count());
878     CPPUNIT_ASSERT_EQUAL(-1.0, forest->current_base());
879     CPPUNIT_ASSERT_EQUAL(0.0, forest->model().getCurrentSequencePosition());
880     CPPUNIT_ASSERT_EQUAL(0.0, forest->model().getCurrentTime());
881   }
882 };
883 
884 
885 //Uncomment this to activate the test
886 CPPUNIT_TEST_SUITE_REGISTRATION( TestForest );
887