Work around to allow Demography init from provenance with preset population IDs

Author: hyanwong

CHANGELOG.md

@@ -13,6 +13,7 @@
 - Support for simulations after local mrca,
   ({issue}`2157`, {pr}`2396`, {user}`jeromekelleher`, {user}`hossam26644`).
 - Add wheels on Windows ({pr}`2414`, {issue}`2200`,{user}`benjeffery`)
+- Demography objects can now be created from provenance entries ({pr}`{2369}`, {user}`hyanwong`)
 
  **Performance improvements**
 

msprime/demography.py

@@ -292,10 +292,10 @@ def __post_init__(self):
 
         # Assign the IDs and default names, if needed.
         for j, population in enumerate(self.populations):
-            if population.id is not None:
+            if population.id is not None and population.id != j:
                 raise ValueError(
-                    "Population ID should not be set before using to create "
-                    "a Demography"
+                    "Population ID should be unset or set to its index within"
+                    f" the Demography (expected None or {j}, got {population.id})"
                 )
             population.id = j
             if population.name is None:

tests/test_demography.py

@@ -4312,9 +4312,11 @@ def test_duplicate_populations(self):
         with pytest.raises(ValueError, match="must be distinct"):
             msprime.Demography([pop] * 2)
 
-    def test_population_ids_set_on_init(self):
-        pop = msprime.Population(10, id=0)
-        with pytest.raises(ValueError, match="ID should not be set"):
+    def test_population_ids_bad_on_init(self):
+        pop = msprime.Population(10, id=1)
+        with pytest.raises(
+            ValueError, match="ID should be unset or set to its index.*got 1"
+        ):
             msprime.Demography([pop])
 
     def test_add_population_error(self):

tests/test_provenance.py

@@ -244,6 +244,39 @@ def test_current_ts(self):
         assert command == "sim_mutations"
         assert prov["tree_sequence"] == ts1
 
+    def test_demography(self):
+        demography = msprime.Demography.island_model([1, 1], 1 / 3)
+        ts = msprime.sim_ancestry(
+            demography=demography,
+            samples=[
+                msprime.SampleSet(1, population=0),
+                msprime.SampleSet(1, population=1),
+            ],
+            random_seed=3,
+        )
+        command, prov = msprime.provenance.parse_provenance(ts.provenance(-1), ts)
+        assert command == "sim_ancestry"
+        assert prov["demography"] == demography
+
+    def test_bad_demography(self):
+        demography = msprime.Demography.island_model([1, 1], 1 / 3)
+        ts = msprime.sim_ancestry(
+            demography=demography,
+            samples=[
+                msprime.SampleSet(1, population=0),
+                msprime.SampleSet(1, population=1),
+            ],
+            random_seed=3,
+        )
+        prov = ts.provenance(-1)
+        record = json.loads(prov.record)
+        # Corrupt the provenance record
+        assert len(record["parameters"]["demography"]["migration_matrix"]) == 2
+        record["parameters"]["demography"]["migration_matrix"] = [1, 0, 0]
+        prov.record = json.dumps(record)
+        with pytest.raises(ValueError, match="Migration matrix must be square"):
+            msprime.provenance.parse_provenance(prov, ts)
+
 
 class TestRoundTrip:
     """