Coverage for event_normalizer/serializers/flat_dict.py: 100%

1"""Flat dictionary serialization for normalized events.""" 

2 

3from typing import Any 

4 

5from ..models import NormalizedEvent 

6 

7SUPPORTED_INSTRUMENTS = ("H1", "L1", "V1", "K1") 

8 

9 

10def to_flat_dict( 

11 event: NormalizedEvent, 

12 *, 

13 exclude_none: bool = True, 

14) -> dict[str, Any]: 

15 """ 

16 Convert a normalized event into a one-level dictionary. 

17 

18 Nested models already use prefixed fields: 

19 - coinc.coinc_mass -> coinc_mass 

20 - burst.burst_snr -> burst_snr 

21 - mly.mly_SNR -> mly_SNR 

22 - single_H1.single_snr -> single_H1_single_snr 

23 - labels.label_MOCK -> label_MOCK 

24 

25 Detector lists become numeric metric fields: 

26 - ["H1", "L1"] -> 

27 instruments_H1 = 1 

28 instruments_L1 = 1 

29 instruments_V1 = 0 

30 instruments_K1 = 0 

31 """ 

32 raw = event.model_dump( 

33 exclude_none=exclude_none, 

34 mode="json", 

35 ) 

36 

37 result: dict[str, Any] = {} 

38 

39 nested_keys = { 

40 "coinc", 

41 "burst", 

42 "mly", 

43 "single_H1", 

44 "single_L1", 

45 "single_V1", 

46 "single_K1", 

47 "labels", 

48 "p_astro", 

49 "em_bright", 

50 "links", 

51 } 

52 

53 for key, value in raw.items(): 

54 if key in nested_keys: 

55 continue 

56 

57 if key == "instruments": 

58 actual_instruments = set(value or []) 

59 

60 for instrument in SUPPORTED_INSTRUMENTS: 

61 result[f"instruments_{instrument}"] = int( 

62 instrument in actual_instruments 

63 ) 

64 

65 continue 

66 

67 result[key] = value 

68 

69 for nested_key in nested_keys: 

70 nested_value = raw.get(nested_key) 

71 

72 if isinstance(nested_value, dict): 

73 # For single_H1, single_L1, etc., prefix with the IFO name 

74 if nested_key.startswith("single_"): 

75 for sub_key, sub_value in nested_value.items(): 

76 result[f"{nested_key}_{sub_key}"] = sub_value 

77 else: 

78 result.update(nested_value) 

79 

80 return result