Пример #1
0
def test_io_mmap():
    header_size0, nhalos_per_tree0 = lhtutils.read_header_default(SMT)
    data0 = lhtutils._read_from_mmap(SMT, header_size=header_size0)
    with tempfile.NamedTemporaryFile(delete=False) as tmp:
        tmp.close()
        lhtutils._save_to_mmap(tmp.name, data0)
        data1 = lhtutils._read_from_mmap(tmp.name)
        os.remove(tmp.name)
    np.testing.assert_array_equal(data1, data0)
Пример #2
0
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)