Coverage for src/km3pipe/mc.py: 69%

1# Filename: mc.py 

2# pylint: disable=C0103 

3""" 

4Monte Carlo related things. 

5 

6""" 

7import particle 

8import numpy as np 

9 

10from .logger import get_logger 

11 

12__author__ = "Tamas Gal and Moritz Lotze" 

13__copyright__ = "Copyright 2016, Tamas Gal and the KM3NeT collaboration." 

14__credits__ = [] 

15__license__ = "MIT" 

16__maintainer__ = "Tamas Gal and Moritz Lotze" 

17__email__ = "tgal@km3net.de" 

18__status__ = "Development" 

19 

20log = get_logger(__name__) # pylint: disable=C0103 

21 

22 

23def geant2pdg(geant_code): 

24 """Convert GEANT particle ID to PDG""" 

25 try: 

26 return particle.Geant3ID(geant_code).to_pdgid() 

27 except KeyError: 

28 return 0 

29 

30 

31def pdg2name(pdg_id): 

32 """Convert PDG ID to human readable names""" 

33 try: 

34 return particle.Particle.from_pdgid(pdg_id).name 

35 except particle.particle.particle.InvalidParticle: 

36 log.debug("No particle with PDG ID %d found!" % pdg_id) 

37 return "N/A" 

38 

39 

40def name2pdg(name): 

41 """Return best match of a PDG ID for the given name""" 

42 return particle.Particle.from_string(name) 

43 

44 

45def most_energetic(df): 

46 """Grab most energetic particle from mc_tracks dataframe.""" 

47 idx = df.groupby(["event_id"])["energy"].transform(max) == df["energy"] 

48 return df[idx].reindex() 

49 

50 

51def leading_particle(df): 

52 """Grab leading particle (neutrino, most energetic bundle muon). 

53 

54 Note: selecting the most energetic mc particle does not always select 

55 the neutrino! In some sub-percent cases, the post-interaction 

56 secondaries can have more energy than the incoming neutrino! 

57 

58 aanet convention: mc_tracks[0] = neutrino 

59 so grab the first row 

60 

61 if the first row is not unique (neutrinos are unique), it's a muon bundle 

62 grab the most energetic then 

63 """ 

64 leading = df.groupby("event_id", as_index=False).first() 

65 unique = leading.type.unique() 

66 

67 if len(unique) == 1 and unique[0] == 0: 

68 leading = most_energetic(df) 

69 return leading 

70 

71 

72def get_flavor(pdg_types): 

73 """Build a 'flavor' from the 'type' column.""" 

74 pd = km3pipe.extras.pandas() 

75 

76 return pd.Series(pdg_types).apply(pdg2name) 

77 

78 

79def _p_eq_nu(pdg_type): 

80 return np.abs(pdg_type) in {12, 14, 16} 

81 

82 

83def _p_eq_mu(pdg_type): 

84 return pdg_type == -13 

85 

86 

87def is_neutrino(pdg_types): 

88 """flavor string -> is_neutrino""" 

89 pd = km3pipe.extras.pandas() 

90 

91 return pd.Series(pdg_types).apply(_p_eq_nu) 

92 

93 

94def is_muon(pdg_types): 

95 """flavor string -> is_neutrino""" 

96 pd = km3pipe.extras.pandas() 

97 

98 return pd.Series(pdg_types).apply(_p_eq_mu) 

99 

100 

101def convert_mc_times_to_jte_times(times_mc, evt_timestamp_in_ns, evt_mc_time): 

102 """ 

103 Function that converts MC times to JTE times. 

104 

105 Parameters 

106 ---------- 

107 times_mc : np.ndarray 

108 Time array with MC times. 

109 evt_timestamp_in_ns : int 

110 Total timestamp of the event in nanoseconds. 

111 evt_mc_time : int 

112 Mc time of the event in nanoseconds. 

113 

114 Returns 

115 ------- 

116 ndarray 

117 Converted time array with JTE times. 

118 """ 

119 # needs to be cast to normal ndarray (not recarray), or else we 

120 # would get invalid type promotion 

121 times_mc = np.array(times_mc).astype(float) 

122 times_jte = times_mc - evt_timestamp_in_ns + evt_mc_time 

123 return times_jte