tskit-dev / tskit

Population-scale genomics
MIT License
147 stars 69 forks source link

Pass a numpy array of booleans to python simplify? #2968

Closed hyanwong closed 1 week ago

hyanwong commented 1 week ago

For teaching purposes, it's nice to be able to do

ts.simplify(ts.nodes_time == 0)

rather than

ts.simplify(np.where(ts.nodes_time == 0)[0])

What do people think of the idea of special-casing the samples argument to simplify() to check if it is of length num_nodes and of dtype=bool, and then simply doing a np.where(samples)[0]? At the moment this fails with a Duplicate sample value error, as it treats the 0s and 1s as sample IDs.

jeromekelleher commented 1 week ago

It's a bunch of work to do, and not that much different? Seems to me that you'd have to explain more by adding the boolean selector for nodes as well as the list of nodes option.

hyanwong commented 1 week ago

It's only a 5 line addition at the top of tables.simplify(), isn't it?

try:
    if len(samples) == ts.num_nodes and samples.dtype == bool:
        samples = np.where(samples)[0]
except AttributeError:
   pass

You're right that it's not that different, and it's not really a priority, but I'm finding that every extra barrier to new tskit users puts some of them off (and takes an extra few minutes to explain). There's a reason why numpy allows both boolean and numerical indexing. I'm not sure I would actually explain in a practical that you can use both: it's obvious from the context, right?

jeromekelleher commented 1 week ago

As usual with these things, implementation is by far the easiest thing and something like 5% of the actual effort. Testing, documenting and making sure there are no regressions against existing code make the majority of the work.

petrelharp commented 1 week ago

I agree. Higher on my list would be adding an individuals argument.

hyanwong commented 1 week ago

I guess part of the problem is nothing to do with tskit, but the very unintuitive np.where syntax which requires the [0] at the end (which I always forget). The other option is

ts.simplify(np.flatnonzero(ts.nodes_time == 0))

but that's hardly more intuitive, IMO. Or alternatively

ts.simplify((ts.nodes_time == 0).nonzero()[0])

Which is equally cryptic, but at least doesn't need you to use the np prefix, which again requires explanation when it's in the first few lines of a beginner's tutorial.

Anyway, if no-one thinks it's a good idea, I'll close this. However, it's worth pointing out that I'm finding it hard to introduce tskit to new users (especially from R, and if they don't know numpy).

petrelharp commented 1 week ago

Another option is

today_nodes = np.arange(ts.num_nodes)[ts.nodes_time == 0]
ts.simplify(today_nodes)
petrelharp commented 1 week ago

FWIW, I can't think of a single function in R that overloads arguments like that - i.e., takes either a vector of indices or a vector of booleans. Indexing works like that, but only indexing. Overloading can lead to nasty corner cases.