Examples/python/visualization.py

0001 import acts
0002 import numpy as np
0003 import matplotlib.pyplot as plt
0004
0005 from matplotlib.patches import Polygon
0006 from matplotlib.collections import PatchCollection, LineCollection
0007
0008
0009 def polyArea(face):
0010     x = [item[0] for item in face]
0011     y = [item[1] for item in face]
0012     # shoelace formula for computing area of polygons
0013     return 0.5 * abs(np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1)))
0014
0015
0016 def computeProjection(char, vector3):
0017     if char == "x":
0018         return vector3[0]
0019     elif char == "y":
0020         return vector3[1]
0021     elif char == "z":
0022         return vector3[2]
0023     elif char == "r":
0024         return np.sqrt(vector3[0] ** 2 + vector3[1] ** 2)
0025     else:
0026         print("Not a valid projection")
0027
0028
0029 # algorithm to visualize trackk
0030 class TrackVisualizerAlg(acts.examples.IAlgorithm):
0031     def __init__(self, name, level, vis):
0032         acts.examples.IAlgorithm.__init__(self, name, level)
0033
0034         self._vis = vis
0035         self.tracks = acts.examples.ReadDataHandle(
0036             self, acts.examples.ConstTrackContainer, "Tracks"
0037         )
0038         self.tracks.initialize("tracks")
0039
0040     def execute(self, context):
0041         tracks = self.tracks(context.eventStore)
0042         print(f"Event {context.eventNumber}: {len(tracks)} tracks")
0043         for track in tracks:
0044             acts.EventDataView3D.drawTrack(
0045                 self._vis, track, context.geoContext
0046             )  # draw track not a free function
0047
0048         return acts.examples.ProcessCode.SUCCESS
0049
0050
0051 class PyVisualization2D(acts.VisualizationBuffer):
0052
0053     def plot(self, projection, filename, linewidth=None, linestyle=None):
0054         import matplotlib.pyplot as plt
0055
0056         # Reduce font size and complexity to avoid raster overflow
0057         plt.rcParams["font.size"] = 8
0058         plt.rcParams["figure.figsize"] = (12, 10)
0059
0060         fig, ax = plt.subplots()
0061
0062         if linewidth == None:
0063             width = 1
0064         else:
0065             width = linewidth
0066
0067         if linestyle == None:
0068             style = "solid"
0069         else:
0070             style = linestyle
0071
0072         proj2D = [char for char in projection]
0073
0074         if len(proj2D) > 2:
0075             print("Only 2D projection supported")
0076
0077         surfaces2D = [
0078             [
0079                 [computeProjection(proj2D[0], v), computeProjection(proj2D[1], v)]
0080                 for v in surface
0081             ]
0082             for surface in self.surfaces
0083         ]
0084
0085         poly_patches = [
0086             Polygon(face, closed=True) for face in surfaces2D
0087         ]  # face = [[x1,y1], [x2,y2],...]
0088         poly_collection = PatchCollection(poly_patches, alpha=0.5)
0089         poly_collection.set_facecolor(self.faceColors / 255)
0090
0091         ax.set_xlabel(proj2D[0])
0092         ax.set_ylabel(proj2D[1])
0093         ax.ticklabel_format(useOffset=False, style="plain")
0094         ax.set_xticks(np.linspace(-1000, 1000, 5))
0095         ax.set_yticks(np.linspace(-1000, 1000, 5))
0096
0097         ax.add_collection(poly_collection)
0098
0099         ax.autoscale()
0100
0101         for n_face, face in enumerate(surfaces2D):
0102             if polyArea(face) < 1e-8:
0103                 ax.plot(
0104                     [item[0] for item in face],
0105                     [item[1] for item in face],
0106                     color=(
0107                         self.faceColors[n_face][0] / 255,
0108                         self.faceColors[n_face][1] / 255,
0109                         self.faceColors[n_face][2] / 255,
0110                         0.5,
0111                     ),
0112                     lw=1,
0113                 )
0114
0115         # Check if there is a track to be drawn
0116         print(len(self.segments))
0117         if len(self.segments) != 0:
0118             line_segments = [
0119                 [
0120                     [
0121                         computeProjection(proj2D[0], segment[0]),
0122                         computeProjection(proj2D[1], segment[0]),
0123                     ],
0124                     [
0125                         computeProjection(proj2D[0], segment[1]),
0126                         computeProjection(proj2D[1], segment[1]),
0127                     ],
0128                 ]
0129                 for segment in self.segments
0130             ]
0131             line_collection = LineCollection(
0132                 line_segments, linewidths=width, linestyles=style
0133             )
0134             line_collection.set_color(self.lineColor / 255)
0135             ax.add_collection(line_collection)
0136
0137         ax.relim()
0138         ax.autoscale_view()
0139         fig.savefig(filename)