Coverage for src/km3modules/mc.py: 73%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# vim:set ts=4 sts=4 sw=4 et: 

4"""MC Helpers. 

5""" 

6 

7import numpy as np 

8 

9from km3pipe import Module 

10from km3pipe.mc import pdg2name, convert_mc_times_to_jte_times 

11from km3pipe.math import zenith, azimuth 

12 

13__author__ = "Michael Moser and Tamas Gal and Moritz Lotze" 

14 

15NEUTRINOS = { 

16 "nu_e", 

17 "anu_e", 

18 "nu_mu", 

19 "anu_mu", 

20 "nu_tau", 

21 "anu_tau", 

22} 

23 

24 

25class GlobalRandomState(Module): 

26 """Sets the global random seed of the numpy random generator 

27 

28 KM3Pipe uses numpy routines exclusively to generate randomness. Setting a 

29 seed to a specific value will create a fully reproducible pipeline as long 

30 as your own modules also utilise numpy for random numbers. 

31 

32 Parameters 

33 ---------- 

34 seed: int, default=42 

35 

36 """ 

37 

38 def configure(self): 

39 self.seed = self.get("seed", default=42) 

40 

41 from numpy.random.mtrand import _rand as global_randstate 

42 

43 global_randstate.seed(self.seed) 

44 

45 self._random_state = np.random.RandomState(self.seed) 

46 

47 @property 

48 def random_state(self): 

49 return self._random_state 

50 

51 

52class McTruth(Module): 

53 """Extract MC info of 1st MC track. 

54 

55 Parameters 

56 ---------- 

57 most_energetic_primary: bool, default=True 

58 """ 

59 

60 def configure(self): 

61 self.most_energetic_primary = bool(self.get("most_energetic_primary")) or True 

62 

63 @classmethod 

64 def t2f(cls, row): 

65 return pdg2name(row["type"]) 

66 

67 @classmethod 

68 def is_nu(cls, flavor): 

69 return flavor in NEUTRINOS 

70 

71 def process(self, blob): 

72 mc = blob["McTracks"].conv_to("pandas") 

73 if self.most_energetic_primary: 

74 mc.sort_values("energy", ascending=False, inplace=True) 

75 mc = mc.head(1) 

76 mc["zenith"] = zenith(mc[["dir_x", "dir_y", "dir_z"]]) 

77 mc["azimuth"] = azimuth(mc[["dir_x", "dir_y", "dir_z"]]) 

78 flavor = mc.apply(self.t2f, axis=1) 

79 mc["is_neutrino"] = flavor.apply(self.is_nu) 

80 blob["McTruth"] = mc 

81 return blob 

82 

83 

84class MCTimeCorrector(Module): 

85 """ 

86 Module that converts MC hit times to JTE times. 

87 Thus, the following tables need to be converted: 

88 - mc_tracks 

89 - mc_hits 

90 

91 Parameters 

92 ---------- 

93 mc_hits_key : str, optional 

94 Name of the mc_hits to convert (default: 'McHits'). 

95 mc_tracks_key : str, optional 

96 Name of the mc_tracks to convert (default: 'McTracks'). 

97 event_info_key : str, optional 

98 Name of the event_info to store this in (default: 'EventInfo'). 

99 """ 

100 

101 def configure(self): 

102 # get event_info, hits and mc_tracks key ; define conversion func 

103 self.event_info_key = self.get("event_info_key", default="EventInfo") 

104 self.mc_tracks_key = self.get("mc_tracks_key", default="McTracks") 

105 self.mc_hits_key = self.get("mc_hits_key", default="McHits") 

106 

107 self.convert_mc_times_to_jte_times = np.frompyfunc( 

108 convert_mc_times_to_jte_times, 3, 1 

109 ) 

110 

111 def process(self, blob): 

112 # convert the mc times to jte times 

113 event_info = blob[self.event_info_key] 

114 mc_tracks = blob[self.mc_tracks_key] 

115 mc_hits = blob[self.mc_hits_key] 

116 timestamp_in_ns = event_info.timestamp * 1e9 + event_info.nanoseconds 

117 

118 mc_tracks["time"] = self.convert_mc_times_to_jte_times( 

119 mc_tracks.time, timestamp_in_ns, event_info.mc_time 

120 ) 

121 mc_hits["time"] = self.convert_mc_times_to_jte_times( 

122 mc_hits.time, timestamp_in_ns, event_info.mc_time 

123 ) 

124 

125 blob[self.mc_tracks_key] = mc_tracks 

126 blob[self.mc_hits_key] = mc_hits 

127 

128 return blob