handle invalid particles and events without history, improved tooltips (

#27)
scikit-hep · Aug 23, 2022 · af9e753 · af9e753
1 parent a0a1a25
commit af9e753
Show file tree

Hide file tree

Showing 5 changed files with 221 additions and 47 deletions.
diff --git a/examples/visualization.ipynb b/examples/visualization.ipynb
diff --git a/src/pyhepmc/view.py b/src/pyhepmc/view.py
@@ -1,7 +1,8 @@
 from graphviz import Digraph
-from particle import Particle, ParticleNotFound
+from particle import Particle, ParticleNotFound, InvalidParticle
 from pyhepmc._prettify import db as _prettify
 from pyhepmc import Units
+import numpy as np
 
 
 def to_dot(evt, style=None):
@@ -17,19 +18,28 @@ def to_dot(evt, style=None):
     for v in evt.vertices:
         d.node(f"{v.id}", height="0.02", color=gray, tooltip=f"status={v.status}")
 
-    vi = 0
-    vo = 0
     for p in evt.particles:
         try:
             pdb = Particle.from_pdgid(p.pid)
             charge = pdb.charge
             pname = pdb.name
             pname = _prettify.get(pdb.pdgid, pname)
-            tooltip = str(pdb)
+            tooltip = pdb.name
+            if pdb.quarks:
+                tooltip += f" ({pdb.quarks})"
+            tooltip += f"\nQ = {pdb.charge:.3g}"
+            if pdb.mass:
+                tooltip += f"\nm = {pdb.mass:.4g} MeV/c2"
+            if pdb.ctau and np.isfinite(pdb.ctau):
+                tooltip += f"\ncτ = {pdb.ctau:.3g} mm"
         except ParticleNotFound:
             pname = f"Internal({p.pid})"
             charge = 0
-            tooltip = ""
+            tooltip = pname
+        except InvalidParticle:
+            pname = f"Invalid({p.pid})"
+            charge = 0
+            tooltip = pname
         en = p.momentum.e / GeV
         unit = "GeV"
         if en > 1e2:
@@ -39,44 +49,44 @@ def to_dot(evt, style=None):
             en *= 1e3
             unit = "MeV"
 
-        label = f" {pname} {en:.2g} {unit}"
-        style = "solid" if charge else "dashed"
-        if not p.parents:  # initial particles
-            vid = f"in_{vi}"
-            vi += 1
-            d.node(vid, style="invis")
-            assert p.end_vertex
-            d.edge(
-                vid,
-                f"{p.end_vertex.id}",
-                label=label,
-                arrowsize="1",
-                style="bold",
-                tooltip=tooltip,
-            )
-        elif not p.children:  # final state particles
-            vid = f"out_{vo}"
-            vo += 1
-            d.node(vid, style="invis")
-            assert p.production_vertex
-            d.edge(
-                f"{p.production_vertex.id}",
-                vid,
-                label=label,
-                style=style,
-                tooltip=tooltip,
-            )
+        label = f"{pname} {en:.2g} {unit}"
+
+        # do not connect initial particles to root vertex
+        if p.production_vertex and p.production_vertex.id != 0:
+            vin = f"{p.production_vertex.id}"
+        else:
+            vin = f"in_{p.id}"
+            d.node(vin, style="invis")
+
+        if p.end_vertex:
+            vout = f"{p.end_vertex.id}"
         else:
-            # intermediate particles
-            d.edge(
-                f"{p.production_vertex.id}",
-                f"{p.end_vertex.id}",
-                label=label,
-                color=gray,
-                fontcolor=gray,
-                style=style,
-                dir="none",
-                tooltip=tooltip,
-            )
+            vout = f"out_{p.id}"
+            d.node(vout, style="invis")
+
+        style = "solid" if charge else "dashed"
+        if not p.production_vertex or p.production_vertex.id == 0 and p.end_vertex:
+            # initial particle which is not also final state is beam particle
+            style = "bold"
+            edir = "forward"
+            color = None
+        elif p.status == 1:  # final state
+            edir = "forward"
+            color = None
+        else:  # intermediate
+            edir = "none"
+            color = gray
+
+        d.edge(
+            vin,
+            vout,
+            label=label,
+            dir=edir,
+            style=style,
+            color=color,
+            fontcolor=color,
+            tooltip=tooltip,
+            labeltooltip=tooltip,
+        )
 
     return d
diff --git a/tests/pythia6_simple.dat b/tests/pythia6_simple.dat
@@ -0,0 +1,51 @@
+HepMC::Version 3.02.05
+HepMC::Asciiv3-START_EVENT_LISTING
+E 0 19 40
+U GEV MM
+P 1 0 2212 0.0000000000000000e+00 0.0000000000000000e+00 4.9991195718917339e+01 4.9999999999999993e+01 9.3827000000000005e-01 3
+P 2 0 2212 0.0000000000000000e+00 0.0000000000000000e+00 -4.9991195718917339e+01 4.9999999999999993e+01 9.3827000000000005e-01 3
+P 3 1 1 -2.1806871592392363e-02 5.8805329337607859e-01 -4.8466219038425701e-01 7.6235139816699049e-01 -0.0000000000000000e+00 3
+P 4 2 1 -2.1399492116645962e+00 6.3066455563989210e-01 -5.5599470506278008e+00 5.9908373050875925e+00 -0.0000000000000000e+00 3
+P 5 3 1 -2.1806871592392363e-02 5.8805329337607881e-01 -4.8466219038425762e-01 7.6235139816699082e-01 0.0000000000000000e+00 3
+P 6 4 1 -2.1399492116645962e+00 6.3066455563989210e-01 -5.5599470506278017e+00 5.9908373050875934e+00 0.0000000000000000e+00 3
+V -5 0 [5,6]
+P 7 -5 1 -2.2428932413666449e+00 8.0921610252646115e-01 -5.4337023479006294e+00 5.9430145547808575e+00 3.3000000000000002e-01 3
+P 8 -5 1 8.1137158109656382e-02 4.0950174648950965e-01 -6.1090689311142943e-01 8.1017414847372760e-01 3.3000000000000002e-01 3
+P 9 7 1 -2.2428932413666449e+00 8.0921610252646115e-01 -5.4337023479006294e+00 5.9430145547808566e+00 3.3000000000000002e-01 2
+P 10 1 2203 2.1806871592392363e-02 -5.8805329337607859e-01 4.9609014856830349e+01 4.9618500453456001e+01 7.7132999999999996e-01 2
+P 11 8 1 8.1137158109656382e-02 4.0950174648950965e-01 -6.1090689311142943e-01 8.1017414847372760e-01 3.3000000000000002e-01 2
+P 12 2 2203 2.1399492116645962e+00 -6.3066455563989210e-01 -4.3564405615818288e+01 4.3628310843289412e+01 7.7132999999999996e-01 2
+V -8 0 [9,10]
+P 13 -8 92 -2.2210863697742527e+00 2.2116280915038256e-01 4.4175312508929721e+01 5.5561515008236860e+01 3.3625014157599679e+01 2
+P 14 13 331 -1.1738012688921793e+00 8.0765689386644379e-01 -3.0798593525447195e+00 3.5261610954816991e+00 9.5820632578202913e-01 2
+P 15 13 111 -3.1519450654286596e-01 -3.5980562516633882e-02 2.9900482170948783e-01 4.5636132794019346e-01 1.3497999999999999e-01 2
+P 16 13 -211 -7.0637494262006184e-01 3.1031715494884905e-02 -2.1658732822757143e+00 2.2826334323986108e+00 1.3957000000000000e-01 1
+P 17 13 211 -4.5612897859630697e-01 -4.7514717593545824e-01 8.1214679879915241e+00 8.1493276132933712e+00 1.3957000000000000e-01 1
+P 18 13 2112 9.0044398831400763e-01 -2.5126197906059283e-01 6.0832400438584600e+00 6.2259564063937489e+00 9.3957000000000002e-01 1
+P 19 13 211 -4.7003066143684613e-01 1.4486391730173881e-01 3.4917332290190672e+01 3.4921075132729221e+01 1.3957000000000000e-01 1
+V -9 0 [11,12]
+P 20 -9 92 2.2210863697742527e+00 -2.2116280915038244e-01 -4.4175312508929714e+01 4.4438484991763140e+01 4.2823562966441093e+00 2
+P 21 20 113 5.9967849939287710e-01 3.4510556587052538e-01 -3.3810019804438327e+00 3.5124880762951549e+00 6.5397701803779784e-01 2
+P 22 20 221 1.0839047546205272e-01 -4.2027376334831035e-01 -1.2331149026154447e+01 1.2350923703879788e+01 5.4744999999999999e-01 2
+P 23 20 2214 1.2790389523841685e+00 4.0035046648034817e-01 -2.0044276688780990e+01 2.0134156202149974e+01 1.3472181925873288e+00 2
+P 24 20 111 2.3397844253515424e-01 -5.4634507815294564e-01 -8.4188848135504522e+00 8.4409170094382269e+00 1.3497999999999999e-01 2
+P 25 14 22 -3.7866530724722947e-01 1.3804341125450670e-01 -5.7270017461794431e-01 7.0030628178736276e-01 0.0000000000000000e+00 1
+P 26 14 113 -7.9513596164494982e-01 6.6961348261193698e-01 -2.5071591779267748e+00 2.8258548136943360e+00 7.8675591578348014e-01 2
+V -12 0 [15] @ -4.1740936725268328e-06 -4.7648748698674578e-07 3.9596950722326455e-06 6.0435537161941038e-06
+P 27 -12 22 -6.5398954996743419e-02 -4.0922281680830369e-02 1.4106331269034939e-01 1.6078095235383630e-01 0.0000000000000000e+00 1
+P 28 -12 22 -2.4979555154612260e-01 4.9417191641964835e-03 1.5794150901913848e-01 2.9558037558635719e-01 0.0000000000000000e+00 1
+P 29 21 211 -9.9665405922456762e-02 6.1820172224834033e-02 -5.5589781563603535e-01 5.8502742940866614e-01 1.3957000000000000e-01 1
+P 30 21 -211 6.9934390531533386e-01 2.8328539364569133e-01 -2.8251041648077972e+00 2.9274606468864888e+00 1.3957000000000000e-01 1
+P 31 22 22 -7.5654433284149183e-02 -4.8561402797491426e-01 -7.7089885292840288e+00 7.7246390674305490e+00 0.0000000000000000e+00 1
+P 32 22 22 1.8404490874620189e-01 6.5340264626603889e-02 -4.6221604968704186e+00 4.6262846364492392e+00 0.0000000000000000e+00 1
+P 33 23 2212 1.3146957052863393e+00 5.0503668067441621e-01 -1.6752612574782226e+01 1.6837869981680480e+01 9.3827000000000005e-01 1
+P 34 23 111 -3.5656752902170630e-02 -1.0468621419406804e-01 -3.2916641139987641e+00 3.2962862204694972e+00 1.3497999999999999e-01 2
+V -17 0 [24] @ 3.9111126905523849e-05 -9.1325386451522039e-05 -1.4072752548401441e-03 1.4109580898912608e-03
+P 35 -17 22 1.0231666161604391e-01 -3.9715864280443258e-01 -5.8051273894918847e+00 5.8195968670542273e+00 0.0000000000000000e+00 1
+P 36 -17 22 1.3166178091911035e-01 -1.4918643534851311e-01 -2.6137574240585675e+00 2.6213201423839996e+00 0.0000000000000000e+00 1
+P 37 26 -211 -4.6274315409360200e-01 -5.9074803679170951e-02 -8.3698279598279901e-01 9.6831866901418662e-01 1.3957000000000000e-01 1
+P 38 26 211 -3.3239280755134815e-01 7.2868828629110816e-01 -1.6701763819439761e+00 1.8575361446801499e+00 1.3957000000000000e-01 1
+V -19 0 [34] @ -1.1940887063064449e-06 -3.5057770520518805e-06 -1.1023266628523709e-04 1.1038745337847554e-04
+P 39 -19 22 -7.4358455762204817e-02 -1.3769292753516078e-02 -1.4720686737358513e+00 1.4740098213922117e+00 0.0000000000000000e+00 1
+P 40 -19 22 3.8701702860034173e-02 -9.0916921440551968e-02 -1.8195954402629129e+00 1.8222763990772854e+00 0.0000000000000000e+00 1
+HepMC::Asciiv3-END_EVENT_LISTING
diff --git a/tests/sibyll21.dat b/tests/sibyll21.dat
@@ -0,0 +1,26 @@
+HepMC::Version 3.02.05
+HepMC::Asciiv3-START_EVENT_LISTING
+E 0 0 21
+U GEV MM
+P 1 0 0 -6.3597112894058228e-01 -6.7266426980495453e-02 1.5901356935501099e+00 1.8782439231872559e+00 7.6829999685287476e-01 2
+P 2 0 0 3.0926612019538879e-01 2.1742968261241913e-01 1.8657438755035400e+00 2.0577671527862549e+00 7.6859998703002930e-01 2
+P 3 0 211 -5.7559907436370850e-01 -4.3126103281974792e-01 8.6977729797363281e+00 8.7294225692749023e+00 1.3956999778747559e-01 1
+P 4 0 2212 4.6384686231613159e-01 1.7648369073867798e-01 2.5262069702148438e+01 2.5284357070922852e+01 9.3826997280120850e-01 1
+P 5 0 0 5.1563632488250732e-01 -2.6430293917655945e-01 5.8055062294006348e+00 5.9959793090820312e+00 1.3827999830245972e+00 2
+P 6 0 321 1.4977657794952393e-01 7.2727096080780029e-01 -2.5128841400146484e-02 9.1774952411651611e-01 4.9364998936653137e-01 1
+P 7 0 0 -3.4058013558387756e-01 -1.4105777442455292e-01 2.3133645057678223e+00 2.4741170406341553e+00 7.6829999685287476e-01 2
+P 8 0 0 3.9163086563348770e-02 -3.8725398480892181e-02 3.5180878639221191e+00 3.6043610572814941e+00 7.8194999694824219e-01 2
+P 9 0 211 -2.3349599540233612e-01 -2.8010216355323792e-01 -3.8940856456756592e+00 3.9136114120483398e+00 1.3956999778747559e-01 1
+P 10 0 2112 3.0795735120773315e-01 1.0153134167194366e-01 -4.5133453369140625e+01 4.5144393920898438e+01 9.3957000970840454e-01 1
+P 11 0 -211 -5.7138837873935699e-02 1.5356467664241791e-01 6.2309533357620239e-02 2.2407411038875580e-01 1.3956999778747559e-01 1
+P 12 0 111 -5.7883763313293457e-01 -2.2083166241645813e-01 1.5278395414352417e+00 1.6541855335235596e+00 1.3496999442577362e-01 1
+P 13 0 211 4.8268994688987732e-01 2.3095960915088654e-01 1.6581695079803467e+00 1.7511726617813110e+00 1.3956999778747559e-01 1
+P 14 0 -211 -1.7342118918895721e-01 -1.3528063893318176e-02 2.0759028196334839e-01 3.0661198496818542e-01 1.3956999778747559e-01 1
+P 15 0 3122 4.7702565789222717e-01 -3.4917265176773071e-02 4.3904485702514648e+00 4.5551543235778809e+00 1.1156300306320190e+00 1
+P 16 0 211 3.8611851632595062e-02 -2.2938628494739532e-01 1.4150713682174683e+00 1.4408383369445801e+00 1.3956999778747559e-01 1
+P 17 0 -211 -4.5473271608352661e-01 2.6130743324756622e-02 2.0567431449890137e+00 2.1173090934753418e+00 1.3956999778747559e-01 1
+P 18 0 111 1.1414974182844162e-01 -1.6718970239162445e-01 2.5664079189300537e-01 3.5682868957519531e-01 1.3496999442577362e-01 1
+P 19 0 211 2.3528023064136505e-01 -8.2750454545021057e-02 2.3999428749084473e+00 2.4169008731842041e+00 1.3956999778747559e-01 1
+P 20 0 -211 -6.7817553877830505e-02 2.0679535344243050e-02 2.2056096792221069e-01 2.7046951651573181e-01 1.3956999778747559e-01 1
+P 21 0 111 -1.2830749154090881e-01 2.3347064852714539e-02 8.9765828847885132e-01 9.1706871986389160e-01 1.3496999442577362e-01 1
+HepMC::Asciiv3-END_EVENT_LISTING
diff --git a/tests/test_view.py b/tests/test_view.py
@@ -8,12 +8,14 @@
 def test_dot(evt):  # noqa
     d = view.to_dot(evt)
     s = str(d)
+    print(s)
     assert s.startswith('digraph "event 1')
-    assert "in_0" in s
     assert "in_1" in s
-    assert "out_0" in s
-    assert "out_1" in s
-    assert "out_2" in s
+    assert "in_2" in s
+    assert "in_3" not in s
+    assert "out_6" in s
+    assert "out_7" in s
+    assert "out_8" in s
 
 
 def test_dot_2():