vectorised metadata extraction

[petrelharp]

A common thing people - at least with SLiM tree sequences - want to do is "get a vector of values from metadata where every object has the thing in metadata", e.g.

np.array([ind.metadata['slim_time'] for ind in ts.individuals()])

We do in pyslim, and @bodkan does this also (somewhat painfully?) in slendr using reticulate. The list comprehension is pretty easy to do, but to make it easier, perhaps there's an API we could make up? like

ts.metadata_vector('individuals', 'slim_time')

for instance? Perhaps with an optional default argument that provides the value to be used when the key is not present in a metadata?

[jeromekelleher]

I think this is a useful thing to have in general.

It would be slightly nicer I think if we did this on the TableCollection, because we could then do stuff like

slim_time = tables.individuals.metadata_vector("slim_time")

If we're working on vectorised representations, then we're probably working with the table columns anyway (rather than the tree sequence methods)?

The full signature of the Table method might look like:

def metadata_vector(key, *, dtype=None, shape=None, default_value=None):
    ...

I guess the simplest behaviour is to let np.array decide what the dtype and shape are, which can be overridden with the parameters if the user wants. If the default_value is provided, that fills in any gaps, otherwise we raise an error. We'd have to be careful about how this interacts with the defaults specified in the schemas, though. Also, None might not be a good choice as the value in the signature, as the array might be an object array where None would be a valid value. We should probably make an "UNSPECIFIED" singleton for this, I guess.

What do you think @benjeffery?

[benjeffery]

Great idea, we can implement as the simple list comprehension now, and if perf becomes and issue we can, in future, do things like only decoding the bytes that specify the requested key.

I agree that this makes more sense on a table collection. One other detail to consider is if key is allowed to be a list of strings for accessing metadata that is deeper than the top level. I think this would be useful.

The current metadata codecs transparently deal with defaults from the schema, so any default specified here would only work if there was no default specified in the schema. A singleton does sound best for that as None is often a valid metadata value.

[petrelharp]

One remaining question is whether this should return a numpy array or just a list - I'm thinking probably the latter, as otherwise types that aren't numpy-array-able (eg lists, I think?) will do bad stuff. So, we can let the user wrap np.array( ) around the result, which is no big deal?

[jeromekelleher]

I think it's better to make it a numpy array actually, because then we have the option of doing things efficiently behind the scenes at a later date if we want (suppose we wanted to decode some simple binary encoded data into arrays - we could do this at machine speed in the Python C module). If we tie ourselves to lists, this can never really scale.

I think np.array() is pretty flexible these days - what sort of thing do you think it'll choke on?

[petrelharp]

Ah-ha, you're right, I see there's a dtype="object":

>>> np.array([list(range(k)) for k in range(4)])
<stdin>:1: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray
array([list([]), list([0]), list([0, 1]), list([0, 1, 2])], dtype=object)
>>> np.array([{k: k} for k in range(4)])
array([{0: 0}, {1: 1}, {2: 2}, {3: 3}], dtype=object)

Note that inserting None gets you this dtype (and so we'd want to avoid this for numerical things):

>>> np.array([0,3,4,None])
array([0, 3, 4, None], dtype=object)

So, is the proposal to just wrap np.array( ) around the list comprehension, provide an optional dtype argument, and let the user sort out awkward cases?

[jeromekelleher]

So, is the proposal to just wrap np.array( ) around the list comprehension, provide an optional dtype argument, and let the user sort out awkward cases?

Exactly. We can optimise common use cases later, but this should be the behaviour.