add `.individuals` property to `Population` object?

It is a very common thing to need to (a) first go through the individuals in a tree sequence and find their population, and then (b) subsample according to population. Right now, we might get the list of individuals by population like so:

pop_inds = [[] for _ in range(ts.num_populations)]
for ind in ts.individuals():
  pops = [ts.node(n).population for n in ind.nodes]
  for pop in set(pops):
      pop_inds[pop].append(ind.id)

… and I’ve omitted the error checking enuring no individual has nodes in more than one population. This would be easier if we could do

pop_inds = [pop.individuals for pop in ts.populations()]

This suggestion is somewhat fraught because of the possibility that an individual’s nodes may have different populations (although if they do, it’s probably an error). So, I propose throwing an error if any individual’s nodes have more than one population associated. (This check could be an optional argument, but I can’t think of when this wouldn’t be an error.)

[isaacovercast]

That’s great, thanks @petrelharp, exactly what I needed!

[isaacovercast]

+1

Hello tskit team. Looking for a way to get diversity per population from an msprime simulation. Peter’s code above is already a huge help as it looks like just what I need, but I agree it would be useful to simplify in one of the two ways suggested here. From a user perspective, I think it’s more natural for individuals to be a property of populations, but either way would be fine.

Long time msprime/tskit user here. Thanks for this amazing resource and all your hard work!