utils.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import ROOT
  2. __all__ = ["pdg", "show_function", "show_value"]
  3. db = ROOT.TDatabasePDG()
  4. class PDGParticle:
  5. def __init__(self, tPart):
  6. self.pdgId = tPart.PdgCode()
  7. self.name = tPart.GetName()
  8. self.charge = tPart.Charge() / 3.0
  9. self.mass = tPart.Mass()
  10. self.spin = tPart.Spin()
  11. def __repr__(self):
  12. return (f"<PDGParticle {self.name}:"
  13. f"pdgId={self.pdgId}, charge={self.charge}, mass={self.mass:5.4e} GeV, spin={self.spin}>")
  14. def pdg(pdg_id):
  15. try:
  16. return PDGParticle(db.GetParticle(pdg_id))
  17. except ReferenceError:
  18. raise ValueError(f"unknown pdgId: {pdg_id}")
  19. def show_function(dataset, fname):
  20. from IPython.display import Markdown
  21. def md_single(fname_):
  22. impl = dataset._function_impl_lookup[fname_]
  23. return '*{}*\n-----\n```cpp\n{}\n```\n\n---'.format(fname_, impl)
  24. try:
  25. return Markdown('\n'.join(md_single(fname_) for fname_ in iter(fname)))
  26. except TypeError:
  27. return Markdown(md_single(fname))
  28. def show_value(dataset, container):
  29. from IPython.display import Image
  30. from graph_vals import parse
  31. if type(container) != str:
  32. container = container.GetName().split(':')[1]
  33. g, functions = parse(dataset.values[container], container)
  34. try:
  35. return Image(g.create_gif()), show_function(dataset, functions)
  36. except Exception as e:
  37. print(e)
  38. print(g.to_string())