How to get a numpy array from a ragged column e.g. site ancestral states

[hyanwong]

What's the most efficient way to get the ancestral states of all the sites in a tree sequence as a numpy array of strings? I suspect the following is very inefficient:

np.array([site.ancestral_state for site in ts.sites()])

I couldn't find this documented anywhere, and since the array is ragged, there isn't a site_ancestral_state array available (yet)

[hyanwong]

Ah, I forgot about tskit.unpack_strings(t.ancestral_state, t.ancestral_state_offset). Presumably this is the recommended way? But it would require dumping all the tables first, which is not ideal.

[hyanwong]

It looks like np.split should do it too, but the .tobytes() call still needs to be run in a loop:

byte_arrays = [chunk.tobytes() for chunk in np.split(sites.ancestral_state, sites.ancestral_state_offset[1:-1])]

or

string_arrays = [chunk.tobytes().decode() for chunk in np.split(sites.ancestral_state, sites.ancestral_state_offset[1:-1])]

[hyanwong]

Huh - I just timed this and the np.split approach is 2 times slower than the python loop within tskit.unpack_strings(). So looks like the built-in tskit method will be the most efficient.

[jeromekelleher]

Related: #2632. I'm hoping that we can implement this efficiently now with the numpy 2.0 StringDtype (but will require C coding)