Coverage for rta_reconstruction/dl2_io.py: 90%

21 statements  

« prev     ^ index     » next       coverage.py v7.6.4, created at 2024-11-02 09:59 +0000

1from pathlib import Path 

2 

3import tables 

4from pandas import DataFrame 

5 

6 

7def init_dl2_file(dl2_file_path: Path, dl1_file_path: Path): 

8 if dl2_file_path.exists(): 

9 raise FileExistsError("{} already exists, refusing to overwrite it.".format(dl2_file_path)) 

10 

11 # TODO: make these configurable somehow 

12 not_copy_list = [ 

13 "/dl1/event/telescope/image", 

14 "/dl1/event/telescope/image/LST_LSTCam", 

15 "/dl1/event/telescope/parameters", 

16 "/dl1/event/telescope/parameters/LST_LSTCam", 

17 "/dl1/event/telescope/parameters_src_dependent/LST_LSTCam", 

18 "/dl1/event/telescope/likelihood_parameters/LST_LSTCam", 

19 ] 

20 

21 with tables.open_file(dl1_file_path, "r") as dl1_f, tables.open_file(dl2_file_path, "w") as dl2_f: 

22 for node in dl1_f.walk_nodes(): 

23 if node._v_pathname not in not_copy_list: 

24 if node._v_parent._v_pathname not in dl2_f: 

25 parent = dl2_f.create_group( 

26 node._v_parent._v_parent._v_pathname, node._v_parent._v_name, createparents=True 

27 ) 

28 else: 

29 parent = dl2_f.get_node(node._v_parent._v_pathname) 

30 dl2_f.copy_node(node, parent) 

31 

32 

33# TODO: global metadata in file global attributes ? (soft version, contact) 

34 

35 

36def write_dl2_df( 

37 dl2_file_path: Path, 

38 dl2_df: DataFrame, 

39 dl2_table_path: str = "/dl2/event/telescope/parameters/LST_LSTCam", 

40 attributes=None, 

41 filters=tables.Filters( # TODO: make those configurable, benchmark performances and compare with pandas "to_hdf" 

42 complevel=1, 

43 complib="blosc:zstd", 

44 fletcher32=True, 

45 bitshuffle=False, 

46 ), 

47): 

48 with tables.open_file(dl2_file_path, "a") as dl2_f: 

49 path, table_name = dl2_table_path.rsplit("/", maxsplit=1) # TODO: use pathlib split ? 

50 

51 dl2_table = dl2_f.create_table( 

52 path, 

53 table_name, 

54 dl2_df.to_records(index=False), 

55 createparents=True, 

56 filters=filters, 

57 ) 

58 if attributes is not None: 

59 for key, value in attributes.items(): 

60 dl2_table.attrs[key] = value