Extracting struct metadata efficiently

[hyanwong]

I have a simple struct schema for node metadata, with 3 numbers:

tables.nodes.metadata_schema = tskit.MetadataSchema({
                "codec": "struct",
                "type": "object",
                "properties": {
                    "mn": {"type": "number", "binaryFormat": "f", "default": float("NaN")},
                    "vr": {"type": "number", "binaryFormat": "f", "default": float("NaN")},
                    "original_id": {"type": "number", "binaryFormat": "i", "default": -1}
                },
    })

As I have 20 replicate tree sequences each with 2.8M nodes, it takes quite a while to extract this. I suspect I can treat all the bytes as a numpy record array, which will be fast. How do I do that, though?

[hyanwong]

In this case, it looks like this will do it (the size of the md bytes, i.e. the `f4', 'i4', etc can be obtained from https://tskit.dev/tskit/docs/stable/metadata.html#numeric-and-boolean-types)

md = ts.tables.nodes.metadata
mapping = {"mn": 'f4', 'vr': 'f4', 'original_id': 'i4'}
dtype = [(k, mapping[k]) for k in ts.node(0).metadata.keys()]  # need this to ensure order is correct
x = np.frombuffer(md, dtype=dtype)
assert x.nbytes == len(md)
x["mn"] * 2  # or whatever

Is there a way of getting the numpy dtype from e.g. ts.node(0).metadata. print([type(v) for k, v in ts.node(0).metadata.items()]) gives the python types which don't tell you the size in bytes.

[benjeffery]

Nice!
You can get the sizes from the schema:
size = struct.calcsize("<"+ schema['"properties"][key_name]["binaryFormat"])