Can I extract an msprime.Demography object from the provenance of a tree sequence?

[hyanwong]

e.g. json.loads(ts.provenance(-1).record)['parameters']['demography'] could give something like the following. Is there any built-in way to turn that back into an msprime.Demography object (for plotting in demes)?

{
'populations': [
    {'initial_size': 10000, 'growth_rate': 0, 'name': 'ancestral', 'description': '', 'extra_metadata': {}, 'default_sampling_time': 2000, 'initially_active': False, 'id': 0}, 
    {'initial_size': 10000, 'growth_rate': 0, 'name': 'pop1', 'description': '', 'extra_metadata': {}, 'default_sampling_time': None, 'initially_active': None, 'id': 1}, 
    {'initial_size': 5000, 'growth_rate': 0, 'name': 'pop2', 'description': '', 'extra_metadata': {}, 'default_sampling_time': None, 'initially_active': None, 'id': 2}
  ],
  'events': [
    {'time': 2000, 'derived': ['pop1', 'pop2'], 'ancestral': 'ancestral'}
  ],
  'migration_matrix': [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
  '__class__': 'msprime.demography.Demography'
}

[jeromekelleher]

Try using this function.

[hyanwong]

Perfect! Thanks. I couldn't find that in the sim_ancestry doc page, or an internet search, but hopefully this issue will put it into the internet hive mind.

[hyanwong]

Ah, but that function gives an error. I think the "population"s need converting into an msprime.Population object?

File ~/python3.12/site-packages/msprime/demography.py:280, in Demography.post_init(self)
278 # Assign the IDs and default names, if needed.
279 for j, population in enumerate(self.populations):
--> 280 if population.id is not None: AttributeError: 'dict' object has no attribute 'id'"

demography = msprime.Demography()
demography.add_population(name="ancestral", initial_size=10000)
demography.add_population(name="pop1", initial_size=10000)
demography.add_population(name="pop2", initial_size=5000)
demography.add_population_split(time=2000, derived=["pop1", "pop2"], ancestral="ancestral")

# Run the simulation with this model
ts = msprime.sim_ancestry(
    samples={"pop1": 5, "pop2": 5},
    demography=demography,
    random_seed=42
)

msprime.provenance.parse_provenance(ts.provenance(-1), None)

[jeromekelleher]

Yeah there's a reason we decided not to document this and let go of the goal of fully round-tripping stuff from provenance. I'm sure it's possible, but you'll need to get fiddly I'd imagine.

[hyanwong]

For future readers, updates to msprime now means that this works natively (i.e. the code above works fine)

[hyanwong]

So there are 2 problems with the function mentioned in #2367 (comment) (above). I'm just noting them here for future reference.

Firstly, the __class__ of the internal dictionaries isn't set in the provenance, so we can't really know what e.g. the sort of event was. What we really want to do when outputting the provenance is for ProvenanceEncoderDecoder().encode() to recursively print the __class__ on each dict. E.g. at the moment we have:

{
  'populations': [{'initial_size': 10000, 'growth_rate': 0, 'name': 'ancestral', ... 'id': 0},
     ...
  ],
  'events': [{'time': 2000, 'derived': ['pop1', 'pop2'], 'ancestral': 'ancestral'}],
  '__class__': 'msprime.demography.Demography'
}

But what really should be output is

{
  'populations': [{'__class__': 'msprime.demography.Population', 'initial_size': 10000, 'growth_rate': 0, 'name': 'ancestral', ... 'id': 0},
    ...
  ],
  'events': [{'__class__': 'msprime.demography.PopulationSplit', 'time': 2000, 'derived': ['pop1', 'pop2'], 'ancestral': 'ancestral'}],
  '__class__': 'msprime.demography.Demography'
}

Secondly, the Population() initialiser can't be given an id, but wants to invent one itself. So there's no way of making a population with a given ID, as we would need to do to exactly replicate the Demography model. I think the Demography.copy() function just works around this by re-setting the population IDs once they have all been made, but that's too bespoke for a general converter.

If I modify the JSON to add the correct __class__ keys, and remove the population IDs from the population dictionaries, then the parse_provenance function works fine.