Repository: tskit-dev/pyslim

Issue:
recapitate has surprisingly small diversity if WF simulation is started and saved in `late`
As pointed out by @bodkan in https://github.com/bodkan/slendr/issues/141 - 
with this as `late.slim`:
```
initialize() {
  initializeTreeSeq();
  initializeMutationRate(0);
  initializeMutationType("m0", 0.5, "f", 0.0);
  initializeGenomicElementType("g1", m0, 1.0);
  initializeGenomicElement(g1, 0, 1e+07 - 1);
  initializeRecombinationRate(1e-08);
}
1 late() { // <<<------ NOTE THIS LINE --------
  sim.addSubpop("p0", 1000);
}
1000 late() {
  sim.treeSeqOutput("late.trees");
}
```
and this as `early.slim`:
```
initialize() {
  initializeTreeSeq();
  initializeMutationRate(0);
  initializeMutationType("m0", 0.5, "f", 0.0);
  initializeGenomicElementType("g1", m0, 1.0);
  initializeGenomicElement(g1, 0, 1e+07 - 1);
  initializeRecombinationRate(1e-08);
}
1 early() { // <<<------ NOTE THIS LINE --------
  sim.addSubpop("p0", 1000);
}
1000 late() {
  sim.treeSeqOutput("early.trees");
}
```
after running
```
for x in early late; do slim -s 23 ${x}.slim; done
```
then the following script
```
import msprime, tskit, pyslim

Ne = 1000    # constant Ne of the population
rho  = 1e-8  # recombination rate

ts_early = tskit.load("early.trees")
ts_late = tskit.load("late.trees")

def recap(ts):
    return pyslim.recapitate(ts, ancestral_Ne=Ne, recombination_rate=rho)

div_early = recap(ts_early).diversity(mode='branch')
div_late = recap(ts_late).diversity(mode='branch')

print(f"early: {div_early}, late: {div_late}")
```
gets something like
```
early: 4308.810302, late: 3347.565959
```

And, this is because the two tree sequences have the same `tick` value in metadata, but roots at different times ago:
```
>>> ts_early.metadata['SLiM']['tick'], ts_early.node(ts_early.first().roots[0]).time
(1000, 1000.0)
>>> ts_late.metadata['SLiM']['tick'], ts_late.node(ts_late.first().roots[0]).time
(1000, 999.0)
```
while `pyslim.recapitate` contains these lines:
```
        demography = msprime.Demography.from_tree_sequence(ts)
        # must set pop sizes to >0 even though we merge immediately
        for pop in demography.populations:
            pop.initial_size=1.0
# ...
        # the split has to come slightly longer ago than slim's tick
        # since that's when all the linages are at, and otherwise the event
        # won't apply to them
        demography.add_population_split(
                np.nextafter(
                    ts.metadata['SLiM']['tick'],
                    2 * ts.metadata['SLiM']['tick'],
                ),
                derived=derived_names,
                ancestral=ancestral_name,
        )
```

So, what's happening is that I'm creating a bottleneck with Ne=1 for 1 generation in the case that we started and saved in `late()`. For why this happens, see [the docs](https://tskit.dev/pyslim/docs/latest/time_units.html#converting-from-slim-time-to-tskit-time-and-back).

Discussion:
petrelharp: Might you have a look at this, @bodkan? Warning: understanding what's going on here requires the painfully-acquired info in [this section](https://tskit.dev/pyslim/docs/latest/time_units.html#converting-from-slim-time-to-tskit-time-and-back).
petrelharp: Note: if the roots aren't at the expected time in the first tree, it'll do something like this:
```
  /home/peter/projects/tskit-dev/pyslim/pyslim/methods.py:80: RootTimesMismatchWarning: Not all roots of the provided tree sequence are at the time expected by recapitate(). This could happen if you've simplified before recapitating, or if you added new individuals without parents in SLiM during the course of the simulation (e.g., with sim.addSubPop()). If you wish to suppress this warning, you can use, e.g., warnings.simplefilter('ignore', pyslim.RootTimesMismatchWarning) Expected root time: 10 Observed root times: [8.0] 
    warnings.warn(message, RootTimesMismatchWarning)
```

I'm only checking the *first* tree since (a) that will always catch the sort of error we have here; and (b) checking the roots of *all* trees could take a while.
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/tskit-dev/pyslim/pull/308?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) Report
> Merging [#308](https://app.codecov.io/gh/tskit-dev/pyslim/pull/308?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) (7de997f) into [main](https://app.codecov.io/gh/tskit-dev/pyslim/commit/a47c6fe42c608ded8c5af139e0343499133cbbca?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) (a47c6fe) will **increase** coverage by `0.21%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##             main     #308      +/-   ##
==========================================
+ Coverage   95.43%   95.65%   +0.21%     
==========================================
  Files           6        8       +2     
  Lines         592      622      +30     
  Branches      127      131       +4     
==========================================
+ Hits          565      595      +30     
  Misses         18       18              
  Partials        9        9              
```


| [Impacted Files](https://app.codecov.io/gh/tskit-dev/pyslim/pull/308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) | Coverage Δ | |
|---|---|---|
| [pyslim/methods.py](https://app.codecov.io/gh/tskit-dev/pyslim/pull/308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev#diff-cHlzbGltL21ldGhvZHMucHk=) | `97.48% <100.00%> (+0.08%)` | :arrow_up: |

... and [2 files with indirect coverage changes](https://app.codecov.io/gh/tskit-dev/pyslim/pull/308/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev)

petrelharp: Here's the info on when exactly this bug happens: format is {type of model}_{init stage}_{save stage}: <the value of metadata.tick> <the time of the roots>
```
nonWF_early_early.trees: metadata=10 root=9.0
nonWF_early_first.trees: metadata=10 root=8.0
nonWF_early_late.trees: metadata=10 root=9.0
nonWF_first_early.trees: metadata=10 root=10.0
nonWF_first_first.trees: metadata=10 root=9.0
nonWF_first_late.trees: metadata=10 root=10.0
nonWF_late_early.trees: metadata=10 root=9.0
nonWF_late_first.trees: metadata=10 root=8.0
nonWF_late_late.trees: metadata=10 root=9.0
WF_early_early.trees: metadata=10 root=9.0
WF_early_first.trees: metadata=10 root=9.0
WF_early_late.trees: metadata=10 root=10.0
WF_first_early.trees: metadata=10 root=9.0
WF_first_first.trees: metadata=10 root=9.0
WF_first_late.trees: metadata=10 root=10.0
WF_late_early.trees: metadata=10 root=8.0
WF_late_first.trees: metadata=10 root=8.0
WF_late_late.trees: metadata=10 root=9.0
```
The takeaway is that we seem to only **not** have the bug in:
- nonWF/first/{early,late}
- WF/{first,early}/late

In particular, it does affect sims where the init and emit stages were the same.
bhaller: > The takeaway is that we seem to only **not** have the bug in:
> 
> * nonWF/first/{early,late}
> * WF/{first,early}/late
> 
> In particular, it does affect sims where the init and emit stages were the same.

Ouch.  Since `first()` events are not used very often (since they are a pretty new feature), this means that probably the bug bites almost everyone with nonWF models, and bites everyone with WF models except init in `early()`, save in `late()`.
bodkan: > Might you have a look at this, @bodkan?

Sure thing! I'm happy to take a look and re-run some of the slendr-based models and tests (the ones which I later chiseled down to the `early.slim` and `late.slim` equivalents), just to make sure that
a) what worked before still works (no reason it shouldn't, I would think),
b) the models that returned broken results earlier are fixed now.

>  Warning: understanding what's going on here requires the painfully-acquired info in [this section](https://tskit.dev/pyslim/docs/latest/time_units.html#converting-from-slim-time-to-tskit-time-and-back).

I actually think it's about time I finally read what is actually going on at the lower levels. 🙂 Thanks for the link.

**Also:** Wow, that was extremely fast. I guess once we verified that the `early` and `late` versions do indeed give different results you probably knew quite well where to look, but still -- wow! 👏
petrelharp: Another reason we can't just figure out what the time is supposed to be based on metadata (even if the information was in there) - if it's produced by making and annotatingg an msprime tree sequence, it probably won't have the offset we'd expect.

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index ca38ea1..7d9c441 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -2,6 +2,22 @@
 [UPCOMING.X.X] - XXXX-XX-XX
 ***************************
 
+**Bugfixes**:
+
+- From 1.0.1 back to 0.700, there was a bug in `recapitate` when using the
+    `ancestral_Ne` parameter that introduced a bottleneck to diploid size Ne=1
+    for each SLiM subpopulation for 1 or 2 generations *unless* either (a) it
+    was a WF simulation, with calls to addSubPop() in first() or early() and
+    treeSeqOutput() in late(), or (b) it was a nonWF simulation, with calls to
+    addSubPop() in first() and treeSeqOutput() in early() or late(). The fix
+    correctly starts the msprime population with effective size `ancestral_Ne`
+    at the time of the roots, which might be at the value of
+    `ts.metadata['SLiM']['tick']`, this value minus 1, or this value minus 2.
+    Furthermore, `recapitate` now throws a warning if any roots of any trees
+    are not at the same time as the others. (:user:`petrelharp`, :pr:`308`)
+
+
+
 ***************************
 [1.0.1] - 2022-09-23
 ***************************
diff --git a/pyslim/methods.py b/pyslim/methods.py
index 6d67d50..68e50db 100644
--- a/pyslim/methods.py
+++ b/pyslim/methods.py
@@ -8,6 +8,7 @@ from .slim_metadata import *
 from .provenance import *
 from .util import *
 
+
 def recapitate(ts,
                ancestral_Ne=None,
                **kwargs
@@ -48,6 +49,32 @@ def recapitate(ts,
     if ancestral_Ne is not None:
         if "demography" in kwargs:
             raise ValueError("You cannot specify both `demography` and `ancestral_Ne`.")
+        recap_time = ts.metadata['SLiM']['tick']
+        # In various circumstances depending on the stage in which the simulation was
+        # started and when the tree sequence was written out, the time of the roots
+        # might be one or even two less than the "tick" value. We have access
+        # to the stage the tree sequence was written out, but not the time it
+        # was *started*. So, we'll just check if we need to subtract one or two.
+        # Consistency checking this requires looping over all the trees, unfortunately,
+        # but it avoids some common errors.
+        root_times = list(set([ts.node(n).time for t in ts.trees() for n in t.roots]))
+        for adj in (1, 2):
+            if np.allclose(root_times, recap_time - adj):
+                recap_time -= adj
+        if len(root_times) > 1 or not np.allclose(root_times, recap_time):
+            message = (
+                "Not all roots of the provided tree sequence are at the time expected "
+                "by recapitate(). This could happen if you've simplified in "
+                "python before recapitating (fix: don't simplify first). "
+                "If could also happen in other situations, e.g., "
+                "you added new individuals without parents in SLiM "
+                "during the course of the simulation with sim.addSubPop(), "
+                "in which case you will probably need to recapitate with "
+                "msprime.sim_ancestry(initial_state=ts, ...). "
+                f"(Expected root time: {recap_time}; "
+                f"Observed root times: {root_times})"
+            )
+            raise ValueError(message)
         demography = msprime.Demography.from_tree_sequence(ts)
         # must set pop sizes to >0 even though we merge immediately
         for pop in demography.populations:


Test patch:
diff --git a/tests/test_tree_sequence.py b/tests/test_tree_sequence.py
index 8d1ca41..9e3eefc 100644
--- a/tests/test_tree_sequence.py
+++ b/tests/test_tree_sequence.py
@@ -142,8 +142,18 @@ class TestRecapitate(tests.PyslimTestCase):
                         random_seed=123,
             )
 
-    def test_unique_names(self):
+    def test_root_mismatch_error(self):
         ts = msprime.sim_ancestry(4, sequence_length=10, random_seed=12)
+        recap_time = 100
+        ts = pyslim.annotate(ts, model_type="nonWF", tick=recap_time)
+        assert ts.node(ts.first().roots[0]).time < recap_time
+        with pytest.raises(ValueError, match="at the time expected by recapitate"):
+            rts = self.do_recapitate(ts, ancestral_Ne=10)
+
+    def test_unique_names(self):
+        ts = msprime.sim_ancestry(4, sequence_length=10, random_seed=12, end_time=1.0)
+        # adjust seed if not
+        assert min([t.num_roots for t in ts.trees()]) > 1
         ts = pyslim.annotate(ts, model_type="nonWF", tick=1)
         t = ts.dump_tables()
         md = t.populations[0].metadata