From d46c7ce2cc64909abac141fcedf8be391c3491b9 Mon Sep 17 00:00:00 2001 From: Yan Wong Date: Wed, 17 Jul 2024 09:50:53 +0100 Subject: [PATCH] Coerce NODE_IS_SAMPLE to correct numpy type before negation Fixes #937 --- tsinfer/inference.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tsinfer/inference.py b/tsinfer/inference.py index e8dcb876..0b1baba8 100644 --- a/tsinfer/inference.py +++ b/tsinfer/inference.py @@ -2293,7 +2293,10 @@ def get_samples_tree_sequence(self, map_additional_sites=True): # All true ancestors are samples in the ancestors tree sequence. We unset # the SAMPLE flag but keep other flags intact. - new_flags = np.bitwise_and(tables.nodes.flags, ~tskit.NODE_IS_SAMPLE) + new_flags = tables.nodes.flags + new_flags = np.bitwise_and( + new_flags, ~new_flags.dtype.type(tskit.NODE_IS_SAMPLE) + ) tables.nodes.flags = new_flags.astype(np.uint32) sample_ids = list(self.sample_id_map.values()) assert len(tables.nodes) == sample_ids[0]