Skip to main content

gondola_core/events/
mc_tree.rs

1// The following file is part of gaps-online-software and published 
2// under the GPLv3 license
3
4use crate::prelude::*;
5
6#[derive(Debug, Clone, PartialEq)]
7#[cfg_attr(feature="pybindings", pyclass)] 
8pub struct McTree {
9  //pub tracks    : Vec<McTrack>,
10  pub trackmap  : HashMap<u32, McTrack>,
11  pub max_gen   : usize,
12
13  //pub track_map : HashMap<usize, McTrack>,
14  ////pub n_tracks : usize;
15  //  // number of generations in the tree
16  //  uint NGenerations() const;
17  //  // number of tracks in total
18  //  uint NTracks() const;
19  //  // the initial track of the priamry
20  //  CTrackMc GetPrimary() const;
21  //  // daughters of any track in the tree
22  //  std::vector<CTrackMc> GetDaughters(CTrackMc track) const;
23  //  // get a specific generation
24  //  std::vector<CTrackMc> GetGeneration(uint gen) const;
25  //  // access to the whole tree
26  //  const std::map<uint, std::vector<CTrackMc>>& GetTree() const;
27
28
29  //private:
30  //  void InitializeTree(std::vector<CTrackMc> tracks);
31  //  void MakeGenerationTree();
32  //  bool StillAlive(CTrackMc track) const;
33  //  uint maxGen_;
34  //  uint nTracks_;
35  //  std::map<uint, CTrackMc> tracks_;
36  //  std::map<uint, std::vector<CTrackMc>> generationTree_;
37  //  CTrackMc primary_;
38  //};
39}
40
41impl McTree {
42  pub fn new() -> Self {
43    Self {
44      //tracks   : Vec::<McTrack>::new(),
45      max_gen  : 0,
46      trackmap : HashMap::<u32, McTrack>::new(),
47    }
48  }
49
50  pub fn get_daughters(&self) -> Vec<McTrack> {
51    let mut daughters =  Vec::<McTrack>::new();
52  //for (auto t_pair : tracks_)
53  //  {
54  //    if (t_pair.second.GetParentId() == track.GetTrackId())
55  //      {daughters.push_back(t_pair.second);}
56  //  }
57    return daughters;
58  }
59
60  // create all the tracks and create the actual 
61  // tree
62  pub fn assemble(&mut self, hits : &mut Vec<McHit>) {
63    let mut trackmap = HashMap::<u32, McTrack>::new();
64    let mut nhits = hits.len();
65    while nhits > 0 {
66      let h = hits.pop().unwrap();
67      if trackmap.contains_key(&h.track_id) {
68        trackmap.get_mut(&h.track_id).unwrap().hits.push(h);
69      } else {
70        let mut new_track   = McTrack::new();
71        new_track.track.pdg = h.pdg;
72        new_track.hits      = vec![h];
73        trackmap.insert(h.track_id, new_track);
74      }
75      nhits -= 1;
76    }
77    let mut tm_keys = Vec::<u32>::new();
78    for k in trackmap.keys() {
79      tm_keys.push(*k);
80    }
81    for k in &tm_keys {
82      trackmap.get_mut(&k).unwrap().hits.sort_by(|i,j| i.glob_time.total_cmp(&j.glob_time));  
83    }
84    self.trackmap = trackmap;
85  }
86}
87
88impl fmt::Display for McTree {
89  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
90    let mut repr = String::from("<McTree");
91    let mut sorted_keys : Vec<u32> = self.trackmap.keys().copied().collect();
92    sorted_keys.sort();
93    for k in sorted_keys {
94      repr += &(format!("\n  Track Id: {} [pdg {}] -> NHits : {}",k,self.trackmap[&k].track.pdg, self.trackmap[&k].hits.len()));
95    }
96    repr += ">";
97    write!(f, "{}", repr) 
98  }
99}
100
101#[cfg(feature="pybindings")] 
102#[pymethods] 
103impl McTree {
104
105  #[getter]
106  fn get_trackids(&self) -> Vec<u32> {
107    let mut sorted_keys : Vec<u32> = self.trackmap.keys().copied().collect();
108    sorted_keys.sort();
109    sorted_keys
110  }
111
112  fn get_track(&self, track_id : u32) -> Option<McTrack> {
113    self.trackmap.get(&track_id).clone().cloned()
114  }
115}
116
117
118
119#[cfg(feature="pybindings")]
120pythonize!(McTree);