def verify_num_sites(self, tree_seq, tree_seq2, num_replicates):
if isinstance(tree_seq, collections.abc.Iterable):
with tempfile.TemporaryDirectory() as temp_dir:
ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
with open(ms_file_path, "w") as f:
ts.write_ms(
tree_seq,
f,
num_replicates=num_replicates,
)
with open(ms_file_path) as handle:
num_sites = get_ms_file_quantity(handle, "num_sites")
with open(ms_file_path) as handle:
num_positions = get_ms_file_quantity(
handle, "num_positions")
j = 0
for ts_indv in tree_seq2:
self.assertEqual(ts_indv.num_sites, num_sites[j])
self.assertEqual(ts_indv.num_sites, num_positions[j])
j = j + 1
else:
with tempfile.TemporaryDirectory() as temp_dir:
ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
with open(ms_file_path, "w") as f:
ts.write_ms(tree_seq, f)
with open(ms_file_path) as handle:
num_sites = get_ms_file_quantity(handle, "num_sites")
with open(ms_file_path) as handle:
num_positions = get_ms_file_quantity(
handle, "num_positions")
self.assertEqual(tree_seq.num_sites, num_sites[0])
self.assertEqual(tree_seq.num_sites, num_positions[0])
def verify_num_replicates(self, tree_seq, num_replicates):
if isinstance(tree_seq, collections.abc.Iterable):
with tempfile.TemporaryDirectory() as temp_dir:
ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
with open(ms_file_path, "w") as f:
ts.write_ms(
tree_seq,
f,
num_replicates=num_replicates,
)
with open(ms_file_path) as handle:
num_replicates_file = get_ms_file_quantity(
handle, "num_replicates")
else:
with tempfile.TemporaryDirectory() as temp_dir:
ms_file_path = os.path.join(temp_dir, "testing_ms_file.txt")
with open(ms_file_path, "w") as f:
ts.write_ms(tree_seq, f)
with open(ms_file_path) as handle:
num_replicates_file = get_ms_file_quantity(
handle, "num_replicates")
self.assertEqual(num_replicates, num_replicates_file)