def test_io_header_default():
header_size0, nhalos_per_tree0 = lhtutils.read_header_default(SMT)
with tempfile.NamedTemporaryFile(delete=False) as tmp:
_header_size = lhtutils.save_header_default(tmp, nhalos_per_tree0)
assert_equal(_header_size, header_size0)
tmp.close()
header_size1, nhalos_per_tree1 = lhtutils.read_header_default(tmp.name)
os.remove(tmp.name)
assert_equal(header_size1, header_size0)
np.testing.assert_array_equal(nhalos_per_tree1, nhalos_per_tree0)
def test_io_file():
header_size0, nhalos_per_tree0 = lhtutils.read_header_default(SMT)
data0 = lhtutils._read_from_file(SMT, header_size=header_size0)
with tempfile.NamedTemporaryFile(delete=False) as tmp:
tmp.close()
lhtutils._save_to_file(tmp.name, data0)
data1 = lhtutils._read_from_file(tmp.name)
os.remove(tmp.name)
np.testing.assert_array_equal(data1, data0)
def test_io_trees_default():
header_size0, nhalos_per_tree0 = lhtutils.read_header_default(SMT)
data0 = lhtutils.read_trees_default(SMT, header_size=header_size0)
with tempfile.NamedTemporaryFile(delete=False) as tmp:
tmp.close()
lhtutils.save_trees_default(tmp.name, data0)
data1 = lhtutils.read_trees_default(tmp.name)
os.remove(tmp.name)
for k in data0.keys():
assert (k in data1)
np.testing.assert_array_equal(data1[k], data0[k])
def make_small_example(ntrees_per_file=None):
if ntrees_per_file is None:
ntrees_per_file = np.array([2, 2], 'int')
ntrees_before_file = np.cumsum(ntrees_per_file) - ntrees_per_file
header_size0, nhalos_per_tree0 = lhtutils.read_header_default(MMT)
data0 = lhtutils._read_from_mmap(
MMT, header_size=header_size0,
nhalo=np.sum(nhalos_per_tree0[:np.sum(ntrees_per_file)]))
fname_base = os.path.splitext(SMT)[0]
prev_halos = 0
for i in range(len(ntrees_per_file)):
ifile = f"{fname_base}.{i}"
start_tree = ntrees_before_file[i]
stop_tree = start_tree + ntrees_per_file[i]
nhalos_per_tree1 = nhalos_per_tree0[start_tree:stop_tree]
ntothalos1 = np.sum(nhalos_per_tree1)
data1 = data0[prev_halos:(prev_halos + ntothalos1)]
prev_halos += ntothalos1
header_size1 = lhtutils.save_header_default(ifile, nhalos_per_tree1)
lhtutils._save_to_mmap(ifile, data1, header_size=header_size1)